Update server.py

server.py

@@ -320,57 +320,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
+    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 - EXACTLY as Colab SECTION 8
 print("\n" + "="*80)
-print("LOADING MODELS INTO MEMORY")
+print("LOADING MODELS")
 print("="*80)
 
 sft_model = load_model_from_checkpoint(
@@ -385,17 +362,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 - Exact Colab Match")
 
 app.add_middleware(
     CORSMiddleware,
@@ -405,40 +377,14 @@ 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"
-        }
-    }
-
-
 @app.get("/health")
 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": {k: v for k, v in CONFIG.items() if k != 'device'}
     }
 
 
@@ -446,8 +392,6 @@ def health():
 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
     """
     try:
         # Save uploaded file temporarily
@@ -466,24 +410,18 @@ async def sft_inference(file: UploadFile = File(...)):
         return {
             "report": report,
             "model": "SFT",
-            "source": "best_model.pt",
-            "repository": HF_REPO
+            "method": "generate_report() - exact Colab SECTION 9"
         }
         
     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
     """
     try:
         # Save uploaded file temporarily
@@ -502,16 +440,12 @@ async def ppo_inference(file: UploadFile = File(...)):
         return {
             "report": report,
             "model": "PPO",
-            "source": "rlhf_model.pt",
-            "repository": HF_REPO
+            "method": "generate_report() - exact Colab SECTION 9"
         }
         
     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")
@@ -519,8 +453,6 @@ 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
@@ -541,11 +473,8 @@ 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
+            "method": "generate_report() - exact Colab SECTION 9",
+            "config": {k: v for k, v in CONFIG.items() if k != 'device'}
         }
         
     except Exception as e:
@@ -556,45 +485,43 @@ async def compare_models(file: UploadFile = File(...)):
         }
 
 
-@app.get("/model_info")
-def model_info():
+@app.get("/debug_inference")
+def debug_inference():
     """
-    Get detailed information about loaded models
+    Debug endpoint to verify inference setup matches Colab exactly
     """
     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"
-            }
+        "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'],
         },
-        "architecture": {
-            "vision_encoder": CONFIG['coatnet_model'],
-            "text_model": CONFIG['t5_model'],
-            "image_embedding_dim": CONFIG['img_emb_dim']
+        "tokenizer": CONFIG['t5_model'],
+        "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]
         },
-        "inference_config": {
+        "generation_params": {
             "max_length": CONFIG['max_length'],
             "num_beams": CONFIG['num_beams'],
-            "image_size": CONFIG['image_size']
+            "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
         }
     }
 
@@ -611,19 +538,15 @@ else:
     print("⚠️ Build directory not found, serving API only")
 
 print("\n" + "="*80)
-print("🚀 SERVER READY")
+print("SERVER READY - Using EXACT Colab Inference Code")
 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("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__":