Update app.py

app.py

@@ -10,9 +10,10 @@ import re
 import psutil
 import subprocess
 from collections import defaultdict
+import torch
 
-# Persistent directory
-persistent_dir = "/data/hf_cache"
+# Persistent directory for Hugging Face Space
+persistent_dir = os.getenv("HF_HOME", "/data/hf_cache")
 os.makedirs(persistent_dir, exist_ok=True)
 
 model_cache_dir = os.path.join(persistent_dir, "txagent_models")
@@ -89,47 +90,55 @@ def log_system_usage(tag=""):
 
 def clean_response(text: str) -> str:
     text = sanitize_utf8(text)
-    # Remove all tool-related and reasoning text
-    text = re.sub(r"\[TOOL_CALLS\].*|(?:get_|tool\s|retrieve\s).*?\n", "", text, flags=re.DOTALL | re.IGNORECASE)
+    # Exhaustively remove all unwanted 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 the|will (start|look|use|focus)|since the|no (drug|clinical|information)|none|previous|attempt|involve|check for|explore|manually).*?\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"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",
+        "", text, flags=re.DOTALL
+    )
     text = re.sub(r"\n{3,}", "\n\n", text).strip()
-    # Only keep text under specific headings
-    if not re.search(r"^(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", text, re.MULTILINE | re.IGNORECASE):
-        return ""
-    return text
+    # Only keep lines under headings or bullet points
+    lines = []
+    valid_heading = False
+    for line in text.split("\n"):
+        line = line.strip()
+        if line.lower() in ["missed diagnoses:", "medication conflicts:", "incomplete assessments:", "urgent follow-up:"]:
+            valid_heading = True
+            lines.append(f"**{line[:-1]}**:")
+        elif valid_heading and line.startswith("-"):
+            lines.append(line)
+        else:
+            valid_heading = False
+    return "\n".join(lines).strip()
 
 def consolidate_findings(responses: List[str]) -> str:
-    # Aggregate findings under each heading, removing duplicates
+    # Merge findings, keeping only unique points
     findings = defaultdict(set)
     headings = ["Missed Diagnoses", "Medication Conflicts", "Incomplete Assessments", "Urgent Follow-up"]
     
     for response in responses:
         if not response:
             continue
-        # Split response into sections by heading
         current_heading = None
-        current_points = []
         for line in response.split("\n"):
             line = line.strip()
             if not line:
                 continue
-            if any(line.lower().startswith(h.lower()) for h in headings):
-                if current_heading and current_points:
-                    findings[current_heading].update(current_points)
-                current_heading = next(h for h in headings if line.lower().startswith(h.lower()))
-                current_points = []
+            if line.lower().startswith(tuple(h.lower() + ":" for h in headings)):
+                current_heading = next(h for h in headings if line.lower().startswith(h.lower() + ":"))
             elif current_heading and line.startswith("-"):
-                current_points.append(line)
-        if current_heading and current_points:
-            findings[current_heading].update(current_points)
+                findings[current_heading].add(line)
     
-    # Format consolidated output
+    # Format final output
     output = []
     for heading in headings:
         if findings[heading]:
             output.append(f"**{heading}**:")
-            output.extend(sorted(findings[heading]))
+            output.extend(sorted(findings[heading], key=lambda x: x.lower()))
     return "\n".join(output).strip() if output else "No oversights identified."
 
 def init_agent():
@@ -143,7 +152,8 @@ def init_agent():
         step_rag_num=1,
         seed=100,
     )
-    agent.init_model()
+    # Enable FP16 for A100
+    agent.init_model(dtype=torch.float16)
     log_system_usage("After Load")
     print("✅ Agent Ready")
     return agent
@@ -171,13 +181,14 @@ def create_ui(agent):
                     extracted = "\n".join(results)
                     file_hash_value = file_hash(files[0].name) if files else ""
 
-            # Split into small chunks of 1,500 characters
-            chunk_size = 1500
+            # Split into tiny chunks of 1,000 characters
+            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
 
             prompt_template = """
-List doctor oversights under these headings only, with one brief point each. No tools or reasoning steps.
+Output only oversights under these headings, one brief point each. No tools, reasoning, or extra text.
 
 **Missed Diagnoses**:
 **Medication Conflicts**:
@@ -189,35 +200,39 @@ Records:
 """
 
             try:
-                # Process all chunks, collecting responses
-                for chunk in chunks:
-                    prompt = prompt_template.format(chunk=chunk)
-                    chunk_response = ""
-                    for output in agent.run_gradio_chat(
-                        message=prompt,
-                        history=[],
-                        temperature=0.1,
-                        max_new_tokens=256,
-                        max_token=4096,
-                        call_agent=False,
-                        conversation=[],
-                    ):
-                        if output is None:
-                            continue
-                        if isinstance(output, list):
-                            for m in output:
-                                if hasattr(m, 'content') and m.content:
-                                    cleaned = clean_response(m.content)
-                                    if cleaned:
-                                        chunk_response += cleaned + "\n"
-                        elif isinstance(output, str) and output.strip():
-                            cleaned = clean_response(output)
-                            if cleaned:
-                                chunk_response += cleaned + "\n"
-                    if chunk_response:
-                        chunk_responses.append(chunk_response)
-
-                # Consolidate all responses into one final output
+                # Process chunks in batches
+                for i in range(0, len(chunks), batch_size):
+                    batch = chunks[i:i + batch_size]
+                    batch_responses = []
+                    for chunk in batch:
+                        prompt = prompt_template.format(chunk=chunk)
+                        chunk_response = ""
+                        for output in agent.run_gradio_chat(
+                            message=prompt,
+                            history=[],
+                            temperature=0.1,
+                            max_new_tokens=128,
+                            max_token=8192,
+                            call_agent=False,
+                            conversation=[],
+                        ):
+                            if output is None:
+                                continue
+                            if isinstance(output, list):
+                                for m in output:
+                                    if hasattr(m, 'content') and m.content:
+                                        cleaned = clean_response(m.content)
+                                        if cleaned:
+                                            chunk_response += cleaned + "\n"
+                            elif isinstance(output, str) and output.strip():
+                                cleaned = clean_response(output)
+                                if cleaned:
+                                    chunk_response += cleaned + "\n"
+                        if chunk_response:
+                            batch_responses.append(chunk_response)
+                    chunk_responses.extend(batch_responses)
+
+                # Consolidate into one final result
                 final_response = consolidate_findings(chunk_responses)
                 history[-1]["content"] = final_response
                 yield history, None