Update app.py

app.py

@@ -46,11 +46,27 @@ MEDICAL_KEYWORDS = {
     'conclusion', 'history', 'examination', 'progress', 'discharge'
 }
 TOKENIZER = "cl100k_base"
-MAX_MODEL_LEN = 2048  # Matches your model's actual limit
-TARGET_CHUNK_TOKENS = 1200  # Reduced to ensure room for prompt and response
-PROMPT_RESERVE = 300  # Tokens reserved for prompt structure
+MAX_MODEL_LEN = 2048
+TARGET_CHUNK_TOKENS = 1200
+PROMPT_RESERVE = 300
 MEDICAL_SECTION_HEADER = "=== MEDICAL SECTION ==="
 
+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 sanitize_utf8(text: str) -> str:
     """Ensure text is UTF-8 clean."""
     return text.encode("utf-8", "ignore").decode("utf-8")
@@ -82,7 +98,6 @@ def extract_all_pages_with_token_count(file_path: str) -> Tuple[str, int, int]:
                 page_text = page.extract_text() or ""
                 lower_text = page_text.lower()
                 
-                # Mark medical sections
                 if any(re.search(rf'\b{kw}\b', lower_text) for kw in MEDICAL_KEYWORDS):
                     section_header = f"\n{MEDICAL_SECTION_HEADER} (Page {i+1})\n"
                     text_chunks.append(section_header + page_text.strip())
@@ -236,31 +251,6 @@ def split_content_by_tokens(content: str, max_tokens: int = TARGET_CHUNK_TOKENS)
     
     return chunks
 
-def init_agent():
-    """Initialize the TxAgent with proper configuration."""
-    print("🔁 Initializing model...")
-    log_system_usage("Before Load")
-    
-    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):
-        shutil.copy(default_tool_path, target_tool_path)
-
-    agent = TxAgent(
-        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
-        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-        tool_files_dict={"new_tool": target_tool_path},
-        force_finish=True,
-        enable_checker=True,
-        step_rag_num=2,
-        seed=100,
-        additional_default_tools=[],
-    )
-    agent.init_model()
-    log_system_usage("After Load")
-    print("✅ Agent Ready")
-    return agent
-
 def analyze_complete_document(content: str, filename: str, agent: TxAgent) -> str:
     """Analyze complete document with strict token management"""
     chunks = split_content_by_tokens(content)
@@ -268,17 +258,13 @@ def analyze_complete_document(content: str, filename: str, agent: TxAgent) -> st
     
     for i, chunk in enumerate(chunks):
         try:
-            # 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
+            max_content_tokens = MAX_MODEL_LEN - prompt_tokens - 100
             
-            # Ensure chunk fits
             chunk_tokens = count_tokens(chunk)
             if chunk_tokens > max_content_tokens:
-                # Find last paragraph that fits
                 adjusted_chunk = ""
                 tokens_used = 0
                 paragraphs = re.split(r"\n\s*\n", chunk)
@@ -292,7 +278,6 @@ def analyze_complete_document(content: str, filename: str, agent: TxAgent) -> st
                         break
                 
                 if not adjusted_chunk:
-                    # If even one paragraph is too big, split sentences
                     sentences = re.split(r'(?<=[.!?])\s+', chunk)
                     for sent in sentences:
                         sent_tokens = count_tokens(sent)
@@ -311,7 +296,7 @@ def analyze_complete_document(content: str, filename: str, agent: TxAgent) -> st
                 message=prompt,
                 history=[],
                 temperature=0.1,
-                max_new_tokens=300,  # Keep responses very concise
+                max_new_tokens=300,
                 max_token=MAX_MODEL_LEN,
                 call_agent=False,
                 conversation=[],
@@ -332,6 +317,31 @@ def analyze_complete_document(content: str, filename: str, agent: TxAgent) -> st
     
     return format_final_report(analysis_results, filename)
 
+def init_agent():
+    """Initialize the TxAgent with proper configuration."""
+    print("🔁 Initializing model...")
+    log_system_usage("Before Load")
+    
+    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):
+        shutil.copy(default_tool_path, target_tool_path)
+
+    agent = TxAgent(
+        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
+        tool_files_dict={"new_tool": target_tool_path},
+        force_finish=True,
+        enable_checker=True,
+        step_rag_num=2,
+        seed=100,
+        additional_default_tools=[],
+    )
+    agent.init_model()
+    log_system_usage("After Load")
+    print("✅ Agent Ready")
+    return agent
+
 def create_ui(agent):
     """Create the Gradio interface."""
     with gr.Blocks(theme=gr.themes.Soft(), title="Clinical Oversight Assistant") as demo: