Upload agents/report_generator.py with huggingface_hub

agents/report_generator.py

@@ -9,20 +9,20 @@ from datetime import datetime
 
 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
-    TORCH_AVAILABLE = True
+    from .medgemma_engine import get_engine, MedGemmaEngine
+    ENGINE_AVAILABLE = True
 except ImportError:
-    TORCH_AVAILABLE = False
+    ENGINE_AVAILABLE = False
 
 
 class ReportGeneratorAgent(BaseAgent):
     """
     Agent 3: MedGemma Report Generator
     
-    Generates structured radiology reports from interpreted findings.
+    Generates structured radiology reports from interpreted findings
+    using the unified MedGemma engine.
     """
     
     def __init__(self, demo_mode: bool = False):
@@ -31,39 +31,20 @@ class ReportGeneratorAgent(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 for report generation."""
-        if self.demo_mode:
-            self.is_loaded = True
-            return True
-        
-        if not TORCH_AVAILABLE:
-            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
-            )
-            
-            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 model for report generation: {e}")
+            print(f"Failed to load MedGemma engine: {e}")
             self.demo_mode = True
             self.is_loaded = True
             return True
@@ -95,13 +76,13 @@ class ReportGeneratorAgent(BaseAgent):
         clinical_summary = input_data.get("clinical_summary", "")
         key_concerns = input_data.get("key_concerns", [])
         
-        # Process based on mode
-        if self.demo_mode:
-            report = self._simulate_report_generation(
+        # Process - always try to use real model if available
+        if self.engine and self.engine.is_loaded and self.engine.backend != "demo":
+            report = self._run_model_inference(
                 interpreted_findings, clinical_summary, key_concerns, context
             )
         else:
-            report = self._run_model_inference(
+            report = self._simulate_report_generation(
                 interpreted_findings, clinical_summary, key_concerns, context
             )
         
@@ -121,27 +102,14 @@ class ReportGeneratorAgent(BaseAgent):
         key_concerns: List[str],
         context: Optional[Dict]
     ) -> Dict:
-        """Generate report using MedGemma."""
+        """Generate report using MedGemma via unified engine."""
         try:
             prompt = self._build_report_prompt(
                 interpreted_findings, clinical_summary, key_concerns, context
             )
             
-            inputs = self.tokenizer(prompt, return_tensors="pt")
-            if torch.cuda.is_available():
-                inputs = {k: v.cuda() for k, v in inputs.items()}
-            
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=1024,
-                    temperature=0.3,
-                    top_p=0.9,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id
-                )
-            
-            report_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Use the unified engine to generate report
+            report_text = self.engine.generate(prompt, max_tokens=500)
             
             return self._structure_report(report_text, interpreted_findings, context)