Upload folder using huggingface_hub

agent.py

@@ -3,18 +3,21 @@ 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"
+def get_agent(model_name = None, include_image_tools = False):
+    DEFAULT_MODEL_NAME = "http://localhost:1234/v1"
+    if model_name is None:
+        model_name = DEFAULT_MODEL_NAME
 
     llm_engine = HfApiEngine(model_name)
 
     TASK_SOLVING_TOOLBOX = [
         SquadRetrieverTool(),
-        SquadQueryTool(),
-        TextToImageTool(),
+        SquadQueryTool()
     ]
 
+    if include_image_tools:
+        TASK_SOLVING_TOOLBOX.append(TextToImageTool())
+
     # Initialize the agent with both tools
     agent = ReactCodeAgent(
         tools=TASK_SOLVING_TOOLBOX,

app.py

@@ -13,7 +13,7 @@ load_dotenv()
 sessions_path = "sessions.pkl"
 sessions = pickle.load(open(sessions_path, "rb")) if os.path.exists(sessions_path) else {}
 
-agent = get_agent()
+agent = get_agent(model_name="meta-llama/Meta-Llama-3.1-8B-Instruct", include_image_tools=True)
 
 app = None
 

benchmarking.ipynb

@@ -331,11 +331,11 @@
     "from agent import get_agent\n",
     "\n",
     "benchmarks = [\n",
-    "    (get_agent(), \"baseline\"),\n",
+    "    (get_agent(model_name=\"meta-llama/Meta-Llama-3.1-8B-Instruct\"), \"baseline\"),\n",
     "]\n",
     "\n",
     "for agent, name in tqdm(benchmarks):\n",
-    "    benchmark_agent(agent, dfSample, name)\n"
+    "    benchmark_agent(agent, dfSample, name)"
    ]
   },
   {