Update app.py

app.py

@@ -47,7 +47,8 @@ MEDICAL_KEYWORDS = {
 }
 TOKENIZER = "cl100k_base"
 MAX_MODEL_LEN = 2048  # Matches your model's actual limit
-TARGET_CHUNK_TOKENS = 1500  # Leaves room for prompt and response
+TARGET_CHUNK_TOKENS = 1200  # Reduced to ensure room for prompt and response
+PROMPT_RESERVE = 300  # Tokens reserved for prompt structure
 MEDICAL_SECTION_HEADER = "=== MEDICAL SECTION ==="
 
 def sanitize_utf8(text: str) -> str:
@@ -115,7 +116,6 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
                 "status": "complete"
             })
         elif file_type == "csv":
-            # Read CSV in chunks to handle large files
             chunks = []
             for chunk in pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str,
                                    skip_blank_lines=False, on_bad_lines="skip", chunksize=1000):
@@ -146,32 +146,13 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
     except Exception as e:
         return json.dumps({"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"})
 
-def log_system_usage(tag=""):
-    """Log system resource usage."""
-    try:
-        cpu = psutil.cpu_percent(interval=1)
-        mem = psutil.virtual_memory()
-        print(f"[{tag}] CPU: {cpu}% | RAM: {mem.used // (1024**2)}MB / {mem.total // (1024**2)}MB")
-        result = subprocess.run(
-            ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", "--format=csv,nounits,noheader"],
-            capture_output=True, text=True
-        )
-        if result.returncode == 0:
-            used, total, util = result.stdout.strip().split(", ")
-            print(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%")
-    except Exception as e:
-        print(f"[{tag}] GPU/CPU monitor failed: {e}")
-
 def clean_response(text: str) -> str:
     """Clean and format the model response."""
     text = sanitize_utf8(text)
-    # Remove tool calls and JSON artifacts
     text = re.sub(r"\[TOOL_CALLS\].*", "", text, flags=re.DOTALL)
     text = re.sub(r"\['get_[^\]]+\']\n?", "", text)
     text = re.sub(r"\{'meta':\s*\{.*?\}\s*,\s*'results':\s*\[.*?\]\}\n?", "", text, flags=re.DOTALL)
-    # Remove repetitive phrases
     text = re.sub(r"To analyze the medical records for clinical oversights.*?begin by reviewing.*?\n", "", text, flags=re.DOTALL)
-    # Collapse excessive newlines
     text = re.sub(r"\n{3,}", "\n\n", text).strip()
     return text
 
@@ -183,7 +164,6 @@ def format_final_report(analysis_results: List[str], filename: str) -> str:
     report.append(f"File: {filename}")
     report.append("=" * 80)
     
-    # Extract sections from all chunks
     sections = {
         "CRITICAL FINDINGS": [],
         "MISSED DIAGNOSES": [],
@@ -194,7 +174,6 @@ def format_final_report(analysis_results: List[str], filename: str) -> str:
     
     for result in analysis_results:
         for section in sections:
-            # Find section content using regex
             section_match = re.search(
                 rf"{re.escape(section)}:?\s*\n([^*]+?)(?=\n\*|\n\n|$)", 
                 result, 
@@ -205,13 +184,11 @@ def format_final_report(analysis_results: List[str], filename: str) -> str:
                 if content and content not in sections[section]:
                     sections[section].append(content)
     
-    # Build the final report - prioritize critical findings
     if sections["CRITICAL FINDINGS"]:
         report.append("\n🚨 **CRITICAL FINDINGS** 🚨")
         for content in sections["CRITICAL FINDINGS"]:
             report.append(f"\n{content}")
     
-    # Add other sections
     for section, contents in sections.items():
         if section != "CRITICAL FINDINGS" and contents:
             report.append(f"\n**{section.upper()}**")
@@ -236,7 +213,6 @@ def split_content_by_tokens(content: str, max_tokens: int = TARGET_CHUNK_TOKENS)
     for para in paragraphs:
         para_tokens = count_tokens(para)
         if para_tokens > max_tokens:
-            # Handle very long paragraphs by splitting sentences
             sentences = re.split(r'(?<=[.!?])\s+', para)
             for sent in sentences:
                 sent_tokens = count_tokens(sent)
@@ -286,37 +262,28 @@ def init_agent():
     return agent
 
 def analyze_complete_document(content: str, filename: str, agent: TxAgent) -> str:
-    """Analyze complete document with proper chunking and token management"""
+    """Analyze complete document with strict token management"""
     chunks = split_content_by_tokens(content)
     analysis_results = []
     
     for i, chunk in enumerate(chunks):
         try:
-            # Create minimal prompt to save tokens
-            prompt = f"""
-Analyze this medical record section for:
-1. Critical findings (urgent)
-2. Missed diagnoses (with evidence)
-3. Medication issues
-4. Assessment gaps
-5. Follow-up needs
-
-Content:
-{chunk}
-
-Concise findings only:
-"""
-            # Verify we're within token limits
-            prompt_tokens = count_tokens(prompt)
-            chunk_tokens = count_tokens(chunk)
+            # Ultra-minimal prompt to maximize content space
+            base_prompt = "Analyze for:\n1. Critical\n2. Missed DX\n3. Med issues\n4. Gaps\n5. Follow-up\n\nContent:\n"
+            
+            # Calculate available space for content
+            prompt_tokens = count_tokens(base_prompt)
+            max_content_tokens = MAX_MODEL_LEN - prompt_tokens - 100  # Response buffer
             
-            if prompt_tokens + chunk_tokens > MAX_MODEL_LEN - 512:  # Leave room for response
-                # Find a natural truncation point
+            # Ensure chunk fits
+            chunk_tokens = count_tokens(chunk)
+            if chunk_tokens > max_content_tokens:
+                # Find last paragraph that fits
                 adjusted_chunk = ""
                 tokens_used = 0
-                max_content_tokens = MAX_MODEL_LEN - prompt_tokens - 512
+                paragraphs = re.split(r"\n\s*\n", chunk)
                 
-                for para in re.split(r"\n\s*\n", chunk):
+                for para in paragraphs:
                     para_tokens = count_tokens(para)
                     if tokens_used + para_tokens <= max_content_tokens:
                         adjusted_chunk += "\n\n" + para
@@ -325,7 +292,7 @@ Concise findings only:
                         break
                 
                 if not adjusted_chunk:
-                    # If even one paragraph is too long, split sentences
+                    # If even one paragraph is too big, split sentences
                     sentences = re.split(r'(?<=[.!?])\s+', chunk)
                     for sent in sentences:
                         sent_tokens = count_tokens(sent)
@@ -337,12 +304,14 @@ Concise findings only:
                 
                 chunk = adjusted_chunk.strip()
             
+            prompt = base_prompt + chunk
+            
             response = ""
             for output in agent.run_gradio_chat(
                 message=prompt,
                 history=[],
                 temperature=0.1,
-                max_new_tokens=512,  # Keep responses concise
+                max_new_tokens=300,  # Keep responses very concise
                 max_token=MAX_MODEL_LEN,
                 call_agent=False,
                 conversation=[],
@@ -407,7 +376,6 @@ def create_ui(agent):
             
             yield "", None, "⏳ Processing documents (this may take several minutes for large files)..."
             
-            # Process all files completely
             file_contents = []
             filenames = []
             total_tokens = 0
@@ -445,14 +413,12 @@ def create_ui(agent):
             yield "", None, f"🔍 Analyzing content ({total_tokens//1000}k tokens)..."
             
             try:
-                # Process the complete document
                 full_report = analyze_complete_document(
                     combined_content, 
                     combined_filename, 
                     agent
                 )
                 
-                # Save report to file
                 file_hash_value = hashlib.md5(combined_content.encode()).hexdigest()
                 report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
                 with open(report_path, "w", encoding="utf-8") as f:
@@ -465,7 +431,6 @@ def create_ui(agent):
                 print(error_msg)
                 yield "", None, error_msg
         
-        # UI event handlers
         send_btn.click(
             fn=analyze,
             inputs=[file_upload, msg_input],
@@ -483,7 +448,6 @@ def create_ui(agent):
 
 if __name__ == "__main__":
     print("🚀 Launching app...")
-    # Install tiktoken if not available
     try:
         import tiktoken
     except ImportError: