Update app.py

app.py

@@ -6,11 +6,11 @@ from txagent import TxAgent
 import gradio as gr
 from tooluniverse import ToolUniverse
 
-# Configuration - Using remote Hugging Face models
+# Configuration - Using your existing embedding 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_47dc56b3e3ddeb31af4f19defdd538d984de1500368852a0fab80bc2e826c944.pt",  # Your exact filename
     "tool_files": {
         "new_tool": "./data/new_tool.json"
     }
@@ -35,26 +35,19 @@ def prepare_tool_files():
 
 def load_embeddings(agent):
     embedding_path = CONFIG["embedding_filename"]
-    if os.path.exists(embedding_path):
-        logger.info("✅ Loading pre-generated embeddings file")
-        try:
+    try:
+        if os.path.exists(embedding_path):
+            logger.info(f"✅ Loading existing embeddings from {embedding_path}")
             embeddings = torch.load(embedding_path)
             agent.rag_model.tool_desc_embedding = embeddings
-            return
-        except Exception as e:
-            logger.error(f"Failed to load embeddings: {e}")
-    
-    logger.info("Generating tool embeddings...")
-    try:
-        tools = agent.tooluniverse.get_all_tools()
-        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(f"Embeddings saved to {embedding_path}")
+            return True
+        else:
+            logger.error(f"❌ Embedding file not found at {embedding_path}")
+            logger.info("Please ensure the embedding file is in the root directory")
+            return False
     except Exception as e:
-        logger.error(f"Failed to generate embeddings: {e}")
-        raise
+        logger.error(f"Failed to load embeddings: {str(e)}")
+        return False
 
 class TxAgentApp:
     def __init__(self):
@@ -66,9 +59,9 @@ class TxAgentApp:
             return "✅ Already initialized"
         
         try:
-            logger.info("Initializing TxAgent with remote models...")
+            logger.info("Initializing TxAgent...")
             
-            # Initialize without local_files_only parameter
+            # Initialize TxAgent
             self.agent = TxAgent(
                 CONFIG["model_name"],
                 CONFIG["rag_model_name"],
@@ -80,14 +73,18 @@ class TxAgentApp:
                 additional_default_tools=["DirectResponse", "RequireClarification"]
             )
             
-            logger.info("Loading models from Hugging Face Hub...")
+            # Initialize models
+            logger.info("Loading models...")
             self.agent.init_model()
             
-            logger.info("Preparing embeddings...")
-            load_embeddings(self.agent)
+            # Load embeddings
+            logger.info("Loading embeddings...")
+            if not load_embeddings(self.agent):
+                return "❌ Failed to load embeddings - check logs"
             
             self.is_initialized = True
-            return "✅ TxAgent initialized successfully (using remote models)"
+            return "✅ TxAgent initialized successfully"
+            
         except Exception as e:
             logger.error(f"Initialization failed: {str(e)}")
             return f"❌ Initialization failed: {str(e)}"
@@ -126,7 +123,6 @@ def create_interface():
     ) as demo:
         gr.Markdown("""
         # 🧠 TxAgent: Therapeutic Reasoning AI
-        ### (Running with remote Hugging Face models)
         """)
         
         with gr.Row():
@@ -168,7 +164,14 @@ if __name__ == "__main__":
     try:
         logger.info("Starting application...")
         
-        # Prepare local tool files
+        # Verify embedding file exists
+        if not os.path.exists(CONFIG["embedding_filename"]):
+            logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
+            logger.info("Please ensure the file is in the root directory")
+        else:
+            logger.info(f"Found embedding file: {CONFIG['embedding_filename']}")
+        
+        # Prepare tool files
         prepare_tool_files()
         
         # Launch interface
@@ -179,5 +182,5 @@ if __name__ == "__main__":
             share=False
         )
     except Exception as e:
-        logger.error(f"Fatal error: {str(e)}")
+        logger.error(f"Application failed to start: {str(e)}")
         raise