Update app.py

app.py

@@ -11,17 +11,12 @@ import shutil
 import time
 from functools import lru_cache
 
-# Environment and path setup
-current_dir = os.path.dirname(__file__)
-src_path = os.path.abspath(os.path.join(current_dir, "src"))
-
-print(">> Adding to path:", src_path)
+# ✅ Add src to Python path
+src_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src"))
+print(f"Adding to path: {src_path}")
 sys.path.insert(0, src_path)
 
-# Now import
-
-
-# Configure cache directories
+# ✅ Configure Hugging Face and cache dirs
 base_dir = "/data"
 model_cache_dir = os.path.join(base_dir, "txagent_models")
 tool_cache_dir = os.path.join(base_dir, "tool_cache")
@@ -31,14 +26,14 @@ os.makedirs(model_cache_dir, exist_ok=True)
 os.makedirs(tool_cache_dir, exist_ok=True)
 os.makedirs(file_cache_dir, exist_ok=True)
 
-os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
 os.environ["HF_HOME"] = model_cache_dir
+os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
 
 from txagent.txagent import TxAgent
 
-# Utility functions
+# ✅ Utils
 def sanitize_utf8(text: str) -> str:
     return text.encode("utf-8", "ignore").decode("utf-8")
 
@@ -48,8 +43,7 @@ def file_hash(path: str) -> str:
 
 @lru_cache(maxsize=100)
 def get_cached_response(prompt: str, file_hash: str) -> Optional[str]:
-    """Cache for frequent queries"""
-    return None  # Implement actual cache lookup if needed
+    return None
 
 def convert_file_to_json(file_path: str, file_type: str) -> str:
     try:
@@ -90,7 +84,6 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
         return json.dumps({"error": f"Error reading {os.path.basename(file_path)}: {str(e)}"})
 
 def convert_files_to_json_parallel(uploaded_files: list) -> str:
-    """Process files in parallel using ThreadPool"""
     extracted_text = []
     with ThreadPoolExecutor(max_workers=4) as executor:
         futures = []
@@ -100,14 +93,12 @@ def convert_files_to_json_parallel(uploaded_files: list) -> str:
             path = file.name
             ext = path.split(".")[-1].lower()
             futures.append(executor.submit(convert_file_to_json, path, ext))
-        
+
         for future in as_completed(futures):
             extracted_text.append(sanitize_utf8(future.result()))
     return "\n".join(extracted_text)
 
 def init_agent():
-    """Initialize the TxAgent with optimized settings"""
-    # Copy default tool file if needed
     default_tool_path = os.path.abspath("data/new_tool.json")
     target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
     if not os.path.exists(target_tool_path):
@@ -115,20 +106,16 @@ def init_agent():
 
     model_name = "mims-harvard/TxAgent-T1-Llama-3.1-8B"
     rag_model_name = "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B"
-    
+
     agent = TxAgent(
         model_name=model_name,
         rag_model_name=rag_model_name,
         tool_files_dict={"new_tool": target_tool_path},
         force_finish=True,
         enable_checker=True,
-        step_rag_num=8,  # Reduced from 10
+        step_rag_num=8,
         seed=100,
-        additional_default_tools=[],
-        torch_dtype="auto",
-        device_map="auto",
-        load_in_4bit=False,
-        load_in_8bit=False
+        additional_default_tools=[]
     )
     agent.init_model()
     return agent
@@ -154,12 +141,9 @@ def create_ui(agent: TxAgent):
                 history.append({"role": "assistant", "content": "⏳ Processing your request..."})
                 yield history
 
-                # File processing with timing
-                file_process_time = time.time()
                 extracted_text = ""
                 if uploaded_files and isinstance(uploaded_files, list):
                     extracted_text = convert_files_to_json_parallel(uploaded_files)
-                print(f"File processing took: {time.time() - file_process_time:.2f}s")
 
                 context = (
                     "You are an expert clinical AI assistant. Review this patient's history, "
@@ -168,18 +152,16 @@ def create_ui(agent: TxAgent):
                 )
                 chunked_prompt = f"{context}\n\n--- Patient Record ---\n{extracted_text}\n\n[Final Analysis]"
 
-                # Model processing with timing
-                model_start = time.time()
                 generator = agent.run_gradio_chat(
                     message=chunked_prompt,
                     history=[],
                     temperature=0.3,
-                    max_new_tokens=768,  # Reduced from 1024
-                    max_token=4096,      # Reduced from 8192
+                    max_new_tokens=768,
+                    max_token=4096,
                     call_agent=False,
                     conversation=conversation,
                     uploaded_files=uploaded_files,
-                    max_round=10       # Reduced from 30
+                    max_round=10
                 )
 
                 final_response = []
@@ -190,14 +172,12 @@ def create_ui(agent: TxAgent):
                         final_response.append(update)
                     elif isinstance(update, list):
                         final_response.extend(msg.content for msg in update if hasattr(msg, 'content'))
-                    
-                    # Yield intermediate results periodically
-                    if len(final_response) % 3 == 0:  # More frequent updates
+
+                    if len(final_response) % 3 == 0:
                         history[-1] = {"role": "assistant", "content": "".join(final_response).strip()}
                         yield history
 
                 history[-1] = {"role": "assistant", "content": "".join(final_response).strip() or "❌ No response."}
-                print(f"Model processing took: {time.time() - model_start:.2f}s")
                 yield history
 
             except Exception as chat_error:
@@ -220,27 +200,9 @@ def create_ui(agent: TxAgent):
     return demo
 
 if __name__ == "__main__":
-    # Initialize agent and warm it up
     print("Initializing agent...")
     agent = init_agent()
-    
-    # Warm-up call
-    print("Performing warm-up call...")
-    try:
-        warm_up = agent.run_gradio_chat(
-            message="Warm up",
-            history=[],
-            temperature=0.1,
-            max_new_tokens=10,
-            max_token=100,
-            call_agent=False
-        )
-        for _ in warm_up:
-            pass
-    except:
-        pass
-    
-    # Launch Gradio interface
+
     print("Launching interface...")
     demo = create_ui(agent)
     demo.queue(concurrency_count=3).launch(
@@ -248,4 +210,4 @@ if __name__ == "__main__":
         server_port=7860,
         show_error=True,
         share=True
-    )
+    )