Update app.py

app.py

@@ -4,19 +4,23 @@ import logging
 import torch
 import gradio as gr
 from tooluniverse import ToolUniverse
-from txagent import TxAgent  # Updated import statement
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import warnings
 from typing import List, Dict, Any
 
 # Suppress specific warnings
 warnings.filterwarnings("ignore", category=UserWarning)
 
-# Configuration with hardcoded embedding file
+# 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_47dc56b3e3ddeb31af4f19defdd538d984de1500368852a0fab80bc2e826c944.pt",
     "tool_files": {
+        "opentarget": "opentarget_tools.json",
+        "fda_drug_label": "fda_drug_labeling_tools.json",
+        "special_tools": "special_tools.json",
+        "monarch": "monarch_tools.json",
         "new_tool": "./data/new_tool.json"
     }
 }
@@ -34,7 +38,7 @@ def prepare_tool_files():
     if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
         logger.info("Generating tool list using ToolUniverse...")
         tu = ToolUniverse()
-        tools = tu.get_all_tools()
+        tools = tu.get_all_tools() if hasattr(tu, 'get_all_tools') else []
         with open(CONFIG["tool_files"]["new_tool"], "w") as f:
             json.dump(tools, f, indent=2)
         logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
@@ -46,140 +50,134 @@ def safe_load_embeddings(filepath: str) -> Any:
         return torch.load(filepath, weights_only=True)
     except Exception as e:
         logger.warning(f"Secure load failed, trying with weights_only=False: {str(e)}")
-        # If that fails, try with weights_only=False (less secure)
-        return torch.load(filepath, weights_only=False)
-
-def patch_embedding_loading():
-    """Monkey-patch the embedding loading functionality"""
-    try:
-        from txagent.toolrag import ToolRAGModel
-        
-        original_load = ToolRAGModel.load_tool_desc_embedding
-        
-        def patched_load(self, tooluniverse: ToolUniverse) -> bool:
-            try:
-                if not os.path.exists(CONFIG["embedding_filename"]):
-                    logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
-                    return False
-                
-                # Load embeddings safely
-                self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
-                
-                # Handle tool count mismatch
-                tools = tooluniverse.get_all_tools()
-                current_count = len(tools)
-                embedding_count = len(self.tool_desc_embedding)
-                
-                if current_count != embedding_count:
-                    logger.warning(f"Tool count mismatch (tools: {current_count}, embeddings: {embedding_count})")
-                    
-                    if current_count < embedding_count:
-                        self.tool_desc_embedding = self.tool_desc_embedding[:current_count]
-                        logger.info(f"Truncated embeddings to match {current_count} tools")
-                    else:
-                        last_embedding = self.tool_desc_embedding[-1]
-                        padding = [last_embedding] * (current_count - embedding_count)
-                        self.tool_desc_embedding = torch.cat(
-                            [self.tool_desc_embedding] + padding
-                        )
-                        logger.info(f"Padded embeddings to match {current_count} tools")
-                
-                return True
-                
-            except Exception as e:
-                logger.error(f"Failed to load embeddings: {str(e)}")
-                return False
-        
-        # Apply the patch
-        ToolRAGModel.load_tool_desc_embedding = patched_load
-        logger.info("Successfully patched embedding loading")
-        
-    except Exception as e:
-        logger.error(f"Failed to patch embedding loading: {str(e)}")
-        raise
+        try:
+            # Try with the safe_globals context manager
+            with torch.serialization.safe_globals([torch.serialization._reconstruct]):
+                return torch.load(filepath, weights_only=False)
+        except Exception as e:
+            logger.error(f"Failed to load embeddings even with safe_globals: {str(e)}")
+            return None
 
-class TxAgentApp:
+class TxAgentWrapper:
     def __init__(self):
-        self.agent = None
+        self.model = None
+        self.tokenizer = None
+        self.rag_model = None
+        self.tooluniverse = None
         self.is_initialized = False
+        self.special_tools = ['Finish', 'Tool_RAG', 'DirectResponse', 'RequireClarification']
 
     def initialize(self) -> str:
-        """Initialize the TxAgent with all required components"""
+        """Initialize the model from Hugging Face"""
         if self.is_initialized:
             return "✅ Already initialized"
         
         try:
-            # Apply our patch before initialization
-            patch_embedding_loading()
+            logger.info("Loading models from Hugging Face Hub...")
+            
+            # Initialize ToolUniverse first
+            self.tooluniverse = ToolUniverse(tool_files=CONFIG["tool_files"])
+            if hasattr(self.tooluniverse, 'load_tools'):
+                self.tooluniverse.load_tools()
+                logger.info(f"Loaded {len(self.tooluniverse.tools)} tools")
+            else:
+                logger.error("ToolUniverse doesn't have load_tools method")
+                return "❌ Failed to load tools"
             
-            logger.info("Initializing TxAgent...")
-            self.agent = TxAgent(
-                model_name=CONFIG["model_name"],
-                rag_model_name=CONFIG["rag_model_name"],
-                tool_files_dict=CONFIG["tool_files"],
-                force_finish=True,
-                enable_checker=True,
-                step_rag_num=10,
-                seed=100,
-                additional_default_tools=["DirectResponse", "RequireClarification"]
+            # Load main model
+            self.tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_name"])
+            self.model = AutoModelForCausalLM.from_pretrained(
+                CONFIG["model_name"],
+                device_map="auto",
+                torch_dtype=torch.float16
             )
             
-            logger.info("Loading models...")
-            self.agent.init_model()
+            # Load embeddings if file exists
+            if os.path.exists(CONFIG["embedding_filename"]):
+                self.rag_model = safe_load_embeddings(CONFIG["embedding_filename"])
+                if self.rag_model is None:
+                    return "❌ Failed to load embeddings"
             
             self.is_initialized = True
-            return "✅ TxAgent initialized successfully"
+            return "✅ Model initialized successfully"
             
         except Exception as e:
             logger.error(f"Initialization failed: {str(e)}")
             return f"❌ Initialization failed: {str(e)}"
 
     def chat(self, message: str, history: List[List[str]]) -> List[List[str]]:
-        """
-        Handle chat interactions with the TxAgent
-        
-        Args:
-            message: User input message
-            history: Chat history in format [[user_msg, bot_msg], ...]
-            
-        Returns:
-            Updated chat history
-        """
+        """Handle chat interactions with the model"""
         if not self.is_initialized:
             return history + [["", "⚠️ Please initialize the model first"]]
         
         try:
-            # Convert history to the format TxAgent expects
-            tx_history = []
-            for user_msg, bot_msg in history:
-                tx_history.append({"role": "user", "content": user_msg})
-                if bot_msg:  # Only add bot response if it exists
-                    tx_history.append({"role": "assistant", "content": bot_msg})
+            if len(message) <= 10:
+                return history + [["", "Please provide a more detailed question (at least 10 characters)"]]
+            
+            # Prepare tools prompt
+            tools_prompt = self._prepare_tools_prompt(message)
+            
+            # Format conversation
+            conversation = [
+                {"role": "system", "content": "You are a helpful assistant that will solve problems through detailed, step-by-step reasoning."},
+                *self._format_history(history),
+                {"role": "user", "content": message}
+            ]
             
             # Generate response
-            response = ""
-            for chunk in self.agent.run_gradio_chat(
-                message=message,
-                history=tx_history,
-                temperature=0.3,
+            inputs = self.tokenizer.apply_chat_template(
+                conversation,
+                add_generation_prompt=True,
+                return_tensors="pt"
+            ).to(self.model.device)
+            
+            outputs = self.model.generate(
+                inputs,
                 max_new_tokens=1024,
-                max_token=8192,
-                call_agent=False,
-                conversation=None,
-                max_round=30
-            ):
-                response = chunk  # Get the final response
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+            
+            # Decode and clean response
+            response = self.tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+            response = response.split("[TOOL_CALLS]")[0].strip()
             
-            # Format response for Gradio Chatbot
             return history + [[message, response]]
-                
+            
         except Exception as e:
             logger.error(f"Chat error: {str(e)}")
             return history + [["", f"Error: {str(e)}"]]
 
+    def _prepare_tools_prompt(self, message: str) -> str:
+        """Prepare the tools prompt section"""
+        if not hasattr(self.tooluniverse, 'tools'):
+            return ""
+            
+        tools_prompt = "\n\nYou have access to the following tools:\n"
+        for tool in self.tooluniverse.tools:
+            if tool['name'] not in self.special_tools:
+                tools_prompt += f"- {tool['name']}: {tool['description']}\n"
+        
+        # Add special tools
+        tools_prompt += "\nSpecial tools:\n"
+        tools_prompt += "- Finish: Use when you have the final answer\n"
+        tools_prompt += "- Tool_RAG: Search for additional tools when needed\n"
+        
+        return tools_prompt
+
+    def _format_history(self, history: List[List[str]]) -> List[Dict[str, str]]:
+        """Format chat history for the model"""
+        formatted = []
+        for user_msg, bot_msg in history:
+            formatted.append({"role": "user", "content": user_msg})
+            if bot_msg:
+                formatted.append({"role": "assistant", "content": bot_msg})
+        return formatted
+
 def create_interface() -> gr.Blocks:
     """Create the Gradio interface"""
-    app = TxAgentApp()
+    agent = TxAgentWrapper()
     
     with gr.Blocks(
         title="TxAgent",
@@ -189,7 +187,7 @@ def create_interface() -> gr.Blocks:
     ) as demo:
         gr.Markdown("""
         # 🧠 TxAgent: Therapeutic Reasoning AI
-        ### (Using pre-loaded embeddings)
+        ### (Loading from Hugging Face Hub)
         """)
         
         with gr.Row():
@@ -212,9 +210,8 @@ def create_interface() -> gr.Blocks:
             inputs=msg
         )
         
-        def wrapper_initialize() -> tuple:
-            """Wrapper for initialization with UI updates"""
-            status = app.initialize()
+        def wrapper_initialize():
+            status = agent.initialize()
             return status, gr.update(interactive=False)
         
         init_btn.click(
@@ -223,7 +220,7 @@ def create_interface() -> gr.Blocks:
         )
         
         msg.submit(
-            fn=app.chat,
+            fn=agent.chat,
             inputs=[msg, chatbot],
             outputs=chatbot
         ).then(