Upload folder using huggingface_hub

agent.py

@@ -0,0 +1,26 @@
+from transformers import ReactCodeAgent, HfApiEngine
+from prompts import SQUAD_REACT_CODE_SYSTEM_PROMPT
+from tools.squad_tools import SquadRetrieverTool, SquadQueryTool
+from tools.text_to_image import TextToImageTool
+
+def get_agent():
+    # model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    model_name = "http://localhost:1234/v1"
+
+    llm_engine = HfApiEngine(model_name)
+
+    TASK_SOLVING_TOOLBOX = [
+        SquadRetrieverTool(),
+        SquadQueryTool(),
+        TextToImageTool(),
+    ]
+
+    # Initialize the agent with both tools
+    agent = ReactCodeAgent(
+        tools=TASK_SOLVING_TOOLBOX,
+        llm_engine=llm_engine,
+        system_prompt=SQUAD_REACT_CODE_SYSTEM_PROMPT,
+    )
+
+    return agent
+

app.py

@@ -1,40 +1,21 @@
 import gradio as gr
 from gradio import ChatMessage
-from transformers import ReactCodeAgent, HfApiEngine
 from utils import stream_from_transformers_agent
-from prompts import SQUAD_REACT_CODE_SYSTEM_PROMPT
-from tools.squad_tools import SquadRetrieverTool, SquadQueryTool
-from tools.text_to_image import TextToImageTool
 from gradio.context import Context
 from gradio import Request
 import pickle
 import os
 from dotenv import load_dotenv
+from agent import get_agent
 
 load_dotenv()
 
-TASK_SOLVING_TOOLBOX = [
-    SquadRetrieverTool(),
-    SquadQueryTool(),
-    TextToImageTool(),
-]
-
 sessions_path = "sessions.pkl"
 sessions = pickle.load(open(sessions_path, "rb")) if os.path.exists(sessions_path) else {}
 
-app = None
-
-model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-# model_name = "http://localhost:1234/v1"
+agent = get_agent()
 
-llm_engine = HfApiEngine(model_name)
-
-# Initialize the agent with both tools
-agent = ReactCodeAgent(
-    tools=TASK_SOLVING_TOOLBOX,
-    llm_engine=llm_engine,
-    system_prompt=SQUAD_REACT_CODE_SYSTEM_PROMPT,
-)
+app = None
 
 def append_example_message(x: gr.SelectData, messages):
     if x.value["text"] is not None: