Update app.py

app.py

@@ -10,7 +10,6 @@ import re
 import psutil
 import subprocess
 from collections import defaultdict
-import torch
 
 # Persistent directory for Hugging Face Space
 persistent_dir = os.getenv("HF_HOME", "/data/hf_cache")
@@ -90,18 +89,18 @@ def log_system_usage(tag=""):
 
 def clean_response(text: str) -> str:
     text = sanitize_utf8(text)
-    # Exhaustively remove all unwanted text
+    # Remove all tool and reasoning text
     text = re.sub(r"\[TOOL_CALLS\].*|(?:get_|tool\s|retrieve\s|use\s).*?\n", "", text, flags=re.DOTALL | re.IGNORECASE)
     text = re.sub(r"\{'meta':\s*\{.*?\}\s*,\s*'results':\s*\[.*?\]\}\n?", "", text, flags=re.DOTALL)
     text = re.sub(
-        r"(?i)(to address|analyze|will\s|since\s|no\s|none|previous|attempt|involve|check\s|explore|manually|"
+        r"(?i)(to\s|analyze|will\s|since\s|no\s|none|previous|attempt|involve|check\s|explore|manually|"
         r"start|look|use|focus|retrieve|tool|based\s|overall|indicate|mention|consider|ensure|need\s|"
         r"provide|review|assess|identify|potential|records|patient|history|symptoms|medication|"
-        r"conflict|assessment|follow-up|issue|reasoning|step).*?\n",
+        r"conflict|assessment|follow-up|issue|reasoning|step|prompt|address).*?\n",
         "", text, flags=re.DOTALL
     )
     text = re.sub(r"\n{3,}", "\n\n", text).strip()
-    # Only keep lines under headings or bullet points
+    # Only keep heading lines and bullets
     lines = []
     valid_heading = False
     for line in text.split("\n"):
@@ -116,7 +115,7 @@ def clean_response(text: str) -> str:
     return "\n".join(lines).strip()
 
 def consolidate_findings(responses: List[str]) -> str:
-    # Merge findings, keeping only unique points
+    # Merge unique findings
     findings = defaultdict(set)
     headings = ["Missed Diagnoses", "Medication Conflicts", "Incomplete Assessments", "Urgent Follow-up"]
     
@@ -133,7 +132,7 @@ def consolidate_findings(responses: List[str]) -> str:
             elif current_heading and line.startswith("-"):
                 findings[current_heading].add(line)
     
-    # Format final output
+    # Format output
     output = []
     for heading in headings:
         if findings[heading]:
@@ -149,11 +148,10 @@ def init_agent():
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
         force_finish=True,
         enable_checker=True,
-        step_rag_num=1,
+        step_rag_num=0,  # Disable RAG to prevent tool artifacts
         seed=100,
     )
-    # Enable FP16 for A100
-    agent.init_model(dtype=torch.float16)
+    agent.init_model()  # No dtype argument
     log_system_usage("After Load")
     print("✅ Agent Ready")
     return agent
@@ -181,14 +179,14 @@ def create_ui(agent):
                     extracted = "\n".join(results)
                     file_hash_value = file_hash(files[0].name) if files else ""
 
-            # Split into tiny chunks of 1,000 characters
+            # Tiny chunks for speed
             chunk_size = 1000
             chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
             chunk_responses = []
-            batch_size = 4  # Process 4 chunks at a time on A100
+            batch_size = 4  # Batch for A100
 
             prompt_template = """
-Output only oversights under these headings, one brief point each. No tools, reasoning, or extra text.
+Output only oversights under these headings, one point each. No tools, reasoning, or text beyond headings and bullets.
 
 **Missed Diagnoses**:
 **Medication Conflicts**:
@@ -200,7 +198,7 @@ Records:
 """
 
             try:
-                # Process chunks in batches
+                # Batch process chunks
                 for i in range(0, len(chunks), batch_size):
                     batch = chunks[i:i + batch_size]
                     batch_responses = []
@@ -212,7 +210,7 @@ Records:
                             history=[],
                             temperature=0.1,
                             max_new_tokens=128,
-                            max_token=8192,
+                            max_token=4096,  # Revert to 4096 as 8192 may not be supported
                             call_agent=False,
                             conversation=[],
                         ):
@@ -232,12 +230,12 @@ Records:
                             batch_responses.append(chunk_response)
                     chunk_responses.extend(batch_responses)
 
-                # Consolidate into one final result
+                # Single final output
                 final_response = consolidate_findings(chunk_responses)
                 history[-1]["content"] = final_response
                 yield history, None
 
-                # Generate report file
+                # Report file
                 report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
                 if report_path and final_response != "No oversights identified.":
                     with open(report_path, "w", encoding="utf-8") as f: