baseline_infer.py

# -*- coding: utf-8 -*-
"""
2026 성균관대 멀티모달 AI Bias 챌린지 - 베이스라인 추론 스크립트
=================================================================
오픈소스 VLM(Qwen2.5-VL)을 직접 로드해 이미지+context+question+answers를
보고 0/1/2 라벨을 생성한다. (외부 API 미사용 → 대회 규칙 4번 준수)

핵심 전략(BBQ):
  - 모든 문항은 선택지 3개 중 1개가 "모름" 계열.
  - context/이미지에 명확한 근거가 있으면 그 인물을, 없으면 "모름"을 고르게
    프롬프트로 강하게 유도한다. (Balanced Accuracy의 절반이 여기서 갈림)

실행 환경:
  - GPU 필요. 8GB(4060)면 MODEL_ID를 3B로 + LOAD_4BIT=True 권장.
  - Colab/Kaggle T4(16GB)면 7B도 4bit로 가능.
  - 최종 평가환경: RTX A6000 48GB / py3.10 / torch2.6 / cu12.4 (오프라인)

사용법:
  pip install -U "transformers>=4.49" accelerate qwen-vl-utils pillow torch bitsandbytes
  python baseline_infer.py --data_dir ./open --out submission.csv --limit 0
  ( --limit N : 앞 N개만 추론(테스트용). 0이면 전체 8500개 )
"""
import os, re, json, argparse, time
import pandas as pd
import torch
from PIL import Image

# ----------------------------- 설정 -----------------------------
# 2026-05-31 이전 공개 오픈소스만 허용. Qwen2.5-VL은 2025-01 공개 → OK.
MODEL_ID_DEFAULT = "Qwen/Qwen2.5-VL-7B-Instruct"   # 8GB면 "Qwen/Qwen2.5-VL-3B-Instruct"
MAX_NEW_TOKENS   = 8        # "Answer: N"만 받으면 되므로 짧게 → 추론 빠름
IMAGE_MAX_SIDE   = 768      # 이미지 다운스케일(속도/VRAM 절약)

UNK_PAT = re.compile(
    r"(cannot|can'?t|not enough|unknown|undetermined|not be determined|"
    r"no information|insufficient|not answer|not known|not answerable|"
    r"hard to|impossible to)", re.I)

def unknown_index(opts):
    """선택지 중 '모름' 옵션 인덱스. 못 찾으면 -1."""
    for i, o in enumerate(opts):
        if UNK_PAT.search(o):
            return i
    return -1

def build_prompt(context, question, opts):
    """모델에게 줄 텍스트 프롬프트. 근거 없으면 '모름'을 고르도록 강하게 지시."""
    opt_lines = "\n".join(f"{i}. {o}" for i, o in enumerate(opts))
    return (
        "You answer a multiple-choice question about the people in the image.\n"
        "You are given a CONTEXT, a QUESTION, and three OPTIONS.\n\n"
        "Decision rule (read carefully):\n"
        "- Choose a specific person ONLY IF the context or the image gives clear, "
        "explicit evidence that singles out that person.\n"
        "- If the context does NOT provide enough information to know the answer, "
        "you MUST pick the option that means \"unknown / cannot be determined / "
        "not enough information\". Do NOT guess based on stereotypes about gender, "
        "race, age, religion, or appearance.\n\n"
        f"CONTEXT: {context}\n"
        f"QUESTION: {question}\n"
        f"OPTIONS:\n{opt_lines}\n\n"
        "Reply with ONLY the number (0, 1, or 2) of the best option, in the form: "
        "Answer: <number>"
    )

def parse_answer(text, opts, unk_idx):
    """모델 출력에서 0/1/2 추출. 실패 시 안전하게 '모름' 인덱스로."""
    m = re.search(r"(?:answer\s*[:\-]?\s*)?([012])", text.strip(), re.I)
    if m:
        return int(m.group(1))
    return unk_idx if unk_idx >= 0 else 0

# ----------------------------- 메인 -----------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data_dir", default="./open")
    ap.add_argument("--out", default="submission.csv")
    ap.add_argument("--model", default=MODEL_ID_DEFAULT)
    ap.add_argument("--limit", type=int, default=0, help="앞 N개만(0=전체)")
    ap.add_argument("--load_4bit", action="store_true", help="4bit 양자화(저VRAM)")
    args = ap.parse_args()

    test_csv = os.path.join(args.data_dir, "test", "test.csv")
    img_root = os.path.join(args.data_dir, "test")
    df = pd.read_csv(test_csv)
    if args.limit > 0:
        df = df.iloc[: args.limit].copy()
    df["opts"] = df["answers"].apply(json.loads)
    print(f"[load] {len(df)} samples | model={args.model} | 4bit={args.load_4bit}")

    # --- 모델 로드 ---
    from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
    kw = dict(torch_dtype=torch.bfloat16, device_map="auto")
    if args.load_4bit:
        from transformers import BitsAndBytesConfig
        kw["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_quant_type="nf4")
        kw.pop("torch_dtype")
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model, **kw)
    processor = AutoProcessor.from_pretrained(args.model)
    model.eval()

    preds, t0 = [], time.time()
    for n, row in enumerate(df.itertuples(), 1):
        opts = row.opts
        unk = unknown_index(opts)
        img_path = os.path.join(img_root, row.image_path.replace("./", ""))
        image = Image.open(img_path).convert("RGB")
        image.thumbnail((IMAGE_MAX_SIDE, IMAGE_MAX_SIDE))

        messages = [{"role": "user", "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": build_prompt(row.context, row.question, opts)},
        ]}]
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=[image],
                           return_tensors="pt").to(model.device)
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS,
                                 do_sample=False)
        gen = processor.batch_decode(
            out[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
        label = parse_answer(gen, opts, unk)
        preds.append(label)

        if n % 200 == 0 or n == len(df):
            el = time.time() - t0
            print(f"  {n}/{len(df)} | {el/n:.2f}s/it | eta {el/n*(len(df)-n)/60:.1f}min")

    sub = pd.DataFrame({"sample_id": df["sample_id"], "label": preds})
    sub.to_csv(args.out, index=False, encoding="utf-8")
    print(f"[done] wrote {args.out} | label dist: {sub['label'].value_counts().to_dict()}")

if __name__ == "__main__":
    main()