Update app.py

app.py

@@ -9,11 +9,12 @@ from importlib.resources import files
 from txagent import TxAgent
 from tooluniverse import ToolUniverse
 
-# Patch PyTorch to allow loading old numpy pickles
+# Allow loading old numpy types with torch.load
 torch.serialization.add_safe_globals([
     numpy.core.multiarray._reconstruct,
     numpy.ndarray,
-    numpy.dtype
+    numpy.dtype,
+    numpy.dtypes.Float32DType
 ])
 
 logging.basicConfig(
@@ -29,7 +30,7 @@ current_dir = os.path.dirname(os.path.abspath(__file__))
 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_47dc56b3e3ddeb31af4f19defdd538d984de1500368852a0fab80bc2e826c944.pt",
+    "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding.pt",
     "tool_files": {
         "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
         "fda_drug_label": str(files('tooluniverse.data').joinpath('fda_drug_labeling_tools.json')),
@@ -39,42 +40,15 @@ CONFIG = {
     }
 }
 
-def safe_load_embeddings(filepath):
-    try:
-        return torch.load(filepath, weights_only=True)
-    except Exception as e:
-        logger.warning(f"Retrying with weights_only=False due to: {e}")
-        try:
-            return torch.load(filepath, weights_only=False)
-        except Exception as e:
-            logger.error(f"Failed to load embeddings: {e}")
-            return None
-
-def patch_embedding_loading():
-    from txagent.toolrag import ToolRAGModel
-    def patched_load(self, tooluniverse):
-        try:
-            if not os.path.exists(CONFIG["embedding_filename"]):
-                return False
-            self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
-            if self.tool_desc_embedding is None:
-                logger.error("Tool embedding file could not be loaded.")
-                return False
-
-            tools = tooluniverse.get_all_tools() if hasattr(tooluniverse, "get_all_tools") else getattr(tooluniverse, "tools", [])
-            if len(tools) != len(self.tool_desc_embedding):
-                logger.warning("Tool count mismatch.")
-                if len(self.tool_desc_embedding) > len(tools):
-                    self.tool_desc_embedding = self.tool_desc_embedding[:len(tools)]
-                else:
-                    padding = self.tool_desc_embedding[-1].unsqueeze(0).repeat(len(tools) - len(self.tool_desc_embedding), 1)
-                    self.tool_desc_embedding = torch.cat([self.tool_desc_embedding, padding], dim=0)
-            return True
-        except Exception as e:
-            logger.error(f"Embedding load failed: {e}")
-            return False
-
-    ToolRAGModel.load_tool_desc_embedding = patched_load
+def generate_tool_embeddings(agent):
+    tu = ToolUniverse(tool_files=CONFIG["tool_files"])
+    tu.load_tools()
+    embedding_tensor = agent.rag_model.generate_tool_desc_embedding(tu)
+    if embedding_tensor is not None:
+        torch.save(embedding_tensor, CONFIG["embedding_filename"])
+        logger.info(f"Saved new embedding tensor to {CONFIG['embedding_filename']}")
+    else:
+        logger.warning("Embedding generation returned None")
 
 def prepare_tool_files():
     os.makedirs(os.path.join(current_dir, 'data'), exist_ok=True)
@@ -88,14 +62,8 @@ def prepare_tool_files():
             logger.error(f"Tool generation failed: {e}")
 
 def create_agent():
-    patch_embedding_loading()
     prepare_tool_files()
     try:
-        tu = ToolUniverse()
-        tools = tu.get_all_tools() if hasattr(tu, "get_all_tools") else getattr(tu, "tools", [])
-        available_tool_names = [t["name"] for t in tools]
-        additional_default_tools = [t for t in ["DirectResponse", "RequireClarification"] if t in available_tool_names]
-
         agent = TxAgent(
             CONFIG["model_name"],
             CONFIG["rag_model_name"],
@@ -104,8 +72,10 @@ def create_agent():
             enable_checker=True,
             step_rag_num=10,
             seed=42,
-            additional_default_tools=additional_default_tools
+            additional_default_tools=["DirectResponse", "RequireClarification"]
         )
+        if not os.path.exists(CONFIG["embedding_filename"]):
+            generate_tool_embeddings(agent)
         agent.init_model()
         return agent
     except Exception as e:
@@ -114,8 +84,7 @@ def create_agent():
 
 def respond(msg, chat_history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
     if not isinstance(msg, str) or len(msg.strip()) <= 10:
-        chat_history.append({"role": "assistant", "content": "Hi, I am TxAgent. Please provide a valid message longer than 10 characters."})
-        return chat_history
+        return chat_history + [{"role": "assistant", "content": "Hi, I am TxAgent. Please provide a valid message longer than 10 characters."}]
 
     message = msg.strip()
     chat_history.append({"role": "user", "content": message})