Update app.py

app.py

@@ -18,8 +18,6 @@ from huggingface_hub import hf_hub_download
 from transformers import (
     CLIPProcessor,
     CLIPModel,
-    BlipForConditionalGeneration,
-    AutoProcessor,
     AutoTokenizer,
     AutoModelForSeq2SeqLM,
 )
@@ -46,8 +44,6 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print("Using device:", device)
 
-cap_dtype = torch.float16 if device == "cuda" else torch.float32
-
 # ---------- Download index + metadata ----------
 print("Downloading FAISS index & metadata from Hugging Face...")
 
@@ -80,19 +76,8 @@ clip_model = CLIPModel.from_pretrained(CLIP_MODEL_NAME).to(device)
 clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME)
 clip_model.eval()
 
-# ---------- Load BLIP (image -> draft caption) ----------
-print("Loading BLIP ROCO radiology captioning model...")
-CAPTION_MODEL_ID = "WafaaFraih/blip-roco-radiology-captioning"
-
-caption_processor = AutoProcessor.from_pretrained(CAPTION_MODEL_ID)
-caption_model = BlipForConditionalGeneration.from_pretrained(
-    CAPTION_MODEL_ID,
-    torch_dtype=cap_dtype,
-).to(device)
-caption_model.eval()
-
-# ---------- Load FLAN-T5 (text refinement using similar captions) ----------
-print("Loading FLAN-T5 for caption refinement...")
+# ---------- Load FLAN-T5 (caption synthesis) ----------
+print("Loading FLAN-T5 for diagnosis synthesis from similar captions...")
 REFINER_MODEL_ID = "google/flan-t5-base"
 
 refiner_tokenizer = AutoTokenizer.from_pretrained(REFINER_MODEL_ID)
@@ -274,93 +259,78 @@ def search_similar_by_image(image: Image.Image, k: int = 5) -> pd.DataFrame:
     return rows[["ID", "caption", "concepts_manual", "score", "image_url"]]
 
 
-# ---------- Caption cleaning & generation ----------
+# ---------- Helper: caption cleaning & synthesis ----------
 
 def clean_caption(text: str) -> str:
     """
-    Clean captions:
+    Clean generated caption:
     - strip
-    - split into clauses and remove duplicates
-    - normalize spacing and punctuation
+    - remove obvious prompt leftovers
+    - ensure single sentence, nice punctuation
     """
     if not text:
         return ""
 
     text = text.strip()
 
-    parts = re.split(r"[,.]", text)
-    parts = [p.strip() for p in parts if p.strip()]
+    # Drop any leading instruction-like fragments
+    text = re.sub(
+        r"^(you are an expert radiologist[:,]?\s*)",
+        "",
+        text,
+        flags=re.IGNORECASE,
+    )
+    text = re.sub(
+        r"(findings? from similar radiology cases[:,]?\s*)",
+        "",
+        text,
+        flags=re.IGNORECASE,
+    )
 
-    seen = set()
-    unique_parts = []
-    for p in parts:
-        key = p.lower()
-        if key not in seen:
-            seen.add(key)
-            unique_parts.append(p)
+    # Replace multiple separators
+    text = text.replace(" ;", ";")
+    text = re.sub(r"\s+[,;]\s*", ", ", text)
 
-    if unique_parts:
-        cleaned = ", ".join(unique_parts)
-    else:
-        cleaned = text
+    # Collapse spaces
+    text = " ".join(text.split())
 
-    cleaned = re.sub(
-        r"(respectively,?\s+)+", "respectively ", cleaned, flags=re.IGNORECASE
-    )
+    # If there are multiple sentences, keep only the first one
+    parts = re.split(r"(?<=[.!?])\s+", text)
+    if parts:
+        text = parts[0]
 
-    cleaned = " ".join(cleaned.split())
-    if cleaned and not cleaned.endswith("."):
-        cleaned += "."
-    if cleaned:
-        cleaned = cleaned[0].upper() + cleaned[1:]
-    return cleaned
+    # Ensure period
+    if text and not text.endswith((".", "!", "?")):
+        text += "."
 
+    # Capitalize first letter
+    if text:
+        text = text[0].upper() + text[1:]
 
-def generate_draft_caption(image: Image.Image) -> str:
-    """
-    Draft caption directly from image using BLIP.
-    """
-    inputs = caption_processor(images=image, return_tensors="pt").to(
-        device, dtype=cap_dtype
-    )
-    with torch.no_grad():
-        out_ids = caption_model.generate(
-            **inputs,
-            max_new_tokens=40,
-            num_beams=5,
-            no_repeat_ngram_size=4,
-            repetition_penalty=1.4,
-            length_penalty=0.9,
-            early_stopping=True,
-        )
-    raw = caption_processor.batch_decode(out_ids, skip_special_tokens=True)[0]
-    return clean_caption(raw)
+    return text
 
 
-def refine_caption_with_similar_cases(
-    draft_caption: str,
-    similar_captions: str,
-) -> str:
+def synthesize_caption_from_similar_captions(captions: list[str]) -> str:
     """
-    Use FLAN-T5 to rewrite a final diagnosis sentence based on:
-    - draft caption from BLIP (current image)
-    - captions from similar images
+    Use FLAN-T5 to create a diagnosis sentence from captions of similar images.
     """
-    if not draft_caption:
-        draft_caption = "No draft description available."
-    if not similar_captions:
-        # nothing to refine with; just return draft
-        return draft_caption
+    captions = [c.strip() for c in captions if c and isinstance(c, str)]
+    if not captions:
+        return ""
+
+    # Use at most 5-6 captions to keep prompt short
+    caps = captions[:6]
+
+    numbered = "\n".join(
+        f"{i+1}) {c}" for i, c in enumerate(caps)
+    )
 
     prompt = (
-        "You are an expert radiologist.\n\n"
-        "Draft findings from the current image:\n"
-        f"{draft_caption}\n\n"
-        "Findings from similar radiology cases:\n"
-        f"{similar_captions}\n\n"
-        "Based on all of this, write ONE concise radiology impression "
-        "sentence describing the most probable diagnosis and key findings "
-        "for the current image. Do not mention 'similar cases' or 'draft'."
+        "Radiology findings from similar cases:\n"
+        f"{numbered}\n\n"
+        "Based on these, write ONE concise radiology impression sentence "
+        "describing the most likely diagnosis and key findings for the "
+        "current image. Do not mention numbers or 'similar cases'."
     )
 
     inputs = refiner_tokenizer(
@@ -373,13 +343,14 @@ def refine_caption_with_similar_cases(
     with torch.no_grad():
         out_ids = refiner_model.generate(
             **inputs,
-            max_new_tokens=64,
+            max_new_tokens=48,
             num_beams=4,
             length_penalty=0.9,
             no_repeat_ngram_size=4,
         )
-    refined = refiner_tokenizer.decode(out_ids[0], skip_special_tokens=True)
-    return clean_caption(refined)
+
+    raw = refiner_tokenizer.decode(out_ids[0], skip_special_tokens=True)
+    return clean_caption(raw)
 
 
 # ---------- Routes ----------
@@ -388,7 +359,7 @@ def refine_caption_with_similar_cases(
 def root():
     return {
         "status": "ok",
-        "message": "Radiology retrieval + BLIP + FLAN-T5 refinement API",
+        "message": "Radiology retrieval + FLAN-T5 synthesis from similar captions",
     }
 
 
@@ -397,7 +368,7 @@ async def search_by_image(file: UploadFile = File(...), k: int = 5):
     """
     Upload a radiology image.
     Returns:
-      - query_caption: refined caption using draft + similar cases
+      - query_caption: synthesized diagnosis from captions of similar images
       - modality: detected imaging modality
       - scores: random quality metrics
       - results: similar images (similarity + concepts + image_url)
@@ -411,27 +382,18 @@ async def search_by_image(file: UploadFile = File(...), k: int = 5):
     results_df = search_similar_by_image(image, k=k)
     results = results_df.to_dict(orient="records")
 
-    # similar captions context (take up to 5)
+    # 2) Synthesize caption only from similar image captions
     similar_caps_list = results_df["caption"].astype(str).tolist()
-    similar_caps_short = "; ".join(similar_caps_list[:5])
-
-    # 2) Draft caption from BLIP
-    try:
-        draft_caption = generate_draft_caption(image)
-    except Exception as e:
-        print("Error generating draft caption:", e)
-        draft_caption = ""
 
-    # 3) Refine caption with similar case captions
     try:
-        final_caption = refine_caption_with_similar_cases(
-            draft_caption, similar_caps_short
+        final_caption = synthesize_caption_from_similar_captions(
+            similar_caps_list
         )
     except Exception as e:
-        print("Error refining caption:", e)
-        final_caption = draft_caption or None
+        print("Error synthesizing caption:", e)
+        final_caption = ""
 
-    # 4) Modality & scores
+    # 3) Modality & scores
     modality = detect_modality(final_caption or "")
     scores = generate_random_scores()