Update server.py

server.py

@@ -320,34 +320,57 @@ def generate_report(
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# LOAD MODELS FROM HUGGINGFACE
+# LOAD MODELS FROM HUGGINGFACE - Shree2604/BioStack
 # ─────────────────────────────────────────────────────────────────────────────
 print("\n" + "="*80)
-print("LOADING MODELS FROM HUGGINGFACE")
+print("DOWNLOADING 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="vinaykumarhs2020/RLHF_radiology_model",
+        repo_id=HF_REPO,
         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="vinaykumarhs2020/RLHF_radiology_model",
+        repo_id=HF_REPO,
         filename="rlhf_model.pt"
     )
-    print(f"✓ Downloaded SFT model: {SFT_MODEL_PATH}")
-    print(f"✓ Downloaded PPO model: {PPO_MODEL_PATH}")
+    print(f"   ✓ PPO model downloaded: {PPO_MODEL_PATH}")
+    
+    print("\n✅ All models downloaded successfully!")
+    
 except Exception as e:
-    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")
+    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("\n" + "="*80)
-print("LOADING MODELS")
+print("LOADING MODELS INTO MEMORY")
 print("="*80)
 
 sft_model = load_model_from_checkpoint(
@@ -362,12 +385,17 @@ ppo_model = load_model_from_checkpoint(
     CONFIG
 )
 
-print("\n✓ Both models loaded successfully!")
+print("\n✅ Both models loaded successfully!")
+print("="*80)
 
 # ─────────────────────────────────────────────────────────────────────────────
 # FASTAPI APP
 # ─────────────────────────────────────────────────────────────────────────────
-app = FastAPI(title="Medical Report Generation - Exact Colab Match")
+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.add_middleware(
     CORSMiddleware,
@@ -377,14 +405,40 @@ 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": True,
-        "config": {k: v for k, v in CONFIG.items() if k != 'device'}
+        "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)
+        }
     }
 
 
@@ -392,6 +446,8 @@ 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
@@ -410,18 +466,24 @@ async def sft_inference(file: UploadFile = File(...)):
         return {
             "report": report,
             "model": "SFT",
-            "method": "generate_report() - exact Colab SECTION 9"
+            "source": "best_model.pt",
+            "repository": HF_REPO
         }
         
     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
@@ -440,12 +502,16 @@ async def ppo_inference(file: UploadFile = File(...)):
         return {
             "report": report,
             "model": "PPO",
-            "method": "generate_report() - exact Colab SECTION 9"
+            "source": "rlhf_model.pt",
+            "repository": HF_REPO
         }
         
     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")
@@ -453,6 +519,8 @@ 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
@@ -473,8 +541,11 @@ async def compare_models(file: UploadFile = File(...)):
         return {
             "sft_report": sft_report,
             "ppo_report": ppo_report,
-            "method": "generate_report() - exact Colab SECTION 9",
-            "config": {k: v for k, v in CONFIG.items() if k != 'device'}
+            "models": {
+                "sft": "best_model.pt",
+                "ppo": "rlhf_model.pt"
+            },
+            "repository": HF_REPO
         }
         
     except Exception as e:
@@ -485,43 +556,45 @@ async def compare_models(file: UploadFile = File(...)):
         }
 
 
-@app.get("/debug_inference")
-def debug_inference():
+@app.get("/model_info")
+def model_info():
     """
-    Debug endpoint to verify inference setup matches Colab exactly
+    Get detailed information about loaded models
     """
     return {
-        "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'],
+        "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"
+            }
         },
-        "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]
+        "architecture": {
+            "vision_encoder": CONFIG['coatnet_model'],
+            "text_model": CONFIG['t5_model'],
+            "image_embedding_dim": CONFIG['img_emb_dim']
         },
-        "generation_params": {
+        "inference_config": {
             "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
+            "image_size": CONFIG['image_size']
         }
     }
 
@@ -538,15 +611,19 @@ else:
     print("⚠️ Build directory not found, serving API only")
 
 print("\n" + "="*80)
-print("SERVER READY - Using EXACT Colab Inference Code")
+print("🚀 SERVER READY")
 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(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__":