Update server.py

server.py

@@ -113,7 +113,7 @@ class SFTVisionT5Model(nn.Module):
             p.requires_grad = False
 
     def generate_reports(self, pixel_values, max_length=100):
-        self.eval()
+        # Removed self.eval() to match Colab behavior
         with torch.no_grad():
             # Extract + project image features
             img_feats = self.img_encoder(pixel_values)          # [B, feature_dim]
@@ -167,7 +167,6 @@ class PPOVisionT5Model(nn.Module):
         self.img_proj = nn.Linear(img_emb_dim, self.txt_model.config.d_model)
 
     def generate_reports(self, images, max_length=128):
-        self.eval()
         with torch.no_grad():
             img_features = self.img_encoder(images)             # [B, feature_dim]
             img_emb = self.img_proj(img_features).unsqueeze(1) # [B, 1, d_model]
@@ -427,10 +426,22 @@ def health():
 @app.post("/sft")
 async def sft_inference(file: UploadFile = File(...)):
     try:
-        tensor = preprocess(await file.read())
-        report = sft_model.generate_reports(tensor)[0]
-        print(f"[SFT] Generated: {report}")
-        return {"report": report[:81]}
+        # Save uploaded file to temp path (matching Colab approach)
+        import tempfile
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+            temp_file.write(await file.read())
+            temp_path = temp_file.name
+
+        try:
+            # Use file path preprocessing (exact Colab match)
+            tensor = preprocess_image(temp_path)
+            report = sft_model.generate_reports(tensor)[0]
+            print(f"[SFT] Generated: {report}")
+            return {"report": report[:81]}
+        finally:
+            # Clean up temp file
+            os.unlink(temp_path)
+
     except Exception as e:
         traceback.print_exc()
         return {"report": f"ERROR: {str(e)}"}
@@ -439,57 +450,67 @@ async def sft_inference(file: UploadFile = File(...)):
 @app.post("/reward")
 async def reward_inference(file: UploadFile = File(...)):
     try:
-        tensor = preprocess(await file.read())
-
-        # First get the SFT report to score
-        sft_report = sft_model.generate_reports(tensor)[0]
-        print(f"[REWARD] Scoring SFT report: {sft_report}")
-
-        if not sft_report.strip():
-            return {"score": 0.0, "feedback": "", "sft_report": ""}
-
-        enc = tokenizer(
-            [sft_report],
-            max_length=128,
-            padding="max_length",
-            truncation=True,
-            return_tensors="pt"
-        )
-        input_ids      = enc.input_ids.to(device)
-        attention_mask = enc.attention_mask.to(device)
+        # Save uploaded file to temp path (matching Colab approach)
+        import tempfile
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+            temp_file.write(await file.read())
+            temp_path = temp_file.name
 
-        with torch.no_grad():
-            raw_score = reward_model(tensor, input_ids, attention_mask).item()
-
-        # Detailed debug logging
-        print(f"[REWARD] Raw neural network output: {raw_score:.6f}")
-        print(f"[REWARD] Clamping to [0,1] range: max(0.0, min(1.0, {raw_score:.6f})) = {max(0.0, min(1.0, raw_score)):.6f}")
-
-        # Quality assessment details
-        rl = sft_report.lower()
-        present = [t for t in KEY_MEDICAL_TERMS if t in rl]
-        missing  = [t for t in KEY_MEDICAL_TERMS if t not in rl]
-        words    = len(sft_report.split())
-        length_q = "good" if 50 <= words <= 150 else ("too short" if words < 50 else "too long")
-
-        print(f"[REWARD] Report analysis:")
-        print(f"         - Total words: {words} ({length_q})")
-        print(f"         - Medical terms present ({len(present)}/{len(KEY_MEDICAL_TERMS)}): {present}")
-        print(f"         - Medical terms missing: {missing}")
-        print(f"         - Key terms list: {KEY_MEDICAL_TERMS}")
-
-        # Reward model architecture details
-        print(f"[REWARD] Model architecture:")
-        print(f"         - CoAtNet feature dim: {reward_model.img_encoder.feature_dim}")
-        print(f"         - T5 d_model: {reward_model.txt_encoder.config.d_model}")
-        print(f"         - Combined feature dim: 1024 (512 img + 512 text)")
-        print(f"         - Reward head: 1024→512→256→1")
-
-        # Clamped score for display
-        score = float(max(0.0, min(1.0, raw_score)))
-        feedback = reward_feedback(sft_report, score)
-        print(f"[REWARD] Final Score={score:.3f}")
-        return {"score": score, "feedback": feedback, "sft_report": sft_report}
+        try:
+            # Use file path preprocessing (exact Colab match)
+            tensor = preprocess_image(temp_path)
+            # First get the SFT report to score
+            sft_report = sft_model.generate_reports(tensor)[0]
+            print(f"[REWARD] Scoring SFT report: {sft_report}")
+
+            if not sft_report.strip():
+                return {"score": 0.0, "feedback": "", "sft_report": ""}
+
+            enc = tokenizer(
+                [sft_report],
+                max_length=128,
+                padding="max_length",
+                truncation=True,
+                return_tensors="pt"
+            )
+            input_ids      = enc.input_ids.to(device)
+            attention_mask = enc.attention_mask.to(device)
+
+            with torch.no_grad():
+                raw_score = reward_model(tensor, input_ids, attention_mask).item()
+
+            # Detailed debug logging
+            print(f"[REWARD] Raw neural network output: {raw_score:.6f}")
+            print(f"[REWARD] Clamping to [0,1] range: max(0.0, min(1.0, {raw_score:.6f})) = {max(0.0, min(1.0, raw_score)):.6f}")
+
+            # Quality assessment details
+            rl = sft_report.lower()
+            present = [t for t in KEY_MEDICAL_TERMS if t in rl]
+            missing  = [t for t in KEY_MEDICAL_TERMS if t not in rl]
+            words    = len(sft_report.split())
+            length_q = "good" if 50 <= words <= 150 else ("too short" if words < 50 else "too long")
+
+            print(f"[REWARD] Report analysis:")
+            print(f"         - Total words: {words} ({length_q})")
+            print(f"         - Medical terms present ({len(present)}/{len(KEY_MEDICAL_TERMS)}): {present}")
+            print(f"         - Medical terms missing: {missing}")
+            print(f"         - Key terms list: {KEY_MEDICAL_TERMS}")
+
+            # Reward model architecture details
+            print(f"[REWARD] Model architecture:")
+            print(f"         - CoAtNet feature dim: {reward_model.img_encoder.feature_dim}")
+            print(f"         - T5 d_model: {reward_model.txt_encoder.config.d_model}")
+            print(f"         - Combined feature dim: 1024 (512 img + 512 text)")
+            print(f"         - Reward head: 1024→512→256→1")
+
+            # Clamped score for display
+            score = float(max(0.0, min(1.0, raw_score)))
+            feedback = reward_feedback(sft_report, score)
+            print(f"[REWARD] Final Score={score:.3f}")
+            return {"score": score, "feedback": feedback, "sft_report": sft_report}
+        finally:
+            # Clean up temp file
+            os.unlink(temp_path)
 
     except Exception as e:
         traceback.print_exc()
@@ -499,10 +520,22 @@ async def reward_inference(file: UploadFile = File(...)):
 @app.post("/ppo")
 async def ppo_inference(file: UploadFile = File(...)):
     try:
-        tensor = preprocess(await file.read())
-        report = ppo_model.generate_reports(tensor)[0]
-        print(f"[PPO] Generated: {report}")
-        return {"report": report}
+        # Save uploaded file to temp path (matching Colab approach)
+        import tempfile
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+            temp_file.write(await file.read())
+            temp_path = temp_file.name
+
+        try:
+            # Use file path preprocessing (exact Colab match)
+            tensor = preprocess_image(temp_path)
+            report = ppo_model.generate_reports(tensor)[0]
+            print(f"[PPO] Generated: {report}")
+            return {"report": report}
+        finally:
+            # Clean up temp file
+            os.unlink(temp_path)
+
     except Exception as e:
         traceback.print_exc()
         return {"report": f"ERROR: {str(e)}"}