Create app.py

app.py

@@ -0,0 +1,158 @@
+import os
+import logging
+import torch
+import gradio as gr
+from txagent import TxAgent
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Configuration
+MODEL_NAME = "mims-harvard/TxAgent-T1-Llama-3.1-8B"
+RAG_MODEL_NAME = "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B"
+TOOL_FILE = "data/new_tool.json"
+
+# Environment setup
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+os.environ["CUDA_MODULE_LOADING"] = "LAZY"
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+class TxAgentSystem:
+    def __init__(self):
+        self.agent = None
+        self.is_initialized = False
+        self.examples = [
+            ["A 68-year-old with CKD prescribed metformin. Safe for renal clearance?"],
+            ["30-year-old on Prozac diagnosed with WHIM. Safe to take Xolremdi?"]
+        ]
+        
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA is not available - GPU required")
+
+        logger.info(f"GPU: {torch.cuda.get_device_name(0)}")
+        logger.info(f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")
+
+        self._initialize_system()
+
+    def _initialize_system(self):
+        try:
+            os.makedirs("data", exist_ok=True)
+            if not os.path.exists(TOOL_FILE):
+                with open(TOOL_FILE, "w") as f:
+                    f.write("[]")
+
+            logger.info("Initializing TxAgent...")
+            
+            # Initialize with RAG disabled first
+            try:
+                self.agent = TxAgent(
+                    model_name=MODEL_NAME,
+                    rag_model_name=RAG_MODEL_NAME,
+                    tool_files_dict={"new_tool": TOOL_FILE},
+                    force_finish=True,
+                    enable_checker=True,
+                    step_rag_num=10,
+                    seed=100,
+                    enable_rag=True
+                )
+            except Exception as e:
+                logger.warning(f"Failed to initialize with RAG: {str(e)}")
+                logger.info("Retrying without RAG...")
+                self.agent = TxAgent(
+                    model_name=MODEL_NAME,
+                    rag_model_name=None,
+                    tool_files_dict={"new_tool": TOOL_FILE},
+                    force_finish=True,
+                    enable_checker=True,
+                    step_rag_num=0,
+                    seed=100,
+                    enable_rag=False
+                )
+
+            logger.info("Loading main model...")
+            self.agent.init_model()
+
+            self.is_initialized = True
+            logger.info("System initialization completed successfully")
+
+        except Exception as e:
+            logger.error(f"System initialization failed: {str(e)}")
+            self.is_initialized = False
+            raise
+
+    def chat_fn(self, message, history, temperature, max_tokens, rag_depth):
+        if not self.is_initialized:
+            return "", history + [(message, "System initialization failed. Please check logs.")]
+
+        try:
+            response = self.agent.run_gradio_chat(
+                message=message,
+                history=history,
+                temperature=temperature,
+                max_new_tokens=max_tokens,
+                max_total_tokens=16384,
+                enable_multi_agent=False,
+                conv_history=history,
+                max_steps=rag_depth,
+                seed=100
+            )
+            new_history = history + [(message, response)]
+            return "", new_history
+
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            return "", history + [(message, "⚠️ GPU memory overflow. Please try a shorter query.")]
+
+        except Exception as e:
+            logger.error(f"Chat error: {str(e)}")
+            return "", history + [(message, f"🚨 Error: {str(e)}")]
+
+    def launch_ui(self):
+        with gr.Blocks(theme=gr.themes.Soft(), title="TxAgent Medical AI") as demo:
+            gr.Markdown("## 🧠 TxAgent (A100/H100 Optimized)")
+
+            status = gr.Textbox(
+                value="✅ System ready" if self.is_initialized else "❌ Initialization failed",
+                label="System Status",
+                interactive=False
+            )
+
+            with gr.Row():
+                with gr.Column(scale=3):
+                    chatbot = gr.Chatbot(height=600, label="Conversation History")
+                    msg = gr.Textbox(label="Enter Medical Query", placeholder="Type your question here...")
+                with gr.Column(scale=1):
+                    temp = gr.Slider(0.1, 1.0, value=0.7, label="Temperature")
+                    max_tokens = gr.Slider(128, 8192, value=2048, label="Max Response Tokens")
+                    rag_depth = gr.Slider(1, 20, value=10, label="RAG Depth")
+                    clear_btn = gr.Button("Clear History")
+
+            gr.Examples(
+                examples=self.examples,
+                inputs=msg,
+                label="Example Queries"
+            )
+
+            msg.submit(
+                self.chat_fn,
+                inputs=[msg, chatbot, temp, max_tokens, rag_depth],
+                outputs=[msg, chatbot]
+            )
+            clear_btn.click(lambda: None, None, chatbot, queue=False)
+
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860
+        )
+
+if __name__ == "__main__":
+    try:
+        system = TxAgentSystem()
+        system.launch_ui()
+    except Exception as e:
+        logger.critical(f"Fatal error: {str(e)}")
+        raise