Blame for scripts/convert_v6.py

#!/usr/bin/env python3
"""Convert the shipping v6 dataset into the trainer's TokenEvalRow schema.

v6 (hallucination_labelling/DATASET_CARD_v6.md) ships in the labelling-pipeline
schema with THREE label sets per row -- the same panel votes under different
aggregation rules. This converter:

  * selects ONE label set as `labels` (--labels, default `main` =
    `token_labels` = sentence:2:union, the card's shipping labels) and carries
    all three in metadata so the choice is auditable and reversible;
  * preserves `extra` verbatim -- replay reads the generation context from
    `extra.context` / `extra.knowledge`, and the card's per-arm analysis needs
    `extra.arm`, `extra.prompt_regime` and `extra.n_judges_resolved`;
  * REJECTS NOTHING on content. build_v6 already applied the card's exclusion
    rules (4-gram degeneracy, panel < 2); adding conversion-time rules here
    would silently fork the shipped resource. In particular, no uncertain-
    fraction filter: 22% of v6 rows are majority-U under the union rule, and U
    tokens are already excluded per-token at metric time;
  * treats any structural defect as a HARD FAILURE of the whole conversion,
    not a skipped row -- the shipping dataset is supposed to be perfect, and a
    silent skip would hide a build regression.

Usage:
    python scripts/convert_v6.py \
        --input  ~/Desktop/Datasets/hallucinations_custom/hallu_dataset_v6.jsonl \
        --output ~/Desktop/Datasets/hallucinations_custom/hallu_v6_eval.jsonl \
        --labels main
"""

from __future__ import annotations

import argparse
import collections
import hashlib
import json
from pathlib import Path

GROUNDED, HALLUCINATED, UNCERTAIN = 0, 1, 2
LETTER = {"G": GROUNDED, "H": HALLUCINATED, "U": UNCERTAIN}
VARIANT_FIELD = {"main": "token_labels", "strict": "token_labels_strict",
                 "plurality": "token_labels_plurality"}
TSTAR_FIELD = {"main": "t_star_index", "strict": "t_star_index_strict",
               "plurality": None}
GENERATOR = "meta-llama/Llama-3.2-1B-Instruct"


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--input", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--labels", choices=sorted(VARIANT_FIELD), default="main",
                   help="which of the card's three label sets becomes `labels`")
    return p.parse_args(argv)


def spans_from_labels(labels, offsets):
    """Contiguous H runs -> character spans; runs close on unusable offsets."""
    spans, start_char, end_char = [], None, None
    for index, label in enumerate(labels):
        offset = offsets[index] if index < len(offsets) else None
        usable = (isinstance(offset, (list, tuple)) and len(offset) == 2
                  and offset[0] is not None and int(offset[1]) > int(offset[0]))
        if label == HALLUCINATED and usable:
            if start_char is None:
                start_char = int(offset[0])
            end_char = int(offset[1])
        elif start_char is not None:
            spans.append({"start": start_char, "end": end_char})
            start_char, end_char = None, None
    if start_char is not None:
        spans.append({"start": start_char, "end": end_char})
    return spans


def convert(record, variant):
    rid = record["id"]

    def fail(reason):
        raise SystemExit(f"STRUCTURAL DEFECT in shipping dataset at {rid}: {reason}. "
                         "Refusing to convert; the v6 build needs fixing, not a skip.")

    token_ids = record.get("gen_token_ids") or fail("no gen_token_ids")
    tokens = record.get("gen_tokens") or fail("no gen_tokens")
    offsets = record.get("gen_token_offsets") or fail("no gen_token_offsets")
    if record.get("generator_model") != GENERATOR:
        fail(f"generator {record.get('generator_model')!r} != {GENERATOR!r}")

    variants = {}
    for name, field in VARIANT_FIELD.items():
        raw = record.get(field) or fail(f"missing {field}")
        if len(raw) != len(token_ids):
            fail(f"{field} length {len(raw)} != spine {len(token_ids)}")
        try:
            variants[name] = [LETTER[v] for v in raw]
        except KeyError as exc:
            fail(f"unknown label value {exc} in {field}")
    if len(offsets) != len(token_ids):
        fail("offsets length != spine")

    labels = variants[variant]
    tstar_field = TSTAR_FIELD[variant]
    t_star = record.get(tstar_field) if tstar_field else None
    if t_star is None:
        t_star = next((i for i, v in enumerate(labels) if v == HALLUCINATED), None)

    return {
        "id": rid,
        # Bare question; the chat template and the generation context (from
        # extra.context / extra.knowledge) are applied at replay time.
        "prompt": record["question"],
        "answer": record["generated_answer"],
        "labels": labels,
        "hallucinated_spans": spans_from_labels(labels, offsets),
        "metadata": {
            "source": record["source"],
            "dataset": "hallu_dataset_v6",
            "generator_model": GENERATOR,
            "generator_backend": record.get("generator_backend"),
            "generated_token_ids": list(token_ids),
            "generated_tokens": list(tokens),
            "generated_logprobs": record.get("gen_logprobs"),
            "auto_label_method": record.get("auto_label_method"),
            "reference_answer": record.get("reference_answer"),
            "t_star_index": t_star,
            "label_variant_used": VARIANT_FIELD[variant],
            "label_variants": {n: v for n, v in variants.items() if n != variant},
            "original_label_counts": dict(collections.Counter(
                {0: "grounded", 1: "hallucinated", 2: "uncertain"}[v] for v in labels)),
            "extra": record.get("extra") or {},
        },
    }


def main(argv=None):
    args = parse_args(argv)
    rows = []
    with args.input.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                rows.append(convert(json.loads(line), args.labels))

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")

    counts = collections.Counter(v for r in rows for v in r["labels"])
    gh = counts[GROUNDED] + counts[HALLUCINATED]
    arms = collections.Counter(r["metadata"]["extra"].get("arm") for r in rows)
    print(f"wrote {len(rows):,} rows -> {args.output}  (variant: {VARIANT_FIELD[args.labels]})")
    print(f"  tokens {sum(counts.values()):,} | G {counts[GROUNDED]:,} "
          f"H {counts[HALLUCINATED]:,} U {counts[UNCERTAIN]:,}")
    print(f"  H prevalence over G/H: {100*counts[HALLUCINATED]/gh:.3f}%")
    print(f"  arms: {dict(arms)}")
    print(f"  sha256: {hashlib.sha256(args.output.read_bytes()).hexdigest()}")
    return 0


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