Update app.py

app.py

@@ -11,14 +11,13 @@ import shutil
 import time
 from functools import lru_cache
 
+# Environment and path setup
 current_dir = os.path.dirname(os.path.abspath(__file__))
-
-# ✅ Add src to Python path
 src_path = os.path.abspath(os.path.join(current_dir, "src"))
 print(f"Adding to path: {src_path}")
 sys.path.insert(0, src_path)
 
-# ✅ Configure Hugging Face and cache dirs
+# Configure cache directories
 base_dir = "/data"
 model_cache_dir = os.path.join(base_dir, "txagent_models")
 tool_cache_dir = os.path.join(base_dir, "tool_cache")
@@ -28,14 +27,13 @@ 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["HF_HOME"] = model_cache_dir
 os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
+os.environ["HF_HOME"] = model_cache_dir
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
 
 from txagent.txagent import TxAgent
 
-# ✅ Utils
 def sanitize_utf8(text: str) -> str:
     return text.encode("utf-8", "ignore").decode("utf-8")
 
@@ -56,8 +54,7 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
             return open(cache_path, "r", encoding="utf-8").read()
 
         if file_type == "csv":
-            df = pd.read_csv(file_path, encoding_errors="replace", header=None, 
-                            dtype=str, skip_blank_lines=False, on_bad_lines="skip")
+            df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str, skip_blank_lines=False, on_bad_lines="skip")
         elif file_type in ["xls", "xlsx"]:
             try:
                 df = pd.read_excel(file_path, engine="openpyxl", header=None, dtype=str)
@@ -143,9 +140,11 @@ def create_ui(agent: TxAgent):
                 history.append({"role": "assistant", "content": "⏳ Processing your request..."})
                 yield history
 
+                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, "
@@ -154,6 +153,7 @@ def create_ui(agent: TxAgent):
                 )
                 chunked_prompt = f"{context}\n\n--- Patient Record ---\n{extracted_text}\n\n[Final Analysis]"
 
+                model_start = time.time()
                 generator = agent.run_gradio_chat(
                     message=chunked_prompt,
                     history=[],
@@ -180,6 +180,7 @@ def create_ui(agent: TxAgent):
                         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:
@@ -205,9 +206,25 @@ if __name__ == "__main__":
     print("Initializing agent...")
     agent = init_agent()
 
+    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,
+            conversation=[]
+        )
+        for _ in warm_up:
+            pass
+    except:
+        pass
+
     print("Launching interface...")
     demo = create_ui(agent)
-    demo.queue(concurrency_count=3).launch(
+    demo.queue().launch(
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True,