Update app.py

app.py

@@ -4,27 +4,28 @@ import logging
 import torch
 from txagent import TxAgent
 import gradio as gr
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import snapshot_download
 from tooluniverse import ToolUniverse
 
 # Configuration
 CONFIG = {
     "model_name": "mims-harvard/TxAgent-T1-Llama-3.1-8B",
     "rag_model_name": "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-    "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding_e27fb393f3144ec28f620f33d4d79911.pt",
+    "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding_47dc56b3e3ddeb31af4f19defdd538d984de1500368852a0fab80bc2e826c944.pt",
     "local_dir": "./models",
     "tool_files": {
         "new_tool": "./data/new_tool.json"
     }
 }
 
-# Logging
+# Logging setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 def prepare_tool_files():
     os.makedirs("./data", exist_ok=True)
     if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
+        logger.info("Generating tool list using ToolUniverse...")
         tu = ToolUniverse()
         tools = tu.get_all_tools()
         with open(CONFIG["tool_files"]["new_tool"], "w") as f:
@@ -33,13 +34,14 @@ def prepare_tool_files():
 
 def download_model_files():
     os.makedirs(CONFIG["local_dir"], exist_ok=True)
-    logger.info("Downloading model files...")
+    print("Downloading model files...")
 
     snapshot_download(
         repo_id=CONFIG["model_name"],
         local_dir=os.path.join(CONFIG["local_dir"], CONFIG["model_name"]),
         resume_download=True
     )
+
     snapshot_download(
         repo_id=CONFIG["rag_model_name"],
         local_dir=os.path.join(CONFIG["local_dir"], CONFIG["rag_model_name"]),
@@ -47,24 +49,23 @@ def download_model_files():
     )
 
 def generate_embeddings(agent):
-    """Generates and assigns embeddings manually"""
-    path = os.path.join(CONFIG["local_dir"], CONFIG["embedding_filename"])
-    try:
-        if os.path.exists(path):
-            logger.info("Embedding file already exists, loading...")
-            agent.rag_model.tool_desc_embedding = torch.load(path)
-            return
+    embedding_path = CONFIG["embedding_filename"]
+
+    if os.path.exists(embedding_path):
+        print("Embeddings file already exists — loading...")
+        agent.rag_model.tool_desc_embedding = torch.load(embedding_path)
+        return
 
-        logger.info("Generating embeddings and saving...")
+    print("Generating missing tool embeddings...")
+    try:
         tools = agent.tooluniverse.get_all_tools()
-        descs = [t["description"] for t in tools]
-        embeddings = agent.rag_model.generate_embeddings(descs)
-        torch.save(embeddings, path)
+        descriptions = [tool["description"] for tool in tools]
+        embeddings = agent.rag_model.generate_embeddings(descriptions)
+        torch.save(embeddings, embedding_path)
         agent.rag_model.tool_desc_embedding = embeddings
-        logger.info("✅ Embeddings generated and saved")
-
+        print(f"Embeddings saved to {embedding_path}")
     except Exception as e:
-        logger.error(f"Failed to generate embeddings: {e}")
+        print(f"Failed to generate embeddings: {e}")
         raise
 
 class TxAgentApp:
@@ -111,6 +112,7 @@ class TxAgentApp:
                 max_round=30
             ):
                 response += chunk
+
             return history + [(message, response)]
         except Exception as e:
             return history + [(message, f"Error: {str(e)}")]
@@ -147,4 +149,4 @@ if __name__ == "__main__":
     prepare_tool_files()
     download_model_files()
     interface = create_interface()
-    interface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    interface.launch(server_name="0.0.0.0", server_port=7860, share=False)