Update app.py

app.py

@@ -4,30 +4,19 @@ import logging
 import torch
 from txagent import TxAgent
 import gradio as gr
-from huggingface_hub import snapshot_download
+from huggingface_hub import hf_hub_download
 from tooluniverse import ToolUniverse
-import time
-from requests.adapters import HTTPAdapter
-from requests import Session
-from urllib3.util.retry import Retry
 from tqdm import tqdm
 
-# Configuration
+# Configuration - Now using remote Hugging Face models
 CONFIG = {
     "model_name": "mims-harvard/TxAgent-T1-Llama-3.1-8B",
-    "rag_model_name": "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
+    "rag_model_name": "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B", 
     "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding_47dc56b3e3ddeb31af4f19defdd538d984de1500368852a0fab80bc2e826c944.pt",
-    "local_dir": "./models",
     "tool_files": {
         "new_tool": "./data/new_tool.json"
     },
-    "download_settings": {
-        "timeout": 600,  # 10 minutes per request
-        "max_retries": 5,
-        "retry_delay": 30,  # seconds between retries
-        "chunk_size": 1024 * 1024 * 10,  # 10MB chunks
-        "max_concurrent": 2  # concurrent downloads
-    }
+    "load_from_hub": True  # Flag to load directly from Hugging Face
 }
 
 # Logging setup
@@ -37,28 +26,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-def create_optimized_session():
-    """Create a session optimized for large file downloads"""
-    session = Session()
-    
-    retry_strategy = Retry(
-        total=CONFIG["download_settings"]["max_retries"],
-        backoff_factor=1,
-        status_forcelist=[408, 429, 500, 502, 503, 504]
-    )
-    
-    adapter = HTTPAdapter(
-        max_retries=retry_strategy,
-        pool_connections=10,
-        pool_maxsize=10,
-        pool_block=True
-    )
-    
-    session.mount("https://", adapter)
-    session.mount("http://", adapter)
-    
-    return session
-
 def prepare_tool_files():
     os.makedirs("./data", exist_ok=True)
     if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
@@ -69,69 +36,6 @@ def prepare_tool_files():
             json.dump(tools, f, indent=2)
         logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
 
-def download_model_with_progress(repo_id, local_dir):
-    custom_session = create_optimized_session()
-    
-    for attempt in range(CONFIG["download_settings"]["max_retries"] + 1):
-        try:
-            logger.info(f"Download attempt {attempt + 1} for {repo_id}")
-            
-            # Create progress bar
-            progress = tqdm(
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1024,
-                miniters=1,
-                desc=f"Downloading {repo_id.split('/')[-1]}"
-            )
-            
-            def update_progress(monitor):
-                progress.update(monitor.bytes_read - progress.n)
-            
-            snapshot_download(
-                repo_id=repo_id,
-                local_dir=local_dir,
-                resume_download=True,
-                local_dir_use_symlinks=False,
-                use_auth_token=True,
-                max_workers=CONFIG["download_settings"]["max_concurrent"],
-                tqdm_class=None,  # We handle progress ourselves
-                session=custom_session
-            )
-            
-            progress.close()
-            return True
-            
-        except Exception as e:
-            logger.error(f"Attempt {attempt + 1} failed: {str(e)}")
-            if attempt < CONFIG["download_settings"]["max_retries"]:
-                wait_time = CONFIG["download_settings"]["retry_delay"] * (attempt + 1)
-                logger.info(f"Waiting {wait_time} seconds before retry...")
-                time.sleep(wait_time)
-            else:
-                progress.close()
-                return False
-
-def download_model_files():
-    os.makedirs(CONFIG["local_dir"], exist_ok=True)
-    logger.info("Starting model downloads...")
-
-    # Download main model
-    if not download_model_with_progress(
-        CONFIG["model_name"],
-        os.path.join(CONFIG["local_dir"], CONFIG["model_name"])
-    ):
-        raise RuntimeError(f"Failed to download {CONFIG['model_name']}")
-
-    # Download RAG model
-    if not download_model_with_progress(
-        CONFIG["rag_model_name"],
-        os.path.join(CONFIG["local_dir"], CONFIG["rag_model_name"])
-    ):
-        raise RuntimeError(f"Failed to download {CONFIG['rag_model_name']}")
-
-    logger.info("All model files downloaded successfully")
-
 def load_embeddings(agent):
     embedding_path = CONFIG["embedding_filename"]
     if os.path.exists(embedding_path):
@@ -143,7 +47,7 @@ def load_embeddings(agent):
         except Exception as e:
             logger.error(f"Failed to load embeddings: {e}")
     
-    logger.info("Generating tool embeddings...")
+    logger.info("Generating tool embeddings from remote model...")
     try:
         tools = agent.tooluniverse.get_all_tools()
         descriptions = [tool["description"] for tool in tools]
@@ -165,33 +69,29 @@ class TxAgentApp:
             return "✅ Already initialized"
         
         try:
-            # Initialize with progress tracking
-            with tqdm(total=4, desc="Initializing TxAgent") as pbar:
-                logger.info("Creating TxAgent instance...")
-                self.agent = TxAgent(
-                    CONFIG["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"]
-                )
-                pbar.update(1)
-                
-                logger.info("Initializing models...")
-                self.agent.init_model()
-                pbar.update(1)
-                
-                logger.info("Loading embeddings...")
-                load_embeddings(self.agent)
-                pbar.update(1)
-                
-                self.is_initialized = True
-                pbar.update(1)
-                
-            return "✅ TxAgent initialized successfully"
+            logger.info("Initializing TxAgent with remote models...")
+            
+            # Initialize with remote models
+            self.agent = TxAgent(
+                CONFIG["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"],
+                local_files_only=False  # Force loading from Hugging Face Hub
+            )
+            
+            logger.info("Loading remote models...")
+            self.agent.init_model()
+            
+            logger.info("Preparing embeddings...")
+            load_embeddings(self.agent)
+            
+            self.is_initialized = True
+            return "✅ TxAgent initialized successfully (using remote models)"
         except Exception as e:
             logger.error(f"Initialization failed: {str(e)}")
             return f"❌ Initialization failed: {str(e)}"
@@ -226,37 +126,30 @@ def create_interface():
         title="TxAgent",
         css="""
         .gradio-container {max-width: 900px !important}
-        .progress-bar {height: 20px !important}
         """
     ) as demo:
         gr.Markdown("""
-        # � TxAgent: Therapeutic Reasoning AI
-        ### Specialized for clinical decision support
+        # 🧠 TxAgent: Therapeutic Reasoning AI
+        ### (Running with remote Hugging Face models)
         """)
         
-        # Initialization section
         with gr.Row():
             init_btn = gr.Button("Initialize Model", variant="primary")
             init_status = gr.Textbox(label="Status", interactive=False)
-            download_progress = gr.Textbox(visible=False)
         
-        # Chat interface
         chatbot = gr.Chatbot(height=500, label="Conversation")
-        msg = gr.Textbox(label="Your clinical question", placeholder="Ask about drug interactions, dosing, etc...")
+        msg = gr.Textbox(label="Your clinical question")
         clear_btn = gr.Button("Clear Chat")
         
-        # Examples
         gr.Examples(
             examples=[
                 "How to adjust Journavx for renal impairment?",
                 "Xolremdi and Prozac interaction in WHIM syndrome?",
                 "Alternative to Warfarin for patient with amiodarone?"
             ],
-            inputs=msg,
-            label="Example Questions"
+            inputs=msg
         )
         
-        # Event handlers
         init_btn.click(
             fn=app.initialize,
             outputs=init_status
@@ -277,14 +170,11 @@ def create_interface():
 
 if __name__ == "__main__":
     try:
-        logger.info("Starting application setup...")
+        logger.info("Starting application...")
         
-        # Prepare files
+        # Prepare local tool files
         prepare_tool_files()
         
-        # Download models with progress tracking
-        download_model_files()
-        
         # Launch interface
         interface = create_interface()
         interface.launch(