Update server.py

server.py

@@ -12,49 +12,53 @@ from transformers import T5ForConditionalGeneration, T5Tokenizer
 from huggingface_hub import hf_hub_download
 
 # ─────────────────────────────────────────────────────────────────────────────
-# CONFIGURATION - Matching Colab Notebook Exactly
+# CONFIGURATION - EXACTLY matching Colab CONFIG from SECTION 4
 # ─────────────────────────────────────────────────────────────────────────────
+print("="*80)
+print("INITIALIZING CONFIGURATION")
+print("="*80)
+
+# Device setup - EXACTLY as Colab SECTION 3
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"GPU Device: {torch.cuda.get_device_name(0)}")
+    torch.cuda.empty_cache()
+print(f"🖥️  Using device: {device}")
+
+# Configuration - EXACTLY matching Colab SECTION 4
 CONFIG = {
+    # Model architecture settings
     'coatnet_model': 'coatnet_1_rw_224',
     't5_model': 't5-small',
     'img_emb_dim': 768,
     'train_last_stages': 2,
+
+    # Image preprocessing
     'image_size': 224,
+
+    # Inference settings
     'max_length': 100,
     'num_beams': 4,
+
+    # Device
+    'device': device
 }
 
-# ─────────────────────────────────────────────────────────────────────────────
-# DEVICE
-# ─────────────────────────────────────────────────────────────────────────────
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"🖥️  Using device: {device}")
+print("\nConfiguration loaded:")
+for key, value in CONFIG.items():
+    if key != 'device':
+        print(f"  {key}: {value}")
 
 # ─────────────────────────────────────────────────────────────────────────────
-# LOAD TOKENIZER - Matching Colab
+# SECTION 6: Model Architecture Definitions - EXACT COPY from Colab
 # ─────────────────────────────────────────────────────────────────────────────
 print("\n" + "="*80)
-print("LOADING TOKENIZER")
+print("DEFINING MODEL ARCHITECTURES")
 print("="*80)
-tokenizer = T5Tokenizer.from_pretrained(CONFIG['t5_model'])
-print(f"✓ Loaded tokenizer: {CONFIG['t5_model']}")
-
-# ─────────────────────────────────────────────────────────────────────────────
-# IMAGE TRANSFORM - Matching Colab Exactly
-# ─────────────────────────────────────────────────────────────────────────────
-transform = transforms.Compose([
-    transforms.Resize((CONFIG['image_size'], CONFIG['image_size'])),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
-    )
-])
-print(f"✓ Image transform defined (size: {CONFIG['image_size']}x{CONFIG['image_size']})")
 
-# ─────────────────────────────────────────────────────────────────────────────
-# ARCHITECTURE 1: CoAtNetEncoder - Exactly from Colab SECTION 6
-# ─────────────────────────────────────────────────────────────────────────────
+# --- Encoder: CoAtNet --- EXACT COPY from Colab SECTION 6
 class CoAtNetEncoder(nn.Module):
     def __init__(self, model_name="coatnet_1_rw_224", pretrained=True, train_last_stages=2):
         super().__init__()
@@ -80,9 +84,7 @@ class CoAtNetEncoder(nn.Module):
         return self.encoder(x)
 
 
-# ─────────────────────────────────────────────────────────────────────────────
-# ARCHITECTURE 2: VisionT5Model - Exactly from Colab SECTION 6
-# ─────────────────────────────────────────────────────────────────────────────
+# --- Vision-T5 Model --- EXACT COPY from Colab SECTION 6
 class VisionT5Model(nn.Module):
     def __init__(self, img_encoder, txt_model_name="t5-small", img_emb_dim=768):
         super().__init__()
@@ -127,9 +129,6 @@ class VisionT5Model(nn.Module):
         return outputs
 
     def generate_reports(self, pixel_values, max_length=100, num_beams=4):
-        """
-        Generate reports - EXACTLY matching Colab SECTION 6
-        """
         # Extract and project image features
         img_feats = self.img_encoder(pixel_values)
         img_feats = self.proj(img_feats)
@@ -140,7 +139,7 @@ class VisionT5Model(nn.Module):
             inputs_embeds=encoder_hidden_states
         )
 
-        # Generate report using beam search - EXACT parameters from Colab
+        # Generate report using beam search
         generated_ids = self.t5.generate(
             encoder_outputs=encoder_outputs,
             attention_mask=torch.ones(
@@ -157,11 +156,42 @@ class VisionT5Model(nn.Module):
 print("✓ Model architecture classes defined")
 
 # ─────────────────────────────────────────────────────────────────────────────
-# MODEL LOADING FUNCTION - Exactly from Colab SECTION 8
+# SECTION 7: Load Tokenizer and Image Transform - EXACT COPY from Colab
+# ─────────────────────────────────────────────────────────────────────────────
+print("\n" + "="*80)
+print("LOADING TOKENIZER AND IMAGE TRANSFORM")
+print("="*80)
+
+# Load tokenizer
+tokenizer = T5Tokenizer.from_pretrained(CONFIG['t5_model'])
+print(f"✓ Loaded tokenizer: {CONFIG['t5_model']}")
+
+# Define image transform - EXACTLY as Colab SECTION 7
+transform = transforms.Compose([
+    transforms.Resize((CONFIG['image_size'], CONFIG['image_size'])),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+print(f"✓ Image transform defined (size: {CONFIG['image_size']}x{CONFIG['image_size']})")
+
+# ─────────────────────────────────────────────────────────────────────────────
+# SECTION 8: Model Loading Functions - EXACT COPY from Colab
 # ─────────────────────────────────────────────────────────────────────────────
 def load_model_from_checkpoint(checkpoint_path: str, model_name: str, config: dict):
     """
-    Load VisionT5Model from checkpoint - EXACT implementation from Colab
+    Load VisionT5Model from checkpoint.
+    EXACT COPY from Colab SECTION 8
+
+    Args:
+        checkpoint_path: Path to .pt checkpoint file
+        model_name: Name for logging (e.g., 'SFT' or 'PPO')
+        config: Configuration dictionary
+
+    Returns:
+        Loaded model
     """
     print(f"\nLoading {model_name} model...")
     print(f"  Checkpoint: {checkpoint_path}")
@@ -242,22 +272,36 @@ def load_model_from_checkpoint(checkpoint_path: str, model_name: str, config: di
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# INFERENCE FUNCTION - Exactly from Colab SECTION 9
+# SECTION 9: Inference Functions - EXACT COPY from Colab
 # ─────────────────────────────────────────────────────────────────────────────
+def preprocess_image(image_path: str) -> torch.Tensor:
+    """Load and preprocess image. EXACT COPY from Colab SECTION 9"""
+    image = Image.open(image_path).convert('RGB')
+    return transform(image)
+
+
 def generate_report(
     image_path: str,
     model: VisionT5Model,
     config: dict
 ) -> str:
     """
-    Generate medical report from X-ray image - EXACT implementation from Colab
+    Generate medical report from X-ray image.
+    EXACT COPY from Colab SECTION 9
+
+    Args:
+        image_path: Path to X-ray image
+        model: VisionT5Model
+        config: Configuration dictionary
+
+    Returns:
+        Generated report text
     """
     try:
         # Preprocess image
-        image = Image.open(image_path).convert('RGB')
-        pixel_values = transform(image).unsqueeze(0).to(device)
+        pixel_values = preprocess_image(image_path).unsqueeze(0).to(device)
 
-        # Generate report - using EXACT parameters from Colab
+        # Generate report
         with torch.no_grad():
             generated_ids = model.generate_reports(
                 pixel_values,
@@ -301,7 +345,7 @@ except Exception as e:
     PPO_MODEL_PATH = "/content/rlhf_model.pt"
     print(f"⚠️ Using local paths instead")
 
-# Load both models
+# Load both models - EXACTLY as Colab SECTION 8
 print("\n" + "="*80)
 print("LOADING MODELS")
 print("="*80)
@@ -323,7 +367,7 @@ print("\n✓ Both models loaded successfully!")
 # ─────────────────────────────────────────────────────────────────────────────
 # FASTAPI APP
 # ─────────────────────────────────────────────────────────────────────────────
-app = FastAPI(title="Medical Report Generation - Matching Colab")
+app = FastAPI(title="Medical Report Generation - Exact Colab Match")
 
 app.add_middleware(
     CORSMiddleware,
@@ -333,46 +377,41 @@ app.add_middleware(
 )
 
 
-def preprocess_bytes(file_bytes: bytes) -> torch.Tensor:
-    """Preprocess image bytes for inference"""
-    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
-    return transform(img).unsqueeze(0).to(device)
-
-
 @app.get("/health")
 def health():
     return {
         "status": "ok",
         "device": str(device),
+        "cuda_available": torch.cuda.is_available(),
         "models_loaded": True,
-        "config": CONFIG
+        "config": {k: v for k, v in CONFIG.items() if k != 'device'}
     }
 
 
 @app.post("/sft")
 async def sft_inference(file: UploadFile = File(...)):
     """
-    SFT model inference - EXACTLY matching Colab behavior
+    SFT model inference - Uses EXACT generate_report() function from Colab SECTION 9
     """
     try:
-        # Preprocess image
-        tensor = preprocess_bytes(await file.read())
+        # Save uploaded file temporarily
+        temp_path = f"/tmp/{file.filename}"
+        with open(temp_path, "wb") as f:
+            f.write(await file.read())
         
-        # Generate report using EXACT Colab parameters
-        with torch.no_grad():
-            generated_ids = sft_model.generate_reports(
-                tensor,
-                max_length=CONFIG['max_length'],
-                num_beams=CONFIG['num_beams']
-            )
+        # Use EXACT generate_report function from Colab
+        report = generate_report(temp_path, sft_model, CONFIG)
         
-        # Decode - EXACTLY as Colab does
-        report = tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
+        # Clean up temp file
+        os.remove(temp_path)
         
-        print(f"[SFT] Generated: {report}")
+        print(f"[SFT] Generated report: {report}")
         
-        # Return FULL report without truncation
-        return {"report": report, "model": "SFT", "config_used": CONFIG}
+        return {
+            "report": report,
+            "model": "SFT",
+            "method": "generate_report() - exact Colab SECTION 9"
+        }
         
     except Exception as e:
         traceback.print_exc()
@@ -382,27 +421,27 @@ async def sft_inference(file: UploadFile = File(...)):
 @app.post("/ppo")
 async def ppo_inference(file: UploadFile = File(...)):
     """
-    PPO model inference - EXACTLY matching Colab behavior
+    PPO model inference - Uses EXACT generate_report() function from Colab SECTION 9
     """
     try:
-        # Preprocess image
-        tensor = preprocess_bytes(await file.read())
+        # Save uploaded file temporarily
+        temp_path = f"/tmp/{file.filename}"
+        with open(temp_path, "wb") as f:
+            f.write(await file.read())
         
-        # Generate report using EXACT Colab parameters
-        with torch.no_grad():
-            generated_ids = ppo_model.generate_reports(
-                tensor,
-                max_length=CONFIG['max_length'],
-                num_beams=CONFIG['num_beams']
-            )
+        # Use EXACT generate_report function from Colab
+        report = generate_report(temp_path, ppo_model, CONFIG)
         
-        # Decode - EXACTLY as Colab does
-        report = tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
+        # Clean up temp file
+        os.remove(temp_path)
         
-        print(f"[PPO] Generated: {report}")
+        print(f"[PPO] Generated report: {report}")
         
-        # Return FULL report without truncation
-        return {"report": report, "model": "PPO", "config_used": CONFIG}
+        return {
+            "report": report,
+            "model": "PPO",
+            "method": "generate_report() - exact Colab SECTION 9"
+        }
         
     except Exception as e:
         traceback.print_exc()
@@ -413,28 +452,20 @@ async def ppo_inference(file: UploadFile = File(...)):
 async def compare_models(file: UploadFile = File(...)):
     """
     Generate reports from both models for comparison
+    Uses EXACT generate_report() function from Colab
     """
     try:
-        file_bytes = await file.read()
-        tensor = preprocess_bytes(file_bytes)
+        # Save uploaded file temporarily
+        temp_path = f"/tmp/{file.filename}"
+        with open(temp_path, "wb") as f:
+            f.write(await file.read())
         
-        # SFT Generation
-        with torch.no_grad():
-            sft_ids = sft_model.generate_reports(
-                tensor,
-                max_length=CONFIG['max_length'],
-                num_beams=CONFIG['num_beams']
-            )
-        sft_report = tokenizer.decode(sft_ids[0], skip_special_tokens=True).strip()
+        # Use EXACT generate_report function from Colab for both models
+        sft_report = generate_report(temp_path, sft_model, CONFIG)
+        ppo_report = generate_report(temp_path, ppo_model, CONFIG)
         
-        # PPO Generation
-        with torch.no_grad():
-            ppo_ids = ppo_model.generate_reports(
-                tensor,
-                max_length=CONFIG['max_length'],
-                num_beams=CONFIG['num_beams']
-            )
-        ppo_report = tokenizer.decode(ppo_ids[0], skip_special_tokens=True).strip()
+        # Clean up temp file
+        os.remove(temp_path)
         
         print(f"[COMPARE] SFT: {sft_report}")
         print(f"[COMPARE] PPO: {ppo_report}")
@@ -442,7 +473,8 @@ async def compare_models(file: UploadFile = File(...)):
         return {
             "sft_report": sft_report,
             "ppo_report": ppo_report,
-            "config_used": CONFIG
+            "method": "generate_report() - exact Colab SECTION 9",
+            "config": {k: v for k, v in CONFIG.items() if k != 'device'}
         }
         
     except Exception as e:
@@ -453,19 +485,43 @@ async def compare_models(file: UploadFile = File(...)):
         }
 
 
-@app.get("/debug_config")
-def debug_config():
-    """Debug endpoint to check configuration"""
+@app.get("/debug_inference")
+def debug_inference():
+    """
+    Debug endpoint to verify inference setup matches Colab exactly
+    """
     return {
-        "config": CONFIG,
         "device": str(device),
+        "cuda_available": torch.cuda.is_available(),
+        "config": {
+            "coatnet_model": CONFIG['coatnet_model'],
+            "t5_model": CONFIG['t5_model'],
+            "img_emb_dim": CONFIG['img_emb_dim'],
+            "train_last_stages": CONFIG['train_last_stages'],
+            "image_size": CONFIG['image_size'],
+            "max_length": CONFIG['max_length'],
+            "num_beams": CONFIG['num_beams'],
+        },
         "tokenizer": CONFIG['t5_model'],
-        "image_size": CONFIG['image_size'],
-        "max_length": CONFIG['max_length'],
-        "num_beams": CONFIG['num_beams'],
+        "transform": {
+            "resize": f"{CONFIG['image_size']}x{CONFIG['image_size']}",
+            "normalize_mean": [0.485, 0.456, 0.406],
+            "normalize_std": [0.229, 0.224, 0.225]
+        },
+        "generation_params": {
+            "max_length": CONFIG['max_length'],
+            "num_beams": CONFIG['num_beams'],
+            "early_stopping": True,
+            "no_extra_penalties": "✓ Exactly as Colab"
+        },
+        "inference_method": "generate_report() from Colab SECTION 9",
         "models_loaded": {
             "sft": sft_model is not None,
             "ppo": ppo_model is not None
+        },
+        "model_state": {
+            "sft_eval_mode": not sft_model.training if sft_model else None,
+            "ppo_eval_mode": not ppo_model.training if ppo_model else None
         }
     }
 
@@ -481,6 +537,17 @@ if os.path.exists("build"):
 else:
     print("⚠️ Build directory not found, serving API only")
 
+print("\n" + "="*80)
+print("SERVER READY - Using EXACT Colab Inference Code")
+print("="*80)
+print("Key points:")
+print("  ✓ Model architecture: VisionT5Model (exact copy from Colab SECTION 6)")
+print("  ✓ Inference method: generate_report() (exact copy from Colab SECTION 9)")
+print("  ✓ Generation params: max_length=100, num_beams=4, early_stopping=True")
+print("  ✓ No extra penalties: NO repetition_penalty, NO no_repeat_ngram_size")
+print("  ✓ Transform: Resize 224x224, Normalize [0.485,0.456,0.406]/[0.229,0.224,0.225]")
+print("  ✓ Device handling: Same as Colab")
+print("="*80)
 
 if __name__ == "__main__":
     import uvicorn