Update app.py

app.py

@@ -1,144 +1,172 @@
-import os
-import time
 import gradio as gr
+import torch
+import json
+import time
+from PIL import Image
+from transformers import AutoProcessor, AutoModel
 from huggingface_hub import InferenceClient
-from mann_engram_en.router import MANNEngramRouter
 
-# 环境配置
-os.environ["HF_HOME"] = "/home/user/.cache/huggingface"
-router_engine = None
+# ==========================================
+# Phase 2: Edge-Side Tensor Router (SiGLIP)
+# ==========================================
+class SiGLIPRouter:
+    def __init__(self):
+        print("Loading Edge Routing Engine (SiGLIP-So400M)...")
+        # Load locally for edge-simulated routing
+        self.processor = AutoProcessor.from_pretrained("google/siglip-so400m-patch14-384")
+        self.model = AutoModel.from_pretrained("google/siglip-so400m-patch14-384")
 
-def init_engine():
-    global router_engine
-    if router_engine is None:
-        print("Initializing SiGLIP Tensor Routing Core...")
-        router_engine = MANNEngramRouter(
-            ckpt_path="./weights/skew_model_v4full_en.pt",
-            enable_local_intent=False 
-        )
-    return router_engine
+    def route_evidence(self, visual_query, image_paths, margin_ratio=0.15, absolute_floor=0.1):
+        """
+        Relative Margin Thresholding Engine
+        Routes images based on dynamic confidence window relative to the best match.
+        """
+        if not image_paths:
+            return [], {}
+            
+        # Convert file paths to PIL Images
+        images = [Image.open(img).convert("RGB") for img in image_paths]
+        
+        # Extract multimodal tensors
+        inputs = self.processor(text=[visual_query], images=images, padding="max_length", return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            
+        # Get Sigmoid matching probabilities
+        logits_per_image = outputs.logits_per_image
+        probs = torch.sigmoid(logits_per_image).squeeze()
+        
+        # Handle single image fallback
+        if probs.dim() == 0:
+            probs = probs.unsqueeze(0)
+            
+        # Calculate dynamic threshold based on the Anchor (Highest Prob)
+        max_prob = torch.max(probs).item()
+        dynamic_threshold = max(max_prob * (1.0 - margin_ratio), absolute_floor)
+        
+        # Precision Pruning
+        routed_images = []
+        for idx, prob in enumerate(probs):
+            if prob.item() >= dynamic_threshold:
+                routed_images.append(image_paths[idx])
+                
+        metrics = {
+            "Anchor_Probability (Max)": round(max_prob, 4),
+            "Dynamic_Threshold": round(dynamic_threshold, 4),
+            "Total_Candidates": len(image_paths),
+            "Images_Retained": len(routed_images)
+        }
+        
+        return routed_images, metrics
+
+# ==========================================
+# Phase 1: Cloud-Side Distillation (Qwen)
+# ==========================================
+def extract_intent_and_query(clinical_text, hf_token):
+    """
+    Dual-track extraction using Qwen-72B via Hugging Face Inference API.
+    Translates messy text into Clinical Intent and Visual Target.
+    """
+    if not hf_token:
+        return {"error": "Missing Hugging Face API Token."}, "Error"
 
-def extract_intent_via_api(messy_text, hf_token):
-    if not hf_token: return "ERROR: MISSING TOKEN"
     client = InferenceClient("Qwen/Qwen2.5-72B-Instruct", token=hf_token)
-    sys_prompt = (
-        "You are an expert clinical triage AI. Extract ONLY the critical medical symptoms and targeted body parts. "
-        "IGNORE all noise. Output a single concise English sentence."
-    )
+    
+    system_prompt = """You are an expert Clinical Triage AI and a Multimodal Routing Specialist. 
+    Your task is to analyze messy, noisy patient narratives and output a strictly formatted JSON object with two fields.
+    1. "Clinical_Intent": A purified medical summary of the core issue.
+    2. "Visual_Query": An extreme extraction of visual, anatomical, and radiological keywords relevant ONLY to the core issue. Think like an image-recognition model. Use nouns and modalities (e.g., 'Brain MRA, skull, Willis circle'). DO NOT use abstract symptoms like 'headache'.
+    Output ONLY valid JSON."""
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": f"Patient Narrative:\n{clinical_text}"}
+    ]
+
     try:
-        res = client.chat_completion(
-            messages=[{"role": "system", "content": sys_prompt},{"role": "user", "content": messy_text}],
-            max_tokens=60
-        )
-        return res.choices[0].message.content.strip()
+        response = client.chat_completion(messages=messages, max_tokens=300)
+        content = response.choices[0].message.content
+        
+        # Clean markdown formatting if present
+        clean_content = content.replace("```json", "").replace("```", "").strip()
+        result = json.loads(clean_content)
+        
+        return result, result.get("Visual_Query", "medical scan")
     except Exception as e:
-        return f"API ERROR: {str(e)}"
+        return {"error": f"Cloud Distillation Failed: {str(e)}"}, "medical scan"
 
-def process_input(query, files, hf_token, top_p_val):
-    if not hf_token:
-        return "❌ Error", "Missing Token", "Status: Failed.", []
-    if not query:
-        return "⚠️ Warning", "Empty Query", "Status: Waiting.", []
-        
-    engine = init_engine()
-    
-    file_paths = []
-    image_paths = []
-    if files:
-        for f in files:
-            ext = f.name.lower().split('.')[-1]
-            if ext in ['jpg', 'jpeg', 'png', 'bmp', 'webp']:
-                image_paths.append(f.name)
-            else:
-                file_paths.append(f.name)
+# Initialize local router
+router = SiGLIPRouter()
 
+# ==========================================
+# Main Execution Pipeline
+# ==========================================
+def execute_pipeline(hf_token, narrative, images, margin_ratio):
     start_time = time.time()
+    
+    if not images:
+        return {"Error": "No images uploaded."}, [], {"Status": "Failed"}
 
-    # 1. 云端解析意图
-    clean_intent = extract_intent_via_api(query, hf_token)
-    if "API ERROR" in clean_intent:
-        return "Cloud API Error", clean_intent, "Check Token.", []
-
-    # 2. 边缘路由 (使用动态调节的 top_p_val)
-    results = engine.compress(
-        query=clean_intent,
-        context_pool=[], 
-        image_pool=image_paths,
-        top_p=top_p_val  # <--- 核心改动：应用滑块数值
-    )
+    # 1. Cloud Intelligence (Linguistic Distillation)
+    cloud_result, visual_query = extract_intent_and_query(narrative, hf_token)
     
-    latency = time.time() - start_time
-    stats_dict = results.get("stats", {})
+    if "error" in cloud_result:
+        return cloud_result, [], {"Status": "Cloud API Error"}
+
+    # 2. Edge Intelligence (Tensor Routing)
+    routed_imgs, metrics = router.route_evidence(visual_query, images, margin_ratio=margin_ratio)
     
-    stats_report = (
-        f"✅ Routing Complete\n"
-        f"---------------------------\n"
-        f"🎯 Active Top_p: {top_p_val}\n"
-        f"🧠 Cloud Brain: Qwen-72B\n"
-        f"🖼️ Images Retained: {stats_dict.get('retained_images', 0)} / {stats_dict.get('original_images', 0)}\n"
-        f"⚡ Total Latency: {latency:.2f}s"
-    )
+    end_time = time.time()
+    metrics["Total_Latency (s)"] = round(end_time - start_time, 2)
+    metrics["Active_Margin_Ratio"] = margin_ratio
     
-    return clean_intent, "", stats_report, results.get("purified_images_pil", [])
+    return cloud_result, routed_imgs, metrics
 
 # ==========================================
-# UI 布局：加入 Slider 控制器
+# Gradio UI Design
 # ==========================================
-with gr.Blocks(title="MANN-Engram Router", css=".gradio-container {max-width: 100% !important;}") as demo:
-    
-    gr.Markdown("# 🧠 MANN-Engram: Edge-Cloud Multimodal Router")
+with gr.Blocks(title="MANN-Engram Showcase", theme=gr.themes.Base()) as demo:
+    gr.Markdown("# 🧠 MANN-Engram: Edge-Cloud Multimodal Semantic Router")
+    gr.Markdown("> **A Privacy-First, Zero-Hallucination Shield for Clinical Vision-Language Models.**")
     
     with gr.Row():
-        # --- 左侧：配置与参数调节 ---
-        with gr.Column(scale=2, min_width=250):
-            gr.Markdown("### ⚙️ Settings")
-            token_input = gr.Textbox(
-                label="Hugging Face API Token",
-                placeholder="hf_xxxxxxxxxxxxxxxxx",
-                type="password"
-            )
+        # Left Column: Inputs & Settings
+        with gr.Column(scale=1):
+            gr.Markdown("### ⚙️ Engine Settings")
+            hf_token = gr.Textbox(label="Hugging Face API Token (For Cloud Brain)", type="password", placeholder="hf_xxxxxxxx...")
             
-            # 💡 核心新增：Top_p 调节滑块
-            top_p_slider = gr.Slider(
-                minimum=0.1, 
-                maximum=1.0, 
-                value=0.85, 
-                step=0.05, 
-                label="Routing Threshold (Top_p)",
-                info="Lower = Stricter Filtering (Precision); Higher = More Context (Recall)"
+            margin_slider = gr.Slider(
+                minimum=0.05, maximum=0.40, step=0.05, value=0.15, 
+                label="Routing Tolerance (Margin Ratio)",
+                info="Lower (0.05) = Sniper Mode (Extreme Precision). Higher (0.30) = Cluster Mode (Recalls multiple related views)."
             )
             
-            gr.Markdown("---")
-            gr.Markdown(
-                "**Quick Tip:**\n"
-                "To get *ONLY* the brain tumor in Case 1, try setting Top_p to **0.5 or 0.6**."
+            gr.Markdown("### 📥 Patient Data Dump")
+            narrative_input = gr.Textbox(
+                label="Messy Clinical Narrative", lines=8,
+                placeholder="Paste the chaotic patient complaint and history here..."
             )
-
-        # --- 中间：输入 ---
-        with gr.Column(scale=4):
-            gr.Markdown("### 📥 Input Console")
-            query_input = gr.Textbox(label="Clinical Complaint", lines=8)
-            file_input = gr.File(label="Patient Data Dump", file_count="multiple")
-            submit_btn = gr.Button("🚀 Execute Routing", variant="primary", size="lg")
+            image_input = gr.File(label="Upload Unorganized Scans (Images)", file_count="multiple", type="filepath")
             
-            gr.Examples(
-                examples=[["Doctor, I have severe stomach pain, but my real concern is the seizure I had today and the numbness in my left arm. Check my head MRI.", None]],
-                inputs=[query_input, file_input],
-                label="Preset Scenario"
-            )
+            run_btn = gr.Button("🚀 Execute Routing Pipeline", variant="primary")
 
-        # --- 右侧：结果 ---
-        with gr.Column(scale=4):
-            gr.Markdown("### 📤 Output Dashboard")
-            out_intent = gr.Textbox(label="🎯 Extracted Clinical Intent", interactive=False)
-            out_stats = gr.Textbox(label="📊 Metrics", interactive=False, lines=6)
-            out_gallery = gr.Gallery(label="🔭 Routed Evidence", columns=2, height="auto")
+        # Right Column: Outputs
+        with gr.Column(scale=1):
+            gr.Markdown("### ☁️ Cloud Output: Dual-Track Distillation")
+            cloud_output = gr.JSON(label="Purified Intent & Visual Query")
+            
+            gr.Markdown("### 🛡️ Edge Output: Routed Core Evidence")
+            routed_gallery = gr.Gallery(label="Surgically Selected Scans", columns=2, object_fit="contain", height=400)
+            
+            gr.Markdown("### 📊 Telemetry & Metrics")
+            metrics_output = gr.JSON(label="Routing Diagnostics")
 
-    # 点击事件绑定，加入 top_p_slider 输入
-    submit_btn.click(
-        fn=process_input,
-        inputs=[query_input, file_input, token_input, top_p_slider],
-        outputs=[out_intent, gr.State(), out_stats, out_gallery]
+    # Wire up the button
+    run_btn.click(
+        fn=execute_pipeline,
+        inputs=[hf_token, narrative_input, image_input, margin_slider],
+        outputs=[cloud_output, routed_gallery, metrics_output]
     )
 
 if __name__ == "__main__":