submission_task1.py

#!/usr/bin/env python
"""
ImageCLEFmed-MEDVQA-GI-2026 Task 1 Submission (aggressive constraint pipeline)
==============================================================================
CSMorgan-MEDVQA / Morgan State University
Peter Ojonugwa Ejiga <ojeji1@morgan.edu>

Strategy
--------
- Base: Qwen/Qwen2.5-VL-7B-Instruct in 4-bit nf4
- Adapter: SimulaMet/Qwen2.5-VL-KvasirVQA-x1-ft (organizer baseline) OR the
  USER_ADAPTER_REPO if it has been pushed.
- Decoding: greedy (T=0, top_k=1, max_tokens=20)
- Per-question constraint: for each test question, look up the top-K training
  answers for that exact question text, fuzzy-match model output against those,
  snap to the best candidate above threshold. Falls back to model output if
  no good match.
- No synonym collapse. Multi-word references like "sigmoid colon", "pink;red",
  "5-10mm" are preserved through the pipeline.
"""

import json
import os
import re
import sys
import tempfile
import time
import subprocess
import platform
from collections import Counter
from difflib import SequenceMatcher

import torch
from tqdm import tqdm
from datasets import load_dataset
from evaluate import load
from transformers import BitsAndBytesConfig

# Memory hints for shared HF-Space GPU
os.environ.setdefault("MAX_PIXELS", "640000")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

bleu   = load("bleu")
rouge  = load("rouge")
meteor = load("meteor")

val_dataset = load_dataset("SimulaMet/Kvasir-VQA-test", split="validation")
predictions = []

gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"
device   = "cuda" if torch.cuda.is_available() else "cpu"

def get_mem():
    return torch.cuda.memory_allocated(device) / (1024 ** 2) \
        if torch.cuda.is_available() else 0

initial_mem = get_mem()

SUBMISSION_INFO = {
    "Participant_Names": "Peter Ojonugwa Ejiga",
    "Affiliations": "Morgan State University, Computer Vision & AI Lab",
    "Contact_emails": ["ojeji1@morgan.edu"],
    "Team_Name": "CSMorgan-MEDVQA",
    "Country": "USA",
    "Notes_to_organizers": (
        "Per-question constrained decoding on top of SimulaMet's "
        "Kvasir-VQA-x1 adapter. Model generates greedily (T=0, max 20 tokens). "
        "Output is fuzzy-matched against top-K training answers for the exact "
        "question text using SequenceMatcher; best match above threshold 0.5 "
        "is returned; otherwise raw model output is used."
    ),
}

HF_REPO_ID = "sageofai/Qwen25VL-MEDVQA-GI-S1-subtask1-v4"
FALLBACK_ADAPTER = "SimulaMet/Qwen2.5-VL-KvasirVQA-x1-ft"

# ===========================================================================
#                         COMPANION FILE FETCHER
# medvqa container only pulls submission_task1.py + requirements.txt.
# Pull the answer bank from the same repo at runtime.
# ===========================================================================
_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path: sys.path.insert(0, _HERE)

QUESTION_BANK = None
_bank_path = os.path.join(_HERE, "qbank.json")
if not os.path.exists(_bank_path):
    try:
        from huggingface_hub import hf_hub_download
        _bank_path = hf_hub_download(
            repo_id=HF_REPO_ID, filename="qbank.json", repo_type="model",
        )
        print(f"[fetch] qbank.json from HF: {_bank_path}")
    except Exception as e:
        print(f"[fetch] qbank.json unavailable: {e}")
        _bank_path = None

if _bank_path and os.path.exists(_bank_path):
    with open(_bank_path) as f:
        QUESTION_BANK = json.load(f)
    print(f"Loaded question bank: {len(QUESTION_BANK)} unique training questions")
else:
    print("WARNING: no question bank. Falling back to raw model output.")
    QUESTION_BANK = {}

# Build flat list of ALL training answers for fallback when test question
# isn't in the question bank.
ALL_TRAINING_ANSWERS = []
for q, answers in QUESTION_BANK.items():
    for ans_obj in answers:
        ALL_TRAINING_ANSWERS.append(ans_obj["answer"])
ALL_TRAINING_ANSWERS = list(set(ALL_TRAINING_ANSWERS))
print(f"  flat training answers: {len(ALL_TRAINING_ANSWERS)}")

# ===========================================================================
#                          ADAPTER SELECTION
# Prefer user-pushed adapter; fall back to SimulaMet's baseline if absent.
# ===========================================================================
def _adapter_exists(repo_id):
    if not repo_id: return False
    try:
        from huggingface_hub import HfApi
        info = HfApi().repo_info(repo_id=repo_id, repo_type="model")
        files = {s.rfilename for s in info.siblings}
        return "adapter_config.json" in files and any(
            f.startswith("adapter_model.") for f in files
        )
    except Exception:
        return False

if _adapter_exists(HF_REPO_ID):
    ADAPTER = HF_REPO_ID
    print(f"[adapter] Using user-fine-tuned: {ADAPTER}")
else:
    ADAPTER = FALLBACK_ADAPTER
    print(f"[adapter] User adapter not found. Falling back: {ADAPTER}")

# ===========================================================================
#                         MODEL LOAD (PtEngine)
# ===========================================================================
from swift.llm import PtEngine, RequestConfig, InferRequest

print(f"Loading base + adapter...")
model_hf = PtEngine(
    model_id_or_path="Qwen/Qwen2.5-VL-7B-Instruct",
    adapters=[ADAPTER],
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.float16,
    ),
    attn_impl="sdpa",
    use_hf=True,
    max_length=2048,
)
req_cfg = RequestConfig(
    max_tokens=20,
    temperature=0.0,
    top_k=1,
    top_p=1.0,
    repetition_penalty=1.0,
)
post_load_mem = get_mem()
print("Model loaded.")

# ===========================================================================
#                  PER-QUESTION CONSTRAINED DECODING
# ===========================================================================

def _normalize_for_match(text):
    """Light normalization for fuzzy comparison. Preserves hyphens and semicolons
    so multi-token refs like '5-10mm', 'pink;red' survive."""
    if not isinstance(text, str): text = str(text)
    t = text.lower().strip()
    if "\n" in t:
        t = t.split("\n", 1)[0].strip()
    for pfx in ("answer:", "final answer:", "the answer is", "it is", "this is"):
        if t.startswith(pfx):
            t = t[len(pfx):].strip()
    t = re.sub(r"[^\w\s\-;]", " ", t)
    t = re.sub(r"\b(a|an|the)\b", " ", t)
    t = re.sub(r"\s+", " ", t).strip()
    return t


def _constrain(pred, question, threshold=0.5):
    """For a model prediction + question, snap to the closest training answer
    seen for that exact question text. Falls back to fuzzy match across all
    training answers if the question is novel. Returns pred unchanged if no
    candidate exceeds threshold."""
    if pred is None:
        pred = ""
    qkey = (question or "").strip().lower()
    pnorm = _normalize_for_match(pred)

    # 1. Direct lookup: exact normalized match against this question's candidates
    candidates = QUESTION_BANK.get(qkey, [])
    if candidates:
        # Try exact normalized match first (strongest signal)
        for cand in candidates[:30]:
            if _normalize_for_match(cand["answer"]) == pnorm:
                return cand["answer"]

    # 2. Empty prediction: fall back to top training answer for this question
    if not pnorm:
        if candidates:
            return candidates[0]["answer"]
        return pred

    # 3. Fuzzy match against top-K candidates for this question
    pool = candidates[:30] if candidates else None
    if pool:
        best_score, best_ans = 0.0, pred
        max_count = max((c["count"] for c in pool), default=1)
        for cand in pool:
            cn = _normalize_for_match(cand["answer"])
            if not cn: continue
            s = SequenceMatcher(None, pnorm, cn).ratio()
            # Bias slightly toward more common answers for this question
            prior = 0.05 * (cand["count"] / max_count)
            score = s + prior
            if score > best_score:
                best_score, best_ans = score, cand["answer"]
        if best_score >= threshold:
            return best_ans

    # 4. Question novel or no good match: light fuzzy across all training answers
    #    Bucket by first word for speed.
    if pnorm:
        first = pnorm.split()[0]
        candidates_fb = [a for a in ALL_TRAINING_ANSWERS
                         if _normalize_for_match(a).startswith(first[:3])]
        if not candidates_fb:
            candidates_fb = ALL_TRAINING_ANSWERS[:5000]  # cap
        best_score, best_ans = 0.0, pred
        for ans in candidates_fb:
            an = _normalize_for_match(ans)
            if not an: continue
            s = SequenceMatcher(None, pnorm, an).ratio()
            if s > best_score:
                best_score, best_ans = s, ans
        if best_score >= threshold + 0.1:  # tighter threshold for unrestricted
            return best_ans

    return pred


# ===========================================================================
#                         INFERENCE LOOP
# ===========================================================================
print("Starting inference...")
_t0 = time.time()
_raw_samples = []

for idx, ex in enumerate(tqdm(val_dataset, desc="Validating")):
    question = ex["question"]
    image    = ex["image"]
    if hasattr(image, "convert") and image.mode != "RGB":
        image = image.convert("RGB")

    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    tmp_path = tmp.name
    tmp.close()
    image.save(tmp_path)

    try:
        req = InferRequest(messages=[{
            "role": "user",
            "content": [
                {"type": "image", "image": tmp_path},
                {"type": "text", "text": question},
            ],
        }])
        resp = model_hf.infer([req], req_cfg)
        raw_answer = resp[0].choices[0].message.content
        if raw_answer is None: raw_answer = ""

        constrained = _constrain(raw_answer, question, threshold=0.5)
    finally:
        try: os.unlink(tmp_path)
        except OSError: pass

    predictions.append({
        "index": idx, "img_id": ex["img_id"],
        "question": question, "answer": constrained,
    })
    if len(_raw_samples) < 30:
        _raw_samples.append((question, raw_answer, constrained))

# ===========================================================================
#                          DIAGNOSTIC
# ===========================================================================
if predictions:
    answers = [p["answer"] for p in predictions]
    word_counts = [len(a.split()) for a in answers]
    freq = Counter(answers)
    print()
    print("=========== TASK 1 DIAGNOSTIC ===========")
    print(f"predictions           : {len(predictions)}")
    print(f"empty answers         : {sum(1 for a in answers if not a.strip())}")
    print(f"avg answer length (w) : {sum(word_counts)/len(word_counts):.2f}")
    print(f"answers with >1 word  : {sum(1 for n in word_counts if n > 1)} ({100*sum(1 for n in word_counts if n > 1)/len(predictions):.1f}%)")
    print(f"answers with >=2 word : {sum(1 for n in word_counts if n >= 2)} ({100*sum(1 for n in word_counts if n >= 2)/len(predictions):.1f}%)")
    print()
    print("Top-30 most frequent predictions:")
    for ans, c in freq.most_common(30):
        print(f"  {c:5d}  {ans!r}")
    print()
    print("Sample of 20 (question | raw | constrained):")
    for q, r, c in _raw_samples[:20]:
        print(f"  Q: {q[:50]!r:55s}  raw: {r!r:50s} -> {c!r}")
    print("===========================================")

# ===========================================================================
#                          FINAL SCORING
# ===========================================================================
preds_texts = [p["answer"] for p in predictions]
refs_texts  = [ex["answer"] for ex in val_dataset]
def _flat(r):
    return "; ".join(str(x) for x in r) if isinstance(r, list) else str(r)
refs_flat = [_flat(r) for r in refs_texts]

# Use BLEU-1 (max_order=1) as the primary BLEU. HuggingFace's default bleu
# is BLEU-4, which is a geometric mean over n=1..4 and collapses to 0 for
# 1-word answers. Per organizer's note (May 21), BLEU-1 is the meaningful
# score for this short-answer task and won't be used as a final ranking
# metric anyway. We still compute BLEU-4 separately for completeness.
bleu_s   = bleu.compute(predictions=preds_texts, references=[[r] for r in refs_flat], max_order=1)
bleu4_s  = bleu.compute(predictions=preds_texts, references=[[r] for r in refs_flat])
rouge_s  = rouge.compute(predictions=preds_texts, references=refs_flat)
meteor_s = meteor.compute(predictions=preds_texts, references=refs_flat)

scores = {
    "bleu":   round(bleu_s["bleu"], 4),    # BLEU-1 (unigram) - meaningful for short answers
    "bleu4":  round(bleu4_s["bleu"], 4),   # BLEU-4 (HF default) - usually 0 for this task
    "rouge1": round(rouge_s["rouge1"], 4),
    "rouge2": round(rouge_s["rouge2"], 4),
    "rougeL": round(rouge_s["rougeL"], 4),
    "meteor": round(meteor_s["meteor"], 4),
}
print(f"\u2728Public scores:  {scores}")

with open("predictions_1.json", "w") as f:
    json.dump(predictions, f, indent=2)

elapsed = time.time() - _t0
print(f"Time: {elapsed:.1f}s | Mem: {get_mem():.2f}MB | "
      f"Model Load Mem: {post_load_mem:.2f}MB | GPU: {gpu_name}")
print("Generation complete. Results saved to 'predictions_1.json'.")
print(f"Run:: medvqa validate_and_submit --competition=gi-2026 --task=1 --repo_id={HF_REPO_ID}")