Deploy BioStack RLHF Medical Demo

server.py

@@ -141,16 +141,7 @@ class SFTVisionT5Model(nn.Module):
                 repetition_penalty=1.3,
             )
 
-        reports = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-        # Strip any leading "Projection: X." prefix that leaked from training data
-        cleaned = []
-        for r in reports:
-            if r.lower().startswith("projection:"):
-                # Remove the first "Projection: X." segment
-                parts = r.split(".", 1)
-                r = parts[1].strip() if len(parts) > 1 else r
-            cleaned.append(r)
-        return cleaned
+        return generated_ids
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -191,16 +182,7 @@ class PPOVisionT5Model(nn.Module):
                 repetition_penalty=1.3,
             )
 
-        reports = tokenizer.batch_decode(generated, skip_special_tokens=True)
-        # Strip any leading "Projection: X." prefix that leaked from training data
-        cleaned = []
-        for r in reports:
-            if r.lower().startswith("projection:"):
-                # Remove the first "Projection: X." segment
-                parts = r.split(".", 1)
-                r = parts[1].strip() if len(parts) > 1 else r
-            cleaned.append(r)
-        return cleaned
+        return generated
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -428,7 +410,12 @@ def health():
 async def sft_inference(file: UploadFile = File(...)):
     try:
         tensor = preprocess(await file.read())
-        report = sft_model.generate_reports(tensor)[0]
+        generated_ids = sft_model.generate_reports(tensor)
+        report = tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
+        # Strip any leading "Projection: X." prefix that leaked from training data
+        if report.lower().startswith("projection:"):
+            parts = report.split(".", 1)
+            report = parts[1].strip() if len(parts) > 1 else report
         print(f"[SFT] Generated: {report}")
         return {"report": report[:81]}
     except Exception as e:
@@ -442,7 +429,12 @@ async def reward_inference(file: UploadFile = File(...)):
         tensor = preprocess(await file.read())
 
         # First get the SFT report to score
-        sft_report = sft_model.generate_reports(tensor)[0]
+        sft_generated_ids = sft_model.generate_reports(tensor)
+        sft_report = tokenizer.decode(sft_generated_ids[0], skip_special_tokens=True).strip()
+        # Strip any leading "Projection: X." prefix that leaked from training data
+        if sft_report.lower().startswith("projection:"):
+            parts = sft_report.split(".", 1)
+            sft_report = parts[1].strip() if len(parts) > 1 else sft_report
         print(f"[REWARD] Scoring SFT report: {sft_report}")
 
         if not sft_report.strip():
@@ -500,7 +492,12 @@ async def reward_inference(file: UploadFile = File(...)):
 async def ppo_inference(file: UploadFile = File(...)):
     try:
         tensor = preprocess(await file.read())
-        report = ppo_model.generate_reports(tensor)[0]
+        generated_ids = ppo_model.generate_reports(tensor)
+        report = tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
+        # Strip any leading "Projection: X." prefix that leaked from training data
+        if report.lower().startswith("projection:"):
+            parts = report.split(".", 1)
+            report = parts[1].strip() if len(parts) > 1 else report
         print(f"[PPO] Generated: {report}")
         return {"report": report}
     except Exception as e: