Upload agents/finding_interpreter.py with huggingface_hub

agents/finding_interpreter.py

@@ -9,13 +9,12 @@ from PIL import Image
 
 from .base_agent import BaseAgent, AgentResult
 
-# Try to import torch and transformers
+# Import the unified MedGemma engine
 try:
-    import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
-    TORCH_AVAILABLE = True
+    from .medgemma_engine import get_engine, MedGemmaEngine
+    ENGINE_AVAILABLE = True
 except ImportError:
-    TORCH_AVAILABLE = False
+    ENGINE_AVAILABLE = False
 
 
 class FindingInterpreterAgent(BaseAgent):
@@ -23,7 +22,7 @@ class FindingInterpreterAgent(BaseAgent):
     Agent 2: MedGemma Finding Interpreter
     
     Takes CXR analysis results and generates clinical interpretations
-    using Google's MedGemma model.
+    using Google's MedGemma model via the unified engine.
     """
     
     def __init__(self, demo_mode: bool = False):
@@ -32,42 +31,20 @@ class FindingInterpreterAgent(BaseAgent):
             model_name="google/medgemma-4b-it"
         )
         self.demo_mode = demo_mode
-        self.tokenizer = None
+        self.engine = None
     
     def load_model(self) -> bool:
-        """Load MedGemma model."""
-        if self.demo_mode:
-            self.is_loaded = True
-            return True
-        
-        if not TORCH_AVAILABLE:
-            print("Warning: PyTorch not available. Running in demo mode.")
-            self.demo_mode = True
+        """Load MedGemma model via unified engine."""
+        if self.demo_mode or not ENGINE_AVAILABLE:
             self.is_loaded = True
             return True
         
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name,
-                trust_remote_code=True
-            )
-            
-            # Load with appropriate settings for memory efficiency
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                trust_remote_code=True,
-                torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto" if torch.cuda.is_available() else None,
-                low_cpu_mem_usage=True
-            )
-            
-            self.model.eval()
-            self.is_loaded = True
+            self.engine = get_engine(force_demo=self.demo_mode)
+            self.is_loaded = self.engine.is_loaded
             return True
-            
         except Exception as e:
-            print(f"Failed to load MedGemma model: {e}")
-            print("Falling back to demo mode.")
+            print(f"Failed to load MedGemma engine: {e}")
             self.demo_mode = True
             self.is_loaded = True
             return True
@@ -98,11 +75,11 @@ class FindingInterpreterAgent(BaseAgent):
         findings = input_data.get("findings", [])
         region_analysis = input_data.get("region_analysis", {})
         
-        # Process based on mode
-        if self.demo_mode:
-            interpretation = self._simulate_interpretation(findings, region_analysis, context)
-        else:
+        # Process - always try to use real model if available
+        if self.engine and self.engine.is_loaded and self.engine.backend != "demo":
             interpretation = self._run_model_inference(findings, region_analysis, context)
+        else:
+            interpretation = self._simulate_interpretation(findings, region_analysis, context)
         
         processing_time = (time.time() - start_time) * 1000
         
@@ -119,37 +96,68 @@ class FindingInterpreterAgent(BaseAgent):
         region_analysis: Dict,
         context: Optional[Dict]
     ) -> Dict:
-        """Run actual MedGemma inference."""
+        """Run actual MedGemma inference using the unified engine."""
         try:
-            # Prepare prompt
-            prompt = self._build_prompt(findings, region_analysis, context)
+            clinical_context = context.get("clinical_history", "Not provided") if context else "Not provided"
             
-            # Tokenize
-            inputs = self.tokenizer(prompt, return_tensors="pt")
-            if torch.cuda.is_available():
-                inputs = {k: v.cuda() for k, v in inputs.items()}
-            
-            # Generate
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=512,
-                    temperature=0.3,
-                    top_p=0.9,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id
-                )
+            # Generate interpretations for each finding using real MedGemma
+            interpreted_findings = []
+            for finding in findings:
+                prompt = f"""As a radiologist, interpret this chest X-ray finding:
+
+Finding: {finding.get('type', 'Unknown')}
+Region: {finding.get('region', 'Unknown')}
+Severity: {finding.get('severity', 'Unknown')}
+Description: {finding.get('description', 'No description')}
+Clinical History: {clinical_context}
+
+Provide:
+1. Clinical significance (1-2 sentences)
+2. Top 3 differential diagnoses
+3. Recommended follow-up
+
+Be concise and clinically relevant."""
+
+                response = self.engine.generate(prompt, max_tokens=200)
+                
+                interpreted = {
+                    "original": finding,
+                    "clinical_significance": self._extract_significance(response, finding),
+                    "differential_diagnoses": self._get_differentials(finding),
+                    "recommended_followup": self._get_followup(finding),
+                    "medgemma_interpretation": response,
+                    "correlation_notes": f"MedGemma analysis: {response[:100]}..."
+                }
+                interpreted_findings.append(interpreted)
             
-            # Decode
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Generate clinical summary
+            clinical_summary = self._generate_clinical_summary(interpreted_findings, clinical_context)
+            key_concerns = self._identify_key_concerns(interpreted_findings)
             
-            # Parse response
-            return self._parse_model_response(response, findings)
+            return {
+                "interpreted_findings": interpreted_findings,
+                "clinical_summary": clinical_summary,
+                "key_concerns": key_concerns,
+                "abnormal_regions": [
+                    region for region, data in region_analysis.items()
+                    if data.get("status") == "abnormal"
+                ],
+                "confidence_level": "high",
+                "model_used": f"MedGemma ({self.engine.backend})"
+            }
             
         except Exception as e:
             print(f"MedGemma inference error: {e}")
             return self._simulate_interpretation(findings, region_analysis, context)
     
+    def _extract_significance(self, response: str, finding: Dict) -> str:
+        """Extract clinical significance from MedGemma response."""
+        # Take first meaningful sentence
+        sentences = response.split('.')
+        if sentences and len(sentences[0]) > 10:
+            return sentences[0].strip() + "."
+        return self._get_significance(finding)
+    
     def _simulate_interpretation(
         self,
         findings: List[Dict],