halueval-llama-20260721/scripts/prepare_halueval_gpu.py

#!/usr/bin/env python3
"""Generate and TrueTeacher-label Mistral summaries for HaluEval.

The output is four prepared JSONL files: 8,000 train-pool and 2,000 test
records, each in full-answer and first-sentence views.  The paper's 7,200/800
probe train/validation split is made later from the 8,000-record train pool.
"""

from __future__ import annotations

import argparse
import gc
import hashlib
import json
import random
import sys
import time
import urllib.request
from collections import Counter
from pathlib import Path

import numpy as np

from repro_common import first_sentence_truncation, seed_everything, write_json


DATA_URL = "https://raw.githubusercontent.com/RUCAIBox/HaluEval/main/data/summarization_data.json"
DEFAULT_MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
DEFAULT_MODEL_REVISION = "c170c708c41dac9275d15a8fff4eca08d52bab71"
JUDGE_ID = "google/t5_11b_trueteacher_and_anli"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--cache-dir", type=Path, default=Path("/tmp/hf-cache"))
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID)
    parser.add_argument("--model-revision", default=DEFAULT_MODEL_REVISION)
    parser.add_argument("--seed", type=int, default=2024)
    parser.add_argument("--generation-batch-size", type=int, default=8)
    parser.add_argument("--judge-batch-size", type=int, default=4)
    parser.add_argument("--max-new-tokens", type=int, default=130)
    parser.add_argument("--max-samples", type=int)
    parser.add_argument("--start-index", type=int, default=0)
    parser.add_argument("--end-index", type=int)
    parser.add_argument("--checkpoint-every", type=int, default=200)
    parser.add_argument("--self-test", action="store_true")
    return parser.parse_args()


def summary_prompt(context: str) -> str:
    return (
        "Summarize the following document in one or two concise sentences.\n"
        f"Document:{context.strip()}\n"
        "Summary:"
    )


def load_records(cache_dir: Path) -> tuple[list[dict], str]:
    cache_dir.mkdir(parents=True, exist_ok=True)
    path = cache_dir / "halueval-summarization.jsonl"
    if not path.exists():
        urllib.request.urlretrieve(DATA_URL, path)
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    records = []
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                row = json.loads(line)
                if not row.get("document"):
                    raise ValueError("HaluEval record is missing `document`")
                records.append(row)
    if len(records) != 10_000:
        raise ValueError(f"expected 10000 HaluEval records, observed {len(records)}")
    return records, digest


def paper_split(total: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
    from sklearn.model_selection import train_test_split

    train, test = train_test_split(
        np.arange(total),
        test_size=0.2,
        random_state=seed,
    )
    return np.asarray(train), np.asarray(test)


def atomic_json(path: Path, value: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, separators=(",", ":")) + "\n",
        encoding="utf-8",
    )
    temporary.replace(path)


def generate_summaries(
    contexts: list[str],
    args: argparse.Namespace,
    checkpoint_path: Path,
) -> tuple[list[str], float]:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(
        args.model_id,
        revision=args.model_revision,
        cache_dir=args.cache_dir,
    )
    tokenizer.padding_side = "left"
    tokenizer.truncation_side = "left"
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        revision=args.model_revision,
        cache_dir=args.cache_dir,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
        device_map="auto",
        attn_implementation="sdpa",
    )
    model.eval()
    prompts = [summary_prompt(context) for context in contexts]
    lengths = tokenizer(
        prompts,
        padding=False,
        truncation=False,
        return_length=True,
    )["length"]
    answers: list[str | None] = [None] * len(prompts)
    if checkpoint_path.exists():
        checkpoint = json.loads(checkpoint_path.read_text(encoding="utf-8"))
        if len(checkpoint.get("answers", [])) != len(answers):
            raise ValueError(f"checkpoint length mismatch: {checkpoint_path}")
        answers = checkpoint["answers"]
        print(f"resuming generation with {sum(value is not None for value in answers)} completed", flush=True)
    order = sorted(
        (index for index in range(len(prompts)) if answers[index] is None),
        key=lambda index: lengths[index],
    )
    started = time.perf_counter()
    input_device = model.get_input_embeddings().weight.device

    completed_since_checkpoint = 0
    for start in range(0, len(order), args.generation_batch_size):
        indices = order[start : start + args.generation_batch_size]
        batch_prompts = [prompts[index] for index in indices]
        encoded = tokenizer(
            batch_prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=4096,
        ).to(input_device)
        with torch.inference_mode():
            generated = model.generate(
                **encoded,
                max_new_tokens=args.max_new_tokens,
                min_new_tokens=1,
                do_sample=True,
                temperature=0.1,
                pad_token_id=tokenizer.pad_token_id,
                use_cache=True,
            )
        continuation = generated[:, encoded["input_ids"].shape[1] :].cpu()
        decoded = tokenizer.batch_decode(continuation, skip_special_tokens=True)
        for index, answer in zip(indices, decoded, strict=True):
            # The released summarization path truncates the ordinary view at
            # the first generated newline.
            answers[index] = answer.split("\n", 1)[0].strip()
        completed_since_checkpoint += len(indices)
        if completed_since_checkpoint >= args.checkpoint_every:
            atomic_json(checkpoint_path, {"answers": answers})
            completed_since_checkpoint = 0
        if start % (args.generation_batch_size * 25) == 0:
            print(f"generation {min(start + len(indices), len(order))}/{len(order)}", flush=True)
        del encoded, generated, continuation

    elapsed = time.perf_counter() - started
    atomic_json(checkpoint_path, {"answers": answers, "complete": True})
    del model, tokenizer
    gc.collect()
    torch.cuda.empty_cache()
    if any(value is None for value in answers):
        raise RuntimeError("generation checkpoint is incomplete")
    return [str(value) for value in answers], elapsed


def trueteacher_labels(
    contexts: list[str],
    summaries: list[str],
    args: argparse.Namespace,
    *,
    view: str,
    checkpoint_path: Path,
) -> tuple[list[int], float]:
    import torch
    from transformers import T5ForConditionalGeneration, T5Tokenizer

    tokenizer = T5Tokenizer.from_pretrained(JUDGE_ID, cache_dir=args.cache_dir)
    model = T5ForConditionalGeneration.from_pretrained(
        JUDGE_ID,
        cache_dir=args.cache_dir,
        low_cpu_mem_usage=True,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    )
    model.eval()
    input_device = next(model.parameters()).device
    labels: list[int | None] = [None] * len(summaries)
    if checkpoint_path.exists():
        checkpoint = json.loads(checkpoint_path.read_text(encoding="utf-8"))
        if len(checkpoint.get("labels", [])) != len(labels):
            raise ValueError(f"checkpoint length mismatch: {checkpoint_path}")
        labels = checkpoint["labels"]
        print(f"resuming judge[{view}] with {sum(value is not None for value in labels)} completed", flush=True)
    started = time.perf_counter()
    pending = [index for index, value in enumerate(labels) if value is None]
    completed_since_checkpoint = 0

    for start in range(0, len(pending), args.judge_batch_size):
        batch_indices = pending[start : start + args.judge_batch_size]
        batch_contexts = [contexts[index] for index in batch_indices]
        batch_summaries = [summaries[index] for index in batch_indices]
        valid_positions = [index for index, value in enumerate(batch_summaries) if value.strip()]
        inputs = [
            f"premise: {batch_contexts[index].strip()} hypothesis: {batch_summaries[index].strip()}"
            for index in valid_positions
        ]
        if inputs:
            encoded = tokenizer(
                inputs,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=2048,
            ).input_ids.to(input_device)
            with torch.inference_mode():
                generated = model.generate(input_ids=encoded)
            decoded = tokenizer.batch_decode(generated, skip_special_tokens=True)
            for position, text in zip(valid_positions, decoded, strict=True):
                labels[batch_indices[position]] = 1 if text.strip().startswith("1") else 0
            del encoded, generated
        for position, summary in enumerate(batch_summaries):
            if not summary.strip():
                labels[batch_indices[position]] = 0
        completed_since_checkpoint += len(batch_indices)
        if completed_since_checkpoint >= args.checkpoint_every:
            atomic_json(checkpoint_path, {"labels": labels})
            completed_since_checkpoint = 0
        if start % (args.judge_batch_size * 50) == 0:
            print(f"judge[{view}] {min(start + len(batch_summaries), len(summaries))}/{len(summaries)}", flush=True)

    elapsed = time.perf_counter() - started
    atomic_json(checkpoint_path, {"labels": labels, "complete": True})
    del model, tokenizer
    gc.collect()
    torch.cuda.empty_cache()
    if any(value is None for value in labels):
        raise RuntimeError(f"judge[{view}] checkpoint is incomplete")
    return [int(value) for value in labels], elapsed


def write_part_view(
    path: Path,
    records: list[dict],
    source_indices: list[int],
    summaries: list[str],
    labels: list[int],
    view: str,
    train_indices: set[int],
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for local_index, source_index in enumerate(source_indices):
            context = records[local_index]["document"].strip()
            answer = summaries[local_index].strip()
            row = {
                "source_index": int(source_index),
                "paper_split": "train" if source_index in train_indices else "test",
                "context": context,
                "question": "Summarize the document.",
                "best_answer": answer,
                "label": int(labels[local_index]),
                "answer_view": view,
                "text": f"{summary_prompt(context)} {answer}",
            }
            handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")


def self_test() -> None:
    train, test = paper_split(10_000, 2024)
    assert len(train) == 8_000 and len(test) == 2_000
    assert len(set(train) & set(test)) == 0
    assert first_sentence_truncation("One fact. Unsupported continuation.") == "One fact."
    assert summary_prompt(" doc ").endswith("Summary:")
    print("self-test passed")


def main() -> None:
    args = parse_args()
    if args.self_test:
        self_test()
        return
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("HaluEval preparation requires a CUDA GPU")
    seed_everything(args.seed)
    random.seed(args.seed)
    records, dataset_sha256 = load_records(args.cache_dir)
    total_records = len(records)
    train_indices, test_indices = paper_split(total_records, args.seed)
    start_index = args.start_index
    end_index = args.end_index if args.end_index is not None else total_records
    if args.max_samples is not None:
        end_index = min(end_index, start_index + args.max_samples)
    if not 0 <= start_index < end_index <= total_records:
        raise ValueError(f"invalid source range [{start_index}, {end_index})")
    source_indices = list(range(start_index, end_index))
    records = records[start_index:end_index]
    contexts = [row["document"] for row in records]
    output = args.output_dir
    output.mkdir(parents=True, exist_ok=True)
    full, generation_seconds = generate_summaries(
        contexts,
        args,
        output / "generation-checkpoint.json",
    )
    fst = [first_sentence_truncation(answer) for answer in full]
    full_labels, full_judge_seconds = trueteacher_labels(
        contexts,
        full,
        args,
        view="full",
        checkpoint_path=output / "judge-full-checkpoint.json",
    )
    fst_labels, fst_judge_seconds = trueteacher_labels(
        contexts,
        fst,
        args,
        view="fst",
        checkpoint_path=output / "judge-fst-checkpoint.json",
    )
    train_index_set = {int(value) for value in train_indices}
    write_part_view(
        output / "part_full.jsonl",
        records,
        source_indices,
        full,
        full_labels,
        "full",
        train_index_set,
    )
    write_part_view(
        output / "part_fst.jsonl",
        records,
        source_indices,
        fst,
        fst_labels,
        "fst",
        train_index_set,
    )
    summary = {
        "status": "complete",
        "scope": "paper-scale HaluEval preparation; 8000 train pool + 2000 fixed test",
        "paper_probe_split": "7200 train + 800 validation drawn later from the 8000 train pool",
        "dataset_url": DATA_URL,
        "dataset_sha256": dataset_sha256,
        "records": len(records),
        "source_range": [start_index, end_index],
        "full_dataset_train_records": int(len(train_indices)),
        "full_dataset_test_records": int(len(test_indices)),
        "model_id": args.model_id,
        "model_revision": args.model_revision,
        "judge_id": JUDGE_ID,
        "seed": args.seed,
        "generation_seconds": generation_seconds,
        "full_judge_seconds": full_judge_seconds,
        "fst_judge_seconds": fst_judge_seconds,
        "full_label_counts": dict(Counter(full_labels)),
        "fst_label_counts": dict(Counter(fst_labels)),
        "python": sys.version,
        "torch": torch.__version__,
        "gpu": torch.cuda.get_device_name(0),
    }
    write_json(output / "summary.json", summary)
    (output / "SUCCESS").write_text("ok\n", encoding="utf-8")
    print(json.dumps(summary, indent=2), flush=True)


if __name__ == "__main__":
    main()