Blame for scripts/analyze_chartqa_privileged_teacher_alignment.py

#!/usr/bin/env python3
"""Compare a privileged teacher with the ordinary student on gold answer tokens.

Gold labels are appended only as teacher-forced diagnostic continuations after
both prompts have been constructed.  They are never visible while building the
student prompt, privileged teacher prompt, or independent teacher-gate answer.
Launch with ``torchrun`` to shard rows over all selected GPUs.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from pathlib import Path
from typing import Any, Mapping, Sequence

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from data_utils.chart.teacher_gate import (
    chartqa_reference_answer,
    chartqa_sample_fingerprint,
    default_manifest_path,
    read_jsonl,
    read_manifest,
    validate_teacher_gate_cache,
)
from data_utils.paths import local_pretrained_kwargs, resolve_model_path, validate_local_model_dir
from opsd_utils.opsd_loss import (
    _teacher_logits_with_oom_retry,
    build_opsd_answer_diagnostic_masks,
)
from opsd_utils.teacher_batching import process_teacher_sample, stack_teacher_processor_batches
from opsd_utils.vocab_align import align_cross_model_logits
from reward_utils.teacher_generate import _align_stacked_batch
from scripts.preprocess_chartqa_teacher_gate import (
    _build_job,
    _load_rows,
    _prepare_sample,
    _selected_indices,
)


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--student-model", required=True)
    parser.add_argument("--teacher-model", required=True)
    parser.add_argument("--student-processor", default="")
    parser.add_argument("--teacher-processor", default="")
    parser.add_argument("--teacher-gate-cache", required=True)
    parser.add_argument("--providers", default="visual_facts_deplot")
    parser.add_argument("--max-samples", type=int, default=0)
    parser.add_argument("--sample-strategy", choices=("head", "spread"), default="spread")
    parser.add_argument("--top-k", type=int, default=16)
    parser.add_argument("--max-relative-change", type=float, default=0.05)
    parser.add_argument("--human-only", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument(
        "--attn-implementation",
        choices=("sdpa", "flash_attention_2"),
        default="flash_attention_2",
    )
    parser.add_argument("--log-every", type=int, default=20)
    return parser.parse_args(argv)


def distributed_context() -> tuple[Any, int, int, int]:
    import torch
    import torch.distributed as dist

    rank = int(os.environ.get("RANK", "0"))
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if not torch.cuda.is_available():
        raise RuntimeError("privileged-teacher alignment analysis requires CUDA")
    torch.cuda.set_device(local_rank)
    if world_size > 1 and not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    return torch, rank, world_size, local_rank


def barrier(world_size: int) -> None:
    if world_size > 1:
        import torch.distributed as dist

        dist.barrier()


def load_model_processor(
    model_source: str,
    processor_source: str,
    *,
    torch: Any,
    local_rank: int,
    attn_implementation: str,
) -> tuple[Any, Any, str, str]:
    from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration

    model_path = validate_local_model_dir(resolve_model_path(model_source), role="model")
    processor_path = resolve_model_path(processor_source or model_source)
    if not os.path.isdir(processor_path):
        raise FileNotFoundError(f"processor directory is incomplete or missing: {processor_path}")
    processor = AutoProcessor.from_pretrained(
        processor_path,
        **local_pretrained_kwargs(processor_path),
    )
    processor.tokenizer.padding_side = "left"
    model = LlavaOnevisionForConditionalGeneration.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
        attn_implementation=attn_implementation,
        **local_pretrained_kwargs(model_path),
    ).to(torch.device("cuda", local_rank))
    model.eval()
    model.requires_grad_(False)
    return model, processor, model_path, processor_path


def prepared_prompt(processor: Any, model: Any, text: str, images: list[Any]) -> dict[str, Any]:
    from opsd_utils.teacher_batching import (
        model_inference_device,
        move_batch_num_images_to_model_device,
        move_pixel_values_to_model_device,
    )

    processed = process_teacher_sample(processor, text, images)
    stacked = stack_teacher_processor_batches(processor, [processed])
    ids, mask, pixel_values, image_sizes, batch_num_images = _align_stacked_batch(
        model, processor, stacked
    )
    if isinstance(pixel_values, list):
        raise RuntimeError("single diagnostic row retained incompatible visual tensors")
    device = model_inference_device(model)
    return {
        "prompt_ids": ids.to(device),
        "prompt_mask": mask.to(device),
        "pixel_values": move_pixel_values_to_model_device(model, pixel_values),
        "image_sizes": image_sizes.to(device) if hasattr(image_sizes, "to") else image_sizes,
        "batch_num_images": move_batch_num_images_to_model_device(model, batch_num_images),
    }


def completion_ids_for_gold(tokenizer: Any, gold: str, torch: Any, device: Any) -> Any:
    token_ids = [
        int(value)
        for value in tokenizer.encode(
            f"\nAnswer: {str(gold).strip()}", add_special_tokens=False
        )
    ]
    eos_token_id = getattr(tokenizer, "eos_token_id", None)
    if eos_token_id is None:
        raise ValueError("diagnostic tokenizer has no eos_token_id")
    token_ids.append(int(eos_token_id))
    return torch.tensor([token_ids], dtype=torch.long, device=device)


def score_completion(
    *,
    torch: Any,
    model: Any,
    processor: Any,
    prompt: Mapping[str, Any],
    completion_ids: Any,
) -> Any:
    completion_mask = torch.ones_like(completion_ids)
    return _teacher_logits_with_oom_retry(
        model,
        processor,
        prompt["prompt_ids"],
        prompt["prompt_mask"],
        completion_ids,
        completion_mask,
        prompt["pixel_values"],
        prompt["image_sizes"],
        int(completion_ids.size(1)),
        teacher_batch_num_images=prompt["batch_num_images"],
    )


def mean_selected(values: Any, mask: Any) -> float:
    selected = values[mask]
    return float(selected.mean().item()) if selected.numel() else 0.0


def pair_metrics(
    *,
    torch: Any,
    student_logits: Any,
    teacher_logits: Any,
    completion_ids: Any,
    tokenizer: Any,
    top_k: int,
) -> dict[str, float]:
    import torch.nn.functional as F

    student_logits, teacher_logits = align_cross_model_logits(
        student_logits.float(), teacher_logits.float()
    )
    teacher_logits = teacher_logits.to(student_logits.device)
    completion_ids = completion_ids.to(student_logits.device)
    completion_mask = torch.ones_like(completion_ids)
    masks = build_opsd_answer_diagnostic_masks(
        completion_ids,
        completion_mask,
        tokenizer,
        answer_flag="Answer:",
    )
    answer_mask = masks["answer_span_mask"].bool().to(student_logits.device)
    eos_mask = masks["eos_mask"].bool().to(student_logits.device)
    first_position = int(masks["first_answer_positions"][0].item())
    if first_position < 0 or not bool(answer_mask.any().item()) or not bool(eos_mask.any().item()):
        raise RuntimeError("gold diagnostic continuation lacks answer span or EOS")

    student_log = F.log_softmax(student_logits, dim=-1)
    teacher_log = F.log_softmax(teacher_logits, dim=-1)
    target = completion_ids.unsqueeze(-1)
    student_gold_logp = student_log.gather(-1, target).squeeze(-1)
    teacher_gold_logp = teacher_log.gather(-1, target).squeeze(-1)

    beta = torch.tensor(0.5, device=student_logits.device)
    mixture = torch.logaddexp(
        student_log + torch.log1p(-beta),
        teacher_log + torch.log(beta),
    )
    jsd = 0.5 * (
        F.kl_div(mixture, teacher_log, reduction="none", log_target=True).sum(-1)
        + F.kl_div(mixture, student_log, reduction="none", log_target=True).sum(-1)
    )
    jsd = jsd.clamp_min(0.0)

    selected_s_log = student_log[answer_mask]
    selected_t_log = teacher_log[answer_mask]
    selected_s_prob = selected_s_log.exp()
    selected_t_prob = selected_t_log.exp()
    k = min(max(1, int(top_k)), int(selected_s_log.size(-1)))
    student_topk = torch.topk(selected_s_log, k=k, dim=-1).indices
    teacher_topk = torch.topk(selected_t_log, k=k, dim=-1).indices
    topk_overlap = (
        student_topk.unsqueeze(-1) == teacher_topk.unsqueeze(-2)
    ).any(dim=-1).float().sum(dim=-1) / float(k)
    shared_mass = torch.minimum(selected_s_prob, selected_t_prob).sum(dim=-1)
    student_entropy = -(selected_s_prob * selected_s_log).sum(dim=-1)
    teacher_entropy = -(selected_t_prob * selected_t_log).sum(dim=-1)
    entropy_gap = teacher_entropy - student_entropy

    first_student_prob = selected_s_prob[0]
    first_teacher_prob = selected_t_prob[0]
    first_student_topk = torch.topk(selected_s_log[0], k=k).indices
    first_teacher_topk = torch.topk(selected_t_log[0], k=k).indices
    first_overlap = (
        first_student_topk.unsqueeze(-1) == first_teacher_topk.unsqueeze(-2)
    ).any(dim=-1).float().sum() / float(k)
    first_shared = torch.minimum(first_student_prob, first_teacher_prob).sum()
    first_student_entropy = -(first_student_prob * selected_s_log[0]).sum()
    first_teacher_entropy = -(first_teacher_prob * selected_t_log[0]).sum()

    return {
        "first_answer_student_gold_logp": float(student_gold_logp[0, first_position].item()),
        "first_answer_teacher_gold_logp": float(teacher_gold_logp[0, first_position].item()),
        "first_answer_teacher_gold_logp_advantage": float(
            (teacher_gold_logp[0, first_position] - student_gold_logp[0, first_position]).item()
        ),
        "answer_span_student_gold_logp": mean_selected(student_gold_logp, answer_mask),
        "answer_span_teacher_gold_logp": mean_selected(teacher_gold_logp, answer_mask),
        "answer_span_teacher_gold_logp_advantage": mean_selected(
            teacher_gold_logp - student_gold_logp, answer_mask
        ),
        "first_answer_jsd": float(jsd[0, first_position].item()),
        "answer_span_jsd": mean_selected(jsd, answer_mask),
        "eos_jsd": mean_selected(jsd, eos_mask),
        "first_answer_topk_overlap": float(first_overlap.item()),
        "answer_span_topk_overlap": float(topk_overlap.mean().item()),
        "first_answer_shared_probability_mass": float(first_shared.item()),
        "answer_span_shared_probability_mass": float(shared_mass.mean().item()),
        "first_answer_teacher_minus_student_entropy": float(
            (first_teacher_entropy - first_student_entropy).item()
        ),
        "answer_span_teacher_minus_student_entropy": float(entropy_gap.mean().item()),
        "answer_span_entropy_abs_gap": float(entropy_gap.abs().mean().item()),
        "answer_token_count": float(answer_mask.sum().item()),
    }


def atomic_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}")
    with temporary.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(dict(row), ensure_ascii=False, sort_keys=True) + "\n")
    os.replace(temporary, path)


def atomic_json(path: Path, payload: Mapping[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}")
    with temporary.open("w", encoding="utf-8") as handle:
        json.dump(dict(payload), handle, ensure_ascii=False, indent=2, sort_keys=True)
        handle.write("\n")
    os.replace(temporary, path)


def summarize(rows: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
    numeric_keys = sorted(
        {
            key
            for row in rows
            for key, value in row.items()
            if isinstance(value, (int, float)) and key not in {"source_index"}
        }
    )

    def group_summary(group: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
        means: dict[str, float] = {}
        for key in numeric_keys:
            values = [float(row[key]) for row in group if key in row and math.isfinite(float(row[key]))]
            if values:
                means[key] = sum(values) / len(values)
        return {"count": len(group), "means": means}

    correct = [row for row in rows if row.get("teacher_correct") is True]
    incorrect = [row for row in rows if row.get("teacher_correct") is False]
    return {
        "all": group_summary(rows),
        "teacher_correct": group_summary(correct),
        "teacher_incorrect": group_summary(incorrect),
        "teacher_correct_rate": len(correct) / len(rows) if rows else 0.0,
    }


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_args(argv)
    if args.max_samples < 0 or args.top_k <= 0 or args.log_every <= 0:
        raise ValueError("max-samples must be non-negative and top-k/log-every positive")
    providers = [value.strip() for value in str(args.providers).split(",") if value.strip()]
    if providers != ["visual_facts_deplot"]:
        raise ValueError("this gold-blind diagnostic currently requires visual_facts_deplot only")

    raw_rows = _load_rows(args.input, human_only=bool(args.human_only))
    gate_manifest = read_manifest(default_manifest_path(args.teacher_gate_cache))
    if int(gate_manifest.get("record_count", -1)) == len(raw_rows):
        gate_records, gate_manifest = validate_teacher_gate_cache(
            raw_rows,
            cache_path=args.teacher_gate_cache,
            require_validation_calibration=False,
        )
    else:
        gate_records = read_jsonl(args.teacher_gate_cache)
        if int(gate_manifest.get("source_row_count", -1)) != len(raw_rows):
            raise ValueError("subset teacher-gate manifest has the wrong source row count")
        manifest_indices = gate_manifest.get("selected_source_indices")
        record_indices = [int(record.get("source_index", -1)) for record in gate_records]
        if not isinstance(manifest_indices, list) or record_indices != [
            int(index) for index in manifest_indices
        ]:
            raise ValueError("subset teacher-gate source indices do not match its manifest")
        for record in gate_records:
            source_index = int(record["source_index"])
            if source_index < 0 or source_index >= len(raw_rows):
                raise ValueError("subset teacher-gate source index is outside the dataset")
            expected_fingerprint = chartqa_sample_fingerprint(raw_rows[source_index])
            if str(record.get("sample_fingerprint") or "") != expected_fingerprint:
                raise ValueError(
                    f"subset teacher-gate fingerprint mismatch at source row {source_index}"
                )
    gate_by_source_index = {
        int(record["source_index"]): record for record in gate_records
    }
    gate_teacher = str(Path(str(gate_manifest.get("teacher_model") or "")).resolve())
    requested_teacher = str(Path(resolve_model_path(args.teacher_model)).resolve())
    if gate_teacher != requested_teacher:
        raise ValueError(
            "teacher-gate cache model does not match diagnostic teacher: "
            f"cache={gate_teacher}, requested={requested_teacher}"
        )

    torch, rank, world_size, local_rank = distributed_context()
    selected = _selected_indices(len(raw_rows), int(args.max_samples), args.sample_strategy)
    missing_gate_rows = [index for index in selected if index not in gate_by_source_index]
    if missing_gate_rows:
        raise ValueError(
            "teacher-gate cache does not cover selected diagnostic rows: "
            f"count={len(missing_gate_rows)}, examples={missing_gate_rows[:8]}"
        )
    local_indices = selected[rank::world_size]
    student_model, student_processor, student_model_path, student_processor_path = (
        load_model_processor(
            args.student_model,
            args.student_processor,
            torch=torch,
            local_rank=local_rank,
            attn_implementation=args.attn_implementation,
        )
    )
    same_model = str(Path(student_model_path).resolve()) == requested_teacher
    if same_model and (not args.teacher_processor or str(Path(student_processor_path).resolve()) == str(Path(resolve_model_path(args.teacher_processor or args.teacher_model)).resolve())):
        teacher_model = student_model
        teacher_processor = student_processor
        teacher_model_path = student_model_path
        teacher_processor_path = student_processor_path
    else:
        teacher_model, teacher_processor, teacher_model_path, teacher_processor_path = (
            load_model_processor(
                args.teacher_model,
                args.teacher_processor,
                torch=torch,
                local_rank=local_rank,
                attn_implementation=args.attn_implementation,
            )
        )
    if len(student_processor.tokenizer) != len(teacher_processor.tokenizer):
        raise ValueError("student and teacher tokenizers have different vocabulary sizes")

    records: list[dict[str, Any]] = []
    started = time.monotonic()
    for local_position, source_index in enumerate(local_indices, start=1):
        sample = _prepare_sample(raw_rows[source_index])
        teacher_job = _build_job(
            sample,
            providers=providers,
            source_index=source_index,
        )
        from opsd_utils.privileged.image_utils import load_rgb

        student_image = load_rgb(sample.get("image"))
        if student_image is None:
            raise ValueError(f"could not load student image for source row {source_index}")
        student_prompt = prepared_prompt(
            student_processor,
            student_model,
            str(sample["prompt"]),
            [student_image],
        )
        teacher_prompt = prepared_prompt(
            teacher_processor,
            teacher_model,
            str(teacher_job["prompt"]),
            list(teacher_job["images"]),
        )
        gold = chartqa_reference_answer(sample)
        student_completion = completion_ids_for_gold(
            student_processor.tokenizer,
            gold,
            torch,
            student_prompt["prompt_ids"].device,
        )
        teacher_completion = completion_ids_for_gold(
            teacher_processor.tokenizer,
            gold,
            torch,
            teacher_prompt["prompt_ids"].device,
        )
        if student_completion.detach().cpu().tolist() != teacher_completion.detach().cpu().tolist():
            raise ValueError("student and teacher tokenize the gold continuation differently")
        with torch.inference_mode():
            student_logits = score_completion(
                torch=torch,
                model=student_model,
                processor=student_processor,
                prompt=student_prompt,
                completion_ids=student_completion,
            )
            teacher_logits = score_completion(
                torch=torch,
                model=teacher_model,
                processor=teacher_processor,
                prompt=teacher_prompt,
                completion_ids=teacher_completion,
            )
        metric_values = pair_metrics(
            torch=torch,
            student_logits=student_logits,
            teacher_logits=teacher_logits,
            completion_ids=student_completion,
            tokenizer=student_processor.tokenizer,
            top_k=int(args.top_k),
        )
        gate = gate_by_source_index[source_index]
        records.append(
            {
                "source_index": int(source_index),
                "teacher_correct": bool(gate.get("teacher_correct") is True),
                "teacher_output": str(gate.get("teacher_output") or ""),
                "gold_answer": gold,
                **metric_values,
            }
        )
        if local_position == len(local_indices) or local_position % int(args.log_every) == 0:
            elapsed = time.monotonic() - started
            print(
                f"[teacher-alignment][rank={rank}/{world_size}] "
                f"{local_position}/{len(local_indices)} elapsed={elapsed:.1f}s "
                f"rate={local_position / max(elapsed, 1e-9):.3f}/s",
                flush=True,
            )

    output = Path(args.output).resolve()
    parts = Path(f"{output}.parts")
    part = parts / f"world{world_size:03d}-rank{rank:03d}.jsonl"
    atomic_jsonl(part, records)
    barrier(world_size)
    if rank == 0:
        merged: list[dict[str, Any]] = []
        for part_rank in range(world_size):
            merged.extend(
                json.loads(line)
                for line in (
                    parts / f"world{world_size:03d}-rank{part_rank:03d}.jsonl"
                ).read_text(encoding="utf-8").splitlines()
                if line.strip()
            )
        merged.sort(key=lambda row: int(row["source_index"]))
        if [int(row["source_index"]) for row in merged] != selected:
            raise ValueError("merged alignment rows are incomplete or duplicated")
        atomic_jsonl(output, merged)
        summary = {
            "input": str(Path(args.input).resolve()),
            "output": str(output),
            "student_model": str(Path(student_model_path).resolve()),
            "teacher_model": str(Path(teacher_model_path).resolve()),
            "student_processor": str(Path(student_processor_path).resolve()),
            "teacher_processor": str(Path(teacher_processor_path).resolve()),
            "teacher_gate_cache": str(Path(args.teacher_gate_cache).resolve()),
            "teacher_gate_manifest": str(
                Path(default_manifest_path(args.teacher_gate_cache)).resolve()
            ),
            "providers": providers,
            "top_k": int(args.top_k),
            "world_size": int(world_size),
            "selected_samples": len(selected),
            "gold_usage": "teacher_forced_diagnostics_only_never_prompt_or_routing",
            "summary": summarize(merged),
        }
        atomic_json(Path(f"{output}.summary.json"), summary)
        print(
            "[teacher-alignment][complete] "
            + json.dumps(summary, ensure_ascii=False, sort_keys=True),
            flush=True,
        )
    barrier(world_size)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())