Update server.py

server.py

@@ -12,53 +12,49 @@ from transformers import T5ForConditionalGeneration, T5Tokenizer
 from huggingface_hub import hf_hub_download
 
 # ─────────────────────────────────────────────────────────────────────────────
-# CONFIGURATION - EXACTLY matching Colab CONFIG from SECTION 4
+# CONFIGURATION - Matching Colab Notebook Exactly
 # ─────────────────────────────────────────────────────────────────────────────
-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
 }
 
-print("\nConfiguration loaded:")
-for key, value in CONFIG.items():
-    if key != 'device':
-        print(f"  {key}: {value}")
+# ─────────────────────────────────────────────────────────────────────────────
+# DEVICE
+# ─────────────────────────────────────────────────────────────────────────────
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"🖥️  Using device: {device}")
 
 # ─────────────────────────────────────────────────────────────────────────────
-# SECTION 6: Model Architecture Definitions - EXACT COPY from Colab
+# LOAD TOKENIZER - Matching Colab
 # ─────────────────────────────────────────────────────────────────────────────
 print("\n" + "="*80)
-print("DEFINING MODEL ARCHITECTURES")
+print("LOADING TOKENIZER")
 print("="*80)
+tokenizer = T5Tokenizer.from_pretrained(CONFIG['t5_model'])
+print(f"✓ Loaded tokenizer: {CONFIG['t5_model']}")
 
-# --- Encoder: CoAtNet --- EXACT COPY from Colab SECTION 6
+# ─────────────────────────────────────────────────────────────────────────────
+# 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
+# ─────────────────────────────────────────────────────────────────────────────
 class CoAtNetEncoder(nn.Module):
     def __init__(self, model_name="coatnet_1_rw_224", pretrained=True, train_last_stages=2):
         super().__init__()
@@ -84,7 +80,9 @@ class CoAtNetEncoder(nn.Module):
         return self.encoder(x)
 
 
-# --- Vision-T5 Model --- EXACT COPY from Colab SECTION 6
+# ─────────────────────────────────────────────────────────────────────────────
+# ARCHITECTURE 2: VisionT5Model - Exactly from Colab SECTION 6
+# ─────────────────────────────────────────────────────────────────────────────
 class VisionT5Model(nn.Module):
     def __init__(self, img_encoder, txt_model_name="t5-small", img_emb_dim=768):
         super().__init__()
@@ -129,6 +127,9 @@ 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)
@@ -139,7 +140,7 @@ class VisionT5Model(nn.Module):
             inputs_embeds=encoder_hidden_states
         )
 
-        # Generate report using beam search
+        # Generate report using beam search - EXACT parameters from Colab
         generated_ids = self.t5.generate(
             encoder_outputs=encoder_outputs,
             attention_mask=torch.ones(
@@ -156,42 +157,11 @@ class VisionT5Model(nn.Module):
 print("✓ Model architecture classes defined")
 
 # ─────────────────────────────────────────────────────────────────────────────
-# 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
+# MODEL LOADING FUNCTION - Exactly from Colab SECTION 8
 # ─────────────────────────────────────────────────────────────────────────────
 def load_model_from_checkpoint(checkpoint_path: str, model_name: str, config: dict):
     """
-    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
+    Load VisionT5Model from checkpoint - EXACT implementation from Colab
     """
     print(f"\nLoading {model_name} model...")
     print(f"  Checkpoint: {checkpoint_path}")
@@ -272,36 +242,22 @@ def load_model_from_checkpoint(checkpoint_path: str, model_name: str, config: di
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# SECTION 9: Inference Functions - EXACT COPY from Colab
+# INFERENCE FUNCTION - Exactly from Colab SECTION 9
 # ─────────────────────────────────────────────────────────────────────────────
-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 COPY from Colab SECTION 9
-
-    Args:
-        image_path: Path to X-ray image
-        model: VisionT5Model
-        config: Configuration dictionary
-
-    Returns:
-        Generated report text
+    Generate medical report from X-ray image - EXACT implementation from Colab
     """
     try:
         # Preprocess image
-        pixel_values = preprocess_image(image_path).unsqueeze(0).to(device)
+        image = Image.open(image_path).convert('RGB')
+        pixel_values = transform(image).unsqueeze(0).to(device)
 
-        # Generate report
+        # Generate report - using EXACT parameters from Colab
         with torch.no_grad():
             generated_ids = model.generate_reports(
                 pixel_values,
@@ -320,57 +276,34 @@ def generate_report(
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# LOAD MODELS FROM HUGGINGFACE - Shree2604/BioStack
+# LOAD MODELS FROM HUGGINGFACE
 # ─────────────────────────────────────────────────────────────────────────────
 print("\n" + "="*80)
-print("DOWNLOADING MODELS FROM HUGGINGFACE")
+print("LOADING MODELS FROM HUGGINGFACE")
 print("="*80)
 
-# Hugging Face repository
-HF_REPO = "Shree2604/BioStack"
-
 # Download model files from Hugging Face
 try:
-    print(f"📦 Downloading from repository: {HF_REPO}")
-    print("This may take a few minutes on first run...\n")
-    
-    # Download SFT model
-    print("1️⃣ Downloading SFT model (best_model.pt)...")
     SFT_MODEL_PATH = hf_hub_download(
-        repo_id=HF_REPO,
+        repo_id="vinaykumarhs2020/RLHF_radiology_model",
         filename="best_model.pt"
     )
-    print(f"   ✓ SFT model downloaded: {SFT_MODEL_PATH}")
-    
-    # Download Reward model
-    print("\n2️⃣ Downloading Reward model (reward_model.pt)...")
-    REWARD_MODEL_PATH = hf_hub_download(
-        repo_id=HF_REPO,
-        filename="reward_model.pt"
-    )
-    print(f"   ✓ Reward model downloaded: {REWARD_MODEL_PATH}")
-    
-    # Download PPO model
-    print("\n3️⃣ Downloading PPO model (rlhf_model.pt)...")
     PPO_MODEL_PATH = hf_hub_download(
-        repo_id=HF_REPO,
+        repo_id="vinaykumarhs2020/RLHF_radiology_model",
         filename="rlhf_model.pt"
     )
-    print(f"   ✓ PPO model downloaded: {PPO_MODEL_PATH}")
-    
-    print("\n✅ All models downloaded successfully!")
-    
+    print(f"✓ Downloaded SFT model: {SFT_MODEL_PATH}")
+    print(f"✓ Downloaded PPO model: {PPO_MODEL_PATH}")
 except Exception as e:
-    print(f"\n❌ Error downloading models from Hugging Face: {e}")
-    print("Please check:")
-    print(f"  - Repository exists: https://huggingface.co/{HF_REPO}")
-    print("  - Files exist: best_model.pt, reward_model.pt, rlhf_model.pt")
-    print("  - You have internet connection")
-    raise
-
-# Load both models - EXACTLY as Colab SECTION 8
+    print(f"❌ Error downloading models: {e}")
+    # Fallback to local paths if downloads fail
+    SFT_MODEL_PATH = "/content/best_model.pt"
+    PPO_MODEL_PATH = "/content/rlhf_model.pt"
+    print(f"⚠️ Using local paths instead")
+
+# Load both models
 print("\n" + "="*80)
-print("LOADING MODELS INTO MEMORY")
+print("LOADING MODELS")
 print("="*80)
 
 sft_model = load_model_from_checkpoint(
@@ -385,17 +318,12 @@ ppo_model = load_model_from_checkpoint(
     CONFIG
 )
 
-print("\n✅ Both models loaded successfully!")
-print("="*80)
+print("\n✓ Both models loaded successfully!")
 
 # ─────────────────────────────────────────────────────────────────────────────
 # FASTAPI APP
 # ─────────────────────────────────────────────────────────────────────────────
-app = FastAPI(
-    title="BioStack Medical Report Generation",
-    description="Medical X-ray report generation using SFT and PPO models from Shree2604/BioStack",
-    version="1.0.0"
-)
+app = FastAPI(title="Medical Report Generation - Matching Colab")
 
 app.add_middleware(
     CORSMiddleware,
@@ -405,23 +333,10 @@ app.add_middleware(
 )
 
 
-@app.get("/")
-def root():
-    return {
-        "message": "BioStack Medical Report Generation API",
-        "repository": "Shree2604/BioStack",
-        "models": {
-            "sft": "best_model.pt",
-            "ppo": "rlhf_model.pt",
-            "reward": "reward_model.pt"
-        },
-        "endpoints": {
-            "health": "GET /health - Check API status",
-            "sft": "POST /sft - Generate report using SFT model",
-            "ppo": "POST /ppo - Generate report using PPO model",
-            "compare": "POST /compare - Compare both models"
-        }
-    }
+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")
@@ -429,111 +344,97 @@ def health():
     return {
         "status": "ok",
         "device": str(device),
-        "cuda_available": torch.cuda.is_available(),
-        "models_loaded": {
-            "sft": sft_model is not None,
-            "ppo": ppo_model is not None
-        },
-        "repository": HF_REPO,
-        "model_files": {
-            "sft": os.path.basename(SFT_MODEL_PATH),
-            "ppo": os.path.basename(PPO_MODEL_PATH)
-        }
+        "models_loaded": True,
+        "config": CONFIG
     }
 
 
 @app.post("/sft")
 async def sft_inference(file: UploadFile = File(...)):
     """
-    SFT model inference - Uses EXACT generate_report() function from Colab SECTION 9
-    
-    Model: best_model.pt from Shree2604/BioStack
+    SFT model inference - EXACTLY matching Colab behavior
     """
     try:
-        # Save uploaded file temporarily
-        temp_path = f"/tmp/{file.filename}"
-        with open(temp_path, "wb") as f:
-            f.write(await file.read())
+        # Preprocess image
+        tensor = preprocess_bytes(await file.read())
         
-        # Use EXACT generate_report function from Colab
-        report = generate_report(temp_path, sft_model, CONFIG)
+        # 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']
+            )
         
-        # Clean up temp file
-        os.remove(temp_path)
+        # Decode - EXACTLY as Colab does
+        report = tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
         
-        print(f"[SFT] Generated report: {report}")
+        print(f"[SFT] Generated: {report}")
         
-        return {
-            "report": report,
-            "model": "SFT",
-            "source": "best_model.pt",
-            "repository": HF_REPO
-        }
+        # Return FULL report without truncation
+        return {"report": report, "model": "SFT", "config_used": CONFIG}
         
     except Exception as e:
         traceback.print_exc()
-        return {
-            "report": f"ERROR: {str(e)}",
-            "model": "SFT"
-        }
+        return {"report": f"ERROR: {str(e)}", "model": "SFT"}
 
 
 @app.post("/ppo")
 async def ppo_inference(file: UploadFile = File(...)):
     """
-    PPO model inference - Uses EXACT generate_report() function from Colab SECTION 9
-    
-    Model: rlhf_model.pt from Shree2604/BioStack
+    PPO model inference - EXACTLY matching Colab behavior
     """
     try:
-        # Save uploaded file temporarily
-        temp_path = f"/tmp/{file.filename}"
-        with open(temp_path, "wb") as f:
-            f.write(await file.read())
+        # Preprocess image
+        tensor = preprocess_bytes(await file.read())
         
-        # Use EXACT generate_report function from Colab
-        report = generate_report(temp_path, ppo_model, CONFIG)
+        # 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']
+            )
         
-        # Clean up temp file
-        os.remove(temp_path)
+        # Decode - EXACTLY as Colab does
+        report = tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
         
-        print(f"[PPO] Generated report: {report}")
+        print(f"[PPO] Generated: {report}")
         
-        return {
-            "report": report,
-            "model": "PPO",
-            "source": "rlhf_model.pt",
-            "repository": HF_REPO
-        }
+        # Return FULL report without truncation
+        return {"report": report, "model": "PPO", "config_used": CONFIG}
         
     except Exception as e:
         traceback.print_exc()
-        return {
-            "report": f"ERROR: {str(e)}",
-            "model": "PPO"
-        }
+        return {"report": f"ERROR: {str(e)}", "model": "PPO"}
 
 
 @app.post("/compare")
 async def compare_models(file: UploadFile = File(...)):
     """
     Generate reports from both models for comparison
-    Uses EXACT generate_report() function from Colab
-    
-    Models: best_model.pt and rlhf_model.pt from Shree2604/BioStack
     """
     try:
-        # Save uploaded file temporarily
-        temp_path = f"/tmp/{file.filename}"
-        with open(temp_path, "wb") as f:
-            f.write(await file.read())
+        file_bytes = await file.read()
+        tensor = preprocess_bytes(file_bytes)
         
-        # 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)
+        # 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()
         
-        # Clean up temp file
-        os.remove(temp_path)
+        # 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()
         
         print(f"[COMPARE] SFT: {sft_report}")
         print(f"[COMPARE] PPO: {ppo_report}")
@@ -541,11 +442,7 @@ async def compare_models(file: UploadFile = File(...)):
         return {
             "sft_report": sft_report,
             "ppo_report": ppo_report,
-            "models": {
-                "sft": "best_model.pt",
-                "ppo": "rlhf_model.pt"
-            },
-            "repository": HF_REPO
+            "config_used": CONFIG
         }
         
     except Exception as e:
@@ -556,45 +453,19 @@ async def compare_models(file: UploadFile = File(...)):
         }
 
 
-@app.get("/model_info")
-def model_info():
-    """
-    Get detailed information about loaded models
-    """
+@app.get("/debug_config")
+def debug_config():
+    """Debug endpoint to check configuration"""
     return {
-        "repository": HF_REPO,
-        "repository_url": f"https://huggingface.co/{HF_REPO}",
-        "models": {
-            "sft": {
-                "filename": "best_model.pt",
-                "url": f"https://huggingface.co/{HF_REPO}/blob/main/best_model.pt",
-                "local_path": SFT_MODEL_PATH,
-                "loaded": sft_model is not None,
-                "in_eval_mode": not sft_model.training if sft_model else None
-            },
-            "ppo": {
-                "filename": "rlhf_model.pt",
-                "url": f"https://huggingface.co/{HF_REPO}/blob/main/rlhf_model.pt",
-                "local_path": PPO_MODEL_PATH,
-                "loaded": ppo_model is not None,
-                "in_eval_mode": not ppo_model.training if ppo_model else None
-            },
-            "reward": {
-                "filename": "reward_model.pt",
-                "url": f"https://huggingface.co/{HF_REPO}/blob/main/reward_model.pt",
-                "local_path": REWARD_MODEL_PATH,
-                "note": "Downloaded but not loaded in this API"
-            }
-        },
-        "architecture": {
-            "vision_encoder": CONFIG['coatnet_model'],
-            "text_model": CONFIG['t5_model'],
-            "image_embedding_dim": CONFIG['img_emb_dim']
-        },
-        "inference_config": {
-            "max_length": CONFIG['max_length'],
-            "num_beams": CONFIG['num_beams'],
-            "image_size": CONFIG['image_size']
+        "config": CONFIG,
+        "device": str(device),
+        "tokenizer": CONFIG['t5_model'],
+        "image_size": CONFIG['image_size'],
+        "max_length": CONFIG['max_length'],
+        "num_beams": CONFIG['num_beams'],
+        "models_loaded": {
+            "sft": sft_model is not None,
+            "ppo": ppo_model is not None
         }
     }
 
@@ -610,21 +481,6 @@ if os.path.exists("build"):
 else:
     print("⚠️ Build directory not found, serving API only")
 
-print("\n" + "="*80)
-print("🚀 SERVER READY")
-print("="*80)
-print(f"Repository: {HF_REPO}")
-print("Models loaded:")
-print(f"  ✓ SFT: best_model.pt")
-print(f"  ✓ PPO: rlhf_model.pt")
-print("\nEndpoints:")
-print("  GET  /              - API info")
-print("  GET  /health        - Health check")
-print("  GET  /model_info    - Model details")
-print("  POST /sft           - SFT inference")
-print("  POST /ppo           - PPO inference")
-print("  POST /compare       - Compare both models")
-print("="*80)
 
 if __name__ == "__main__":
     import uvicorn