experiment/evaluation/eval_pope_val.py

"""
POPE-style yes/no evaluation on the HF val split.

Asks the forced-choice question:
    "Is there a {object} in this image? Please answer yes or no."

Computes POPE accuracy / precision / recall / F1 / yes-ratio plus per-category
breakdown (A_B, A_¬B, ¬A_B, ¬A_¬B). Same image loader as eval_efuf_val.py.

Three model variants supported through one CLI:
  - base                                       (no --efuf_ckpt, --model_path=LLaVA base)
  - EFUF   = base + .pth state-dict overlay    (--efuf_ckpt path/to/epoch_NNN.pth)
  - Nullu  = a complete edited-model dir       (--model_path Nullu/output/edited_model/...)

DDP via torchrun. Outputs JSON with per-sample answers + aggregate metrics.

Usage (single-GPU):
  python experiment/evaluation/eval_pope_val.py \
      --relation kitchen_oven --efuf_ckpt path/to/epoch_009.pth \
      --output_dir results/pope/kitchen_oven_efuf

Usage (DDP, e.g. 4 GPUs):
  torchrun --nproc_per_node=4 experiment/evaluation/eval_pope_val.py \
      --relation kitchen_oven --efuf_ckpt path/to/epoch_009.pth \
      --output_dir results/pope/kitchen_oven_efuf
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from datetime import datetime

import numpy as np
import torch
import torch.distributed as dist
from tqdm import tqdm

LLAVA_PATH = "/home/erwin/.cache/huggingface/hub/models--liuhaotian--llava-v1.5-7b/snapshots/4481d270cc22fd5c4d1bb5df129622006ccd9234"

EFUF_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "EFUF", "efuf")
EXPERIMENT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")

_CHAT_PREFIX = (
    "A chat between a curious user and an artificial intelligence assistant. "
    "The assistant gives helpful, detailed, and polite answers to the user's questions. "
)

POPE_TEMPLATE = "Is there a {object} in this image? Please answer yes or no."


# ----------------------------- generation -----------------------------------

def _configure_efuf_args(llava_path: str, device: str, max_new_tokens: int):
    """Inject EFUF-compatible args before importing common.args."""
    efuf_argv = [
        "--model", "llava",
        "--llava_path", llava_path,
        "--llava_ckpt_load_path", llava_path,
        "--device", device,
        "--max_new_tokens", str(max_new_tokens),
        "--llava_data_size_k", "0",
        "--gold_w", "0",
        "--sent_w", "0",
        "--run_name", "eval_pope",
    ]
    saved = sys.argv[:]
    sys.argv = ["eval_pope_val"] + efuf_argv
    return saved


def _restore_argv(saved):
    sys.argv = saved


def _yesno_generate_batch(llava_model_obj, model, vis_processor, images, prompt, device, max_new_tokens):
    """Greedy short generation for yes/no answers (left-padded batched generate)."""
    texts = [f"{_CHAT_PREFIX}USER: <image>\n{prompt} ASSISTANT:" for _ in images]
    ids_list = [llava_model_obj.tokenize_image(t) for t in texts]
    max_len = max(ids.shape[0] for ids in ids_list)
    pad_id = llava_model_obj.tokenizer.pad_token_id

    padded_ids, attn_masks = [], []
    for ids in ids_list:
        pad_len = max_len - ids.shape[0]
        if pad_len > 0:
            padding = torch.full((pad_len,), pad_id, dtype=ids.dtype)
            padded_ids.append(torch.cat([padding, ids]))
            attn_masks.append(torch.cat([torch.zeros(pad_len, dtype=torch.long),
                                          torch.ones(ids.shape[0], dtype=torch.long)]))
        else:
            padded_ids.append(ids)
            attn_masks.append(torch.ones(ids.shape[0], dtype=torch.long))

    input_ids = torch.stack(padded_ids).to(device)
    attention_mask = torch.stack(attn_masks).to(device)
    pixel_values = torch.stack([vis_processor(img) for img in images]).to(device, model.dtype)

    with torch.inference_mode():
        out_ids = model.generate(
            input_ids=input_ids,
            images=pixel_values,
            attention_mask=attention_mask,
            do_sample=False,
            pad_token_id=pad_id,
            max_new_tokens=max_new_tokens,
        )
    new_ids = out_ids[:, input_ids.shape[1]:]
    return [c.strip() for c in llava_model_obj.tokenizer.batch_decode(new_ids, skip_special_tokens=True)]


# ----------------------------- parsing --------------------------------------

_YES_RE = re.compile(r"\byes\b", re.IGNORECASE)
_NO_RE  = re.compile(r"\bno\b",  re.IGNORECASE)


def parse_yesno(text: str):
    """Return True (yes), False (no), or None (unparseable). Trims leading punctuation."""
    t = text.strip().lower().lstrip(".,!?:;\"'`* ")
    if t.startswith("yes"):
        return True
    if t.startswith("no"):
        return False
    has_yes = bool(_YES_RE.search(t))
    has_no = bool(_NO_RE.search(t))
    if has_yes and not has_no:
        return True
    if has_no and not has_yes:
        return False
    return None


# ----------------------------- metrics --------------------------------------

def _safe_div(a, b):
    return float(a / b) if b else None


def compute_metrics(records: list[dict]) -> dict:
    """records: each has {pred: True/False/None, gt: bool, category: str, scene: int, obj: int}"""
    n_total = len(records)
    n_unparsed = sum(1 for r in records if r["pred"] is None)
    parsed = [r for r in records if r["pred"] is not None]

    tp = sum(1 for r in parsed if r["pred"] and r["gt"])
    fp = sum(1 for r in parsed if r["pred"] and not r["gt"])
    fn = sum(1 for r in parsed if not r["pred"] and r["gt"])
    tn = sum(1 for r in parsed if not r["pred"] and not r["gt"])

    acc = _safe_div(tp + tn, len(parsed))
    prec = _safe_div(tp, tp + fp)
    rec = _safe_div(tp, tp + fn)
    f1 = _safe_div(2 * prec * rec, prec + rec) if (prec is not None and rec is not None and (prec + rec) > 0) else None
    yes_ratio = _safe_div(tp + fp, len(parsed))

    # Per-category — keys: A_B, A_no_B, nonA_B, nonA_no_B
    cat_stats: dict[str, dict] = {}
    for r in records:
        c = r["category"]
        s = cat_stats.setdefault(c, {"n": 0, "n_parsed": 0, "yes": 0, "no": 0, "unparsed": 0, "gt_yes": 0})
        s["n"] += 1
        s["gt_yes"] += int(r["gt"])
        if r["pred"] is None:
            s["unparsed"] += 1
        else:
            s["n_parsed"] += 1
            s["yes"] += int(r["pred"])
            s["no"] += int(not r["pred"])

    for c, s in cat_stats.items():
        n_p = s["n_parsed"] or 1
        s["yes_rate"] = s["yes"] / n_p
        s["no_rate"] = s["no"] / n_p
        # Error rate: depends on GT. For all images in a category the GT is the same.
        gt = (s["gt_yes"] > 0)
        s["gt_label"] = "yes" if gt else "no"
        # error = predicting opposite of GT
        s["error_rate"] = s["no_rate"] if gt else s["yes_rate"]

    return {
        "n_total": n_total,
        "n_parsed": len(parsed),
        "n_unparsed": n_unparsed,
        "accuracy": acc,
        "precision": prec,
        "recall": rec,
        "f1": f1,
        "yes_ratio": yes_ratio,
        "confusion": {"tp": tp, "fp": fp, "fn": fn, "tn": tn},
        "per_category": cat_stats,
    }


# ----------------------------- DDP helpers ----------------------------------

def _setup_dist():
    if "LOCAL_RANK" not in os.environ:
        return 0, 1, None
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    return dist.get_rank(), dist.get_world_size(), local_rank


def _gather_list(local: list, world_size: int) -> list:
    if world_size == 1:
        return local
    bucket = [None] * world_size
    dist.all_gather_object(bucket, local)
    out = []
    for part in bucket:
        out.extend(part)
    return out


# ----------------------------- main -----------------------------------------

def parse_args():
    p = argparse.ArgumentParser(description="POPE yes/no eval on HF val split")
    p.add_argument("--relation", type=str, required=True)
    p.add_argument("--model_path", type=str, default=LLAVA_PATH,
                   help="LLaVA model dir. Use a Nullu edited model dir for Nullu eval.")
    p.add_argument("--efuf_ckpt", type=str, default="",
                   help="Optional EFUF .pth to overlay on the loaded model.")
    p.add_argument("--nullu_model_dir", type=str, default="",
                   help="Nullu edited-model dir. Loads base LLaVA, then splices LLM layers from this dir.")
    p.add_argument("--nullu_layers", type=str, default="8-32",
                   help="Layer range to splice from Nullu (start-end, exclusive end). Default: 8-32.")
    p.add_argument("--method", type=str, default="auto", choices=["auto", "base", "efuf", "nullu"],
                   help="Label written into the output JSON. 'auto' infers from args.")
    p.add_argument("--output_dir", type=str, required=True)
    p.add_argument("--question_object", type=str, default=None,
                   help="Object word for the POPE prompt. Defaults to relation's object_key.")
    p.add_argument("--max_new_tokens", type=int, default=5)
    p.add_argument("--max_samples", type=int, default=0,
                   help="Max samples per category; 0 = full val split.")
    p.add_argument("--split", type=str, default="validation")
    p.add_argument("--device", type=str, default="cuda:0")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--batch_size", type=int, default=8)
    return p.parse_args()


def main():
    args = parse_args()
    rank, world_size, local_rank_ddp = _setup_dist()
    is_main = rank == 0
    device = f"cuda:{local_rank_ddp}" if local_rank_ddp is not None else args.device

    sys.path.insert(0, EXPERIMENT_DIR)
    from config.relation_config import get_relation_config
    from data.hf_loader import load_hf_dataset

    rc = get_relation_config(args.relation)
    scene_col, obj_col = rc.scene_key, rc.object_key
    object_word = args.question_object or rc.object_key
    question = POPE_TEMPLATE.format(object=object_word)
    if is_main:
        print(f"[POPE] relation={args.relation}  question={question!r}")

    ds = load_hf_dataset(rc.dataset_id, split=args.split)
    if is_main:
        print(f"[POPE] loaded {rc.dataset_id} split={args.split}: {len(ds)} samples")

    sc_labels = ds[scene_col]
    ob_labels = ds[obj_col]
    # Stratified subsample if requested
    if args.max_samples <= 0:
        all_indices = list(range(len(ds)))
    else:
        rng = np.random.default_rng(args.seed)
        buckets = {("A_B"): [], ("A_no_B"): [], ("nonA_B"): [], ("nonA_no_B"): []}
        for i, (sc, ob) in enumerate(zip(sc_labels, ob_labels)):
            sc, ob = int(sc), int(ob)
            key = ("A_B" if sc and ob else
                   "A_no_B" if sc and not ob else
                   "nonA_B" if (not sc) and ob else
                   "nonA_no_B")
            buckets[key].append(i)
        all_indices = []
        for b in buckets.values():
            rng.shuffle(b)
            all_indices.extend(b[:args.max_samples])
        all_indices.sort()

    if world_size > 1:
        indices = [all_indices[i] for i in range(rank, len(all_indices), world_size)]
    else:
        indices = all_indices

    n_total = len(all_indices); n = len(indices)
    if is_main:
        print(f"[POPE] world_size={world_size}  total={n_total}  this_rank={n}  batch={args.batch_size}")

    saved = _configure_efuf_args(args.model_path, device, args.max_new_tokens)
    sys.path.insert(0, EFUF_DIR)
    from common.models import LlavaModel
    _restore_argv(saved)

    if is_main:
        print(f"[POPE] loading LLaVA model from {args.model_path}")
    llava_model_obj = LlavaModel()
    model, vis_processor = llava_model_obj.load(args.model_path, str(device), train=False)
    model.eval()

    if args.efuf_ckpt:
        if is_main:
            print(f"[POPE] overlaying EFUF ckpt: {args.efuf_ckpt}")
        ck = torch.load(args.efuf_ckpt, map_location=str(device), weights_only=False)
        sd = ck["model"] if isinstance(ck, dict) and "model" in ck else ck
        model.load_state_dict(sd, strict=False)
        model.eval()

    if args.nullu_model_dir:
        if is_main:
            print(f"[POPE] splicing Nullu LLM layers from: {args.nullu_model_dir}  range={args.nullu_layers}")
        from safetensors.torch import load_file
        ls, le = (int(x) for x in args.nullu_layers.split("-"))
        index_path = os.path.join(args.nullu_model_dir, "model.safetensors.index.json")
        single_path = os.path.join(args.nullu_model_dir, "model.safetensors")
        prefixes = tuple(f"model.layers.{i}." for i in range(ls, le))
        if os.path.exists(index_path):
            with open(index_path) as f:
                weight_map = json.load(f)["weight_map"]
            target_keys = [k for k in weight_map if k.startswith(prefixes)]
            shards: dict[str, list[str]] = {}
            for k in target_keys:
                shards.setdefault(weight_map[k], []).append(k)
            partial_sd: dict = {}
            for shard_file, keys in shards.items():
                tensors = load_file(os.path.join(args.nullu_model_dir, shard_file))
                for k in keys:
                    partial_sd[k] = tensors[k]
        else:
            tensors = load_file(single_path)
            partial_sd = {k: v for k, v in tensors.items() if k.startswith(prefixes)}
        tgt_dtype = next(model.parameters()).dtype
        partial_sd = {k: v.to(tgt_dtype) for k, v in partial_sd.items()}
        _, unexpected = model.load_state_dict(partial_sd, strict=False)
        if unexpected and is_main:
            print(f"[POPE] [warn] unexpected keys during Nullu splice: {unexpected[:3]}")
        if is_main:
            print(f"[POPE] spliced {len(partial_sd)} tensors from layers {ls}..{le - 1}")
        model.eval()

    # Method label
    if args.method == "auto":
        if args.efuf_ckpt:
            method = "efuf"
        elif args.nullu_model_dir:
            method = "nullu"
        else:
            method = "base"
    else:
        method = args.method

    records_local: list[dict] = []
    if is_main:
        print(f"[POPE] generating yes/no answers (method={method})")
    for bs in tqdm(range(0, n, args.batch_size), desc="POPE", unit="batch",
                   dynamic_ncols=True, disable=not is_main):
        idxs = indices[bs:bs + args.batch_size]
        rows = [ds[i] for i in idxs]
        imgs = [r["image"].convert("RGB") for r in rows]
        answers = _yesno_generate_batch(
            llava_model_obj, model, vis_processor, imgs, question, device, args.max_new_tokens
        )
        for ans, idx, row in zip(answers, idxs, rows):
            sc, ob = int(row[scene_col]), int(row[obj_col])
            cat = ("A_B" if sc and ob else
                   "A_no_B" if sc and not ob else
                   "nonA_B" if (not sc) and ob else
                   "nonA_no_B")
            gt = bool(ob)  # POPE question is about the object
            pred = parse_yesno(ans)
            records_local.append({
                "index": int(idx),
                "image_id": row.get("image_id") if hasattr(row, "get") else None,
                "scene": sc, "obj": ob, "category": cat,
                "gt": gt, "answer_raw": ans, "pred": pred,
            })

    records = _gather_list(records_local, world_size)

    if not is_main:
        if dist.is_initialized():
            dist.barrier()
            dist.destroy_process_group()
        return

    records.sort(key=lambda r: r["index"])
    metrics = compute_metrics(records)
    out = {
        "relation": args.relation,
        "method": method,
        "model_path": args.model_path,
        "efuf_ckpt": args.efuf_ckpt or None,
        "question": question,
        "split": args.split,
        "max_samples_per_cat": args.max_samples,
        "n_eval": len(records),
        "metrics": metrics,
        "timestamp": datetime.utcnow().isoformat() + "Z",
    }

    os.makedirs(args.output_dir, exist_ok=True)
    metrics_path = os.path.join(args.output_dir, "pope_metrics.json")
    samples_path = os.path.join(args.output_dir, "pope_samples.json")
    with open(metrics_path, "w") as f:
        json.dump(out, f, indent=2)
    with open(samples_path, "w") as f:
        json.dump(records, f, indent=2)

    print(f"\n[POPE] {args.relation} / {method}")
    print(f"  accuracy={metrics['accuracy']:.4f}  precision={metrics['precision']:.4f}  "
          f"recall={metrics['recall']:.4f}  f1={metrics['f1']:.4f}  "
          f"yes_ratio={metrics['yes_ratio']:.4f}")
    print(f"  confusion: {metrics['confusion']}  unparsed={metrics['n_unparsed']}/{metrics['n_total']}")
    for c, s in metrics["per_category"].items():
        print(f"  {c:>10s} n={s['n']:5d} gt={s['gt_label']}  yes_rate={s['yes_rate']:.4f}  "
              f"error_rate={s['error_rate']:.4f}  unparsed={s['unparsed']}")
    print(f"  -> {metrics_path}")

    if dist.is_initialized():
        dist.barrier()
        dist.destroy_process_group()


if __name__ == "__main__":
    main()