Update app.py

app.py

@@ -30,6 +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.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,6 +40,16 @@ CONFIG = {
     }
 }
 
+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)
     if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
@@ -63,6 +74,10 @@ def create_agent():
             seed=42,
             additional_default_tools=["DirectResponse", "RequireClarification"]
         )
+        if not os.path.exists(CONFIG["embedding_filename"]):
+            generate_tool_embeddings(agent)
+        else:
+            logger.info("Embedding file found. Skipping embedding generation.")
         agent.init_model()
         return agent
     except Exception as e:
@@ -130,8 +145,7 @@ def main():
     global agent
     agent = create_agent()
     demo = create_demo(agent)
-    print("Exiting after embedding generation. Please restart the Space manually.")
-    exit()
+    demo.launch(share=True)
 
 if __name__ == "__main__":
     main()