src/diffusion_lm/dataset_builder.py

"""Build a split-safe, length-bounded instruction-tuning mixture."""
from __future__ import annotations

import hashlib
import json
import random
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Iterable

from .data import apply_neutral_chat_template


DEFAULT_WEIGHTS = {"general": 0.45, "reasoning": 0.18, "math": 0.18, "code": 0.19}
DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."
SYSTEM_PROMPTS = {
    "general": [
        "Respond helpfully, accurately, and clearly.",
        "You are a knowledgeable assistant. Give a direct and useful response.",
        "You are a thoughtful assistant. Follow the user's instructions carefully.",
    ],
    "reasoning": [
        "You are a careful reasoning assistant. Answer accurately and follow the requested format.",
        "Analyze the question carefully and select the best-supported answer.",
    ],
    "math": [
        "You are a careful mathematics tutor. Explain the solution clearly.",
        "Solve mathematical problems accurately and show the relevant reasoning.",
    ],
    "code": [
        "You are an expert programming assistant. Produce correct and readable code.",
        "Help with programming tasks using robust, clear solutions.",
    ],
}


@dataclass
class BuildConfig:
    tokenizer_name: str = "meta-llama/Llama-3.1-8B-Instruct"
    total_examples: int = 500_000
    max_prompt_tokens: int = 256
    max_sequence_tokens: int = 512
    truncate_long_answers: bool = False
    validation_fraction: float = 0.01
    test_fraction: float = 0.01
    seed: int = 42
    weights: dict[str, float] = field(default_factory=lambda: dict(DEFAULT_WEIGHTS))
    cache_dir: str = "data/huggingface"
    exclude_dataset: str | None = None
    allow_excluded_fallback: bool = False
    build_report: dict[str, Any] = field(default_factory=dict, init=False, repr=False)


def normalized_prompt(text: str) -> str:
    """Normalize a prompt for conservative exact-match decontamination."""
    return re.sub(r"\s+", " ", re.sub(r"[^\w\s]", " ", (text or "").casefold())).strip()


def prompt_hash(text: str) -> str:
    return hashlib.sha256(normalized_prompt(text).encode()).hexdigest()


def row_prompt_hashes(row: dict[str, Any]) -> set[str]:
    """Return wrapped and raw-input hashes used for overlap detection."""
    instruction = str(row.get("instruction") or "").strip()
    input_text = str(row.get("input") or "").strip()
    user = "\n\n".join(part for part in (instruction, input_text) if part)
    hashes = {prompt_hash(user)}
    # Code-instruction datasets commonly have an empty separate input field;
    # its hash must not make every such prompt collide with every other one.
    if input_text:
        hashes.add(prompt_hash(input_text))
    return hashes


def format_mc(question: str, choices: list[Any], answer: int | str) -> dict[str, str] | None:
    labels = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[: len(choices)])
    if isinstance(answer, str) and answer in labels:
        index = labels.index(answer)
    else:
        try:
            index = int(answer)
        except (TypeError, ValueError):
            return None
    if not 0 <= index < len(choices):
        return None
    options = "\n".join(f"{label}: {choice}" for label, choice in zip(labels, choices))
    return {
        "instruction": "Answer the following multiple-choice question.",
        "input": f"{question.strip()}\n\n{options}",
        "output": f"{labels[index]}: {choices[index]}",
    }


def format_tulu(row: dict[str, Any]) -> dict[str, str] | None:
    """Convert the final user/assistant exchange; include short prior turns as context."""
    # Tulu is itself a broad mixture. Keep its explicit math/code subsets out of
    # the general bucket because those categories are independently controlled.
    source = str(row.get("source", "")).casefold()
    if "math" in source or "code" in source:
        return None
    messages = [m for m in row.get("messages", []) if (m.get("content") or "").strip()]
    assistant_positions = [i for i, m in enumerate(messages) if m.get("role") == "assistant" and i]
    if not assistant_positions:
        return None
    end = assistant_positions[-1]
    user = next((i for i in range(end - 1, -1, -1) if messages[i].get("role") == "user"), None)
    if user is None:
        return None
    native_system = next((m["content"].strip() for m in messages if m.get("role") == "system"), "")
    history = [m for m in messages[:user] if m.get("role") != "system"]
    history_text = "\n\n".join(f"{m.get('role', 'user').title()}: {m['content'].strip()}" for m in history)
    question = messages[user]["content"].strip()
    return {
        "system": native_system,
        "instruction": "Continue the conversation helpfully." if history_text else "",
        "input": f"Conversation so far:\n{history_text}\n\nUser: {question}" if history_text else question,
        "output": messages[end]["content"].strip(),
    }


def format_hellaswag(row: dict[str, Any]) -> dict[str, str] | None:
    """Make HellaSwag's sentence/event-completion task explicit."""
    choices = row.get("endings", [])
    formatted = format_mc(row.get("ctx", ""), choices, row.get("label"))
    if formatted is None:
        return None
    options = "\n".join(f"{label}: {choice}" for label, choice in zip("ABCDEFGHIJKLMNOPQRSTUVWXYZ", choices))
    return {
        "instruction": "Choose the option that most plausibly continues the described event.",
        "input": (
            f"Beginning of the event:\n{row.get('ctx', '').strip()}\n\n"
            f"What most plausibly happens next?\n{options}"
        ),
        "output": formatted["output"],
    }


def choose_system_prompt(item: dict[str, str], source: str, prompt_key: str) -> str:
    """Preserve native systems; otherwise vary a compatible prompt deterministically."""
    native = (item.get("system") or "").strip()
    if native:
        return native
    category = source.split(":", 1)[0]
    digest = int(prompt_key[:16], 16)
    # Retain the established default for 70% of generated rows. The remaining
    # rows use safe category-specific wording that does not conflict with the
    # expected response style.
    if digest % 10 < 7:
        return DEFAULT_SYSTEM_PROMPT
    variants = SYSTEM_PROMPTS[category]
    return variants[(digest // 10) % len(variants)]


def _targets(total: int, weights: dict[str, float]) -> dict[str, int]:
    if set(weights) != set(DEFAULT_WEIGHTS) or any(v < 0 for v in weights.values()):
        raise ValueError(f"weights must contain exactly {sorted(DEFAULT_WEIGHTS)} with non-negative values")
    scale = sum(weights.values())
    if scale <= 0:
        raise ValueError("at least one mixture weight must be positive")
    targets = {key: int(total * value / scale) for key, value in weights.items()}
    targets[max(targets, key=targets.get)] += total - sum(targets.values())
    return targets


def _take(rows: Iterable[dict[str, Any]], formatter: Callable[[dict[str, Any]], dict[str, str] | None], count: int,
          tokenizer: Any, config: BuildConfig, blocked: set[str], source: str,
          progress: Any | None = None, *, excluded: set[str] | None = None,
          used: set[str] | None = None, stats: dict[str, int] | None = None,
          sample_origin: str = "new_source") -> list[dict[str, Any]]:
    accepted = []
    scanned = 0
    for row in rows:
        scanned += 1
        if progress is not None and scanned % 500 == 0:
            progress.set_postfix_str(f"{source}, scanned={scanned:,}", refresh=True)
        item = formatter(row)
        if not item or not item["output"].strip():
            continue
        user = "\n\n".join(x for x in (item["instruction"].strip(), item["input"].strip()) if x)
        # Reject pathological upstream records before regex normalization or
        # hashing. One observed prompt contains more than a million tokens;
        # running Unicode regexes over it can look like a hung process.
        if len(user) > config.max_prompt_tokens * 50 or len(item["output"]) > config.max_sequence_tokens * 50:
            continue
        key = prompt_hash(user)
        input_text = item["input"].strip()
        # Compare both the raw task text and its instruction-wrapped form: held-out
        # benchmark hashes contain the raw question, while general datasets vary.
        keys = {key}
        if input_text:
            keys.add(prompt_hash(input_text))
        if keys & blocked:
            if stats is not None:
                stats["benchmark_blocked"] = stats.get("benchmark_blocked", 0) + 1
            continue
        if excluded is not None and keys & excluded:
            if stats is not None:
                stats["excluded_overlap"] = stats.get("excluded_overlap", 0) + 1
            continue
        if used is not None and keys & used:
            if stats is not None:
                stats["within_build_duplicate"] = stats.get("within_build_duplicate", 0) + 1
            continue
        resolved_source = str(item.pop("_lad_source", source))
        system = choose_system_prompt(item, resolved_source, key)
        if len(system) > config.max_prompt_tokens * 50:
            continue
        prompt_messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
        # Tokenize to one token beyond the prompt limit. This proves that a row
        # is oversized without creating enormous arrays or triggering the
        # model's max-length warning.
        prompt_ids = apply_neutral_chat_template(
            tokenizer, prompt_messages, tokenize=True, add_generation_prompt=True,
            truncation=True, max_length=config.max_prompt_tokens + 1,
        )
        if len(prompt_ids) > config.max_prompt_tokens:
            continue
        # Do not render and tokenize the prompt a second time. Complete answers
        # reserve one position for EOS. With truncation enabled, an overlong
        # answer may use the whole remaining context and deliberately has no
        # EOS because the source answer did not actually finish.
        answer_capacity = config.max_sequence_tokens - len(prompt_ids)
        complete_answer_limit = answer_capacity - 1
        if complete_answer_limit < 1:
            continue
        answer_ids = tokenizer(
            item["output"].strip(), add_special_tokens=False, truncation=True,
            max_length=answer_capacity + 1,
        )["input_ids"]
        answer_truncated = len(answer_ids) > complete_answer_limit
        if answer_truncated and config.truncate_long_answers:
            # Keep the raw text and stored IDs consistent. Decoding and then
            # re-encoding can occasionally change a boundary token, so shorten
            # until the normalized truncated text fills no more than the
            # remaining context. Truncated answers intentionally omit EOS.
            candidate_ids = list(answer_ids[:answer_capacity])
            while candidate_ids:
                truncated_output = tokenizer.decode(
                    candidate_ids,
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=False,
                ).strip()
                normalized_ids = tokenizer.encode(truncated_output, add_special_tokens=False)
                if truncated_output and len(normalized_ids) <= answer_capacity:
                    item = {**item, "output": truncated_output}
                    answer_ids = normalized_ids
                    break
                candidate_ids.pop()
            else:
                continue
        if answer_truncated and not config.truncate_long_answers:
            continue
        if tokenizer.eos_token_id is None:
            raise ValueError(f"Tokenizer {tokenizer.name_or_path} has no eos_token_id")
        terminal = [] if answer_truncated else [tokenizer.eos_token_id]
        full_ids = list(prompt_ids) + list(answer_ids) + terminal
        # Share one immutable-in-practice Python list while rows are buffered;
        # Dataset.from_list materializes the two required Arrow columns later.
        # Keeping two Python list copies here roughly doubles peak preparation
        # memory and causes severe slowdown from memory pressure on Colab.
        clean_ids = list(full_ids)
        accepted.append({
            **item,
            "system": system,
            "input_ids": clean_ids,
            "labels": clean_ids,
            "category": resolved_source.split(":", 1)[0],
            "source": resolved_source,
            "answer_truncated": answer_truncated,
            "sample_origin": sample_origin,
        })
        if used is not None:
            used.update(keys)
        if progress is not None:
            progress.update(1)
        if len(accepted) >= count:
            break
    return accepted


def _evaluation_hashes(load: Callable[..., Any]) -> set[str]:
    """Hash held-out prompts from every benchmark represented in the mixture."""
    blocked: set[str] = set()
    specs = [
        ("allenai/ai2_arc", "ARC-Easy", ("validation", "test"), "question"),
        ("allenai/ai2_arc", "ARC-Challenge", ("validation", "test"), "question"),
        ("cais/mmlu", "all", ("validation", "test"), "question"),
        ("Rowan/hellaswag", None, ("validation", "test"), "ctx"),
        ("openai/gsm8k", "main", ("test",), "question"),
        ("google-research-datasets/mbpp", "sanitized", ("validation", "test"), "prompt"),
    ]
    for path, name, splits, field_name in specs:
        for split in splits:
            for row in load(path, name, split=split):
                blocked.add(prompt_hash(row.get(field_name, "")))
    return blocked


def _repeat_dataset(dataset: Any, count: int, seed: int, concatenate: Callable) -> Any:
    """Return exactly ``count`` rows, cycling the full unique pool before repeats.

    Arrow datasets are concatenated without expanding repeated rows into a
    giant Python list.  A shuffled partial cycle prevents always favoring the
    beginning of the pool when the target is not an exact multiple.
    """
    if not len(dataset):
        raise RuntimeError("cannot oversample an empty category")
    if len(dataset) >= count:
        return dataset.select(range(count))
    cycles, remainder = divmod(count, len(dataset))
    pieces = [dataset] * cycles
    if remainder:
        pieces.append(dataset.shuffle(seed=seed).select(range(remainder)))
    return concatenate(pieces)


def build_dataset(config: BuildConfig, token: str | None = None):
    """Download, normalize, balance, split, and return a DatasetDict."""
    from datasets import Dataset, DatasetDict, Features, Sequence, Value, concatenate_datasets, load_dataset, load_from_disk
    from transformers import AutoTokenizer
    from tqdm.auto import tqdm

    cache = config.cache_dir
    load = lambda path, name=None, **kw: load_dataset(path, name, cache_dir=cache, token=token, **kw)
    tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name, token=token, cache_dir=cache)
    targets = _targets(config.total_examples, config.weights)
    print(f"Building {config.total_examples:,} examples with category targets: {targets}")
    output_features = Features({
        "system": Value("string"), "instruction": Value("string"), "input": Value("string"), "output": Value("string"),
        # All supported vocabularies fit int32. The collator converts these
        # lists to torch.long, so training behavior is unchanged.
        "input_ids": Sequence(Value("int32")), "labels": Sequence(Value("int32")),
        "category": Value("string"), "source": Value("string"), "answer_truncated": Value("bool"),
        "sample_origin": Value("string"),
    })
    heldout = config.validation_fraction + config.test_fraction
    if not 0 < heldout < 1:
        raise ValueError("validation_fraction + test_fraction must be between zero and one")
    if config.allow_excluded_fallback and not config.exclude_dataset:
        raise ValueError("allow_excluded_fallback requires exclude_dataset")

    excluded_data = None
    excluded_hashes: set[str] = set()
    build_targets = dict(targets)
    reused_heldout = {}
    if config.exclude_dataset:
        exclusion_path = Path(config.exclude_dataset).expanduser()
        print(f"Loading exclusion dataset {config.exclude_dataset}...", flush=True)
        excluded_data = (
            load_from_disk(str(exclusion_path))
            if exclusion_path.is_dir()
            else load_dataset(config.exclude_dataset, cache_dir=cache, token=token)
        )
        required_text = {"instruction", "input"}
        for split_name, rows in excluded_data.items():
            missing = required_text - set(rows.column_names)
            if missing:
                raise ValueError(
                    f"Exclusion dataset split {split_name!r} lacks required columns: {sorted(missing)}"
                )
            for row in rows.select_columns(sorted(required_text)):
                excluded_hashes.update(row_prompt_hashes(row))
        print(f"Loaded {len(excluded_hashes):,} exclusion prompt hashes.", flush=True)

        if config.allow_excluded_fallback:
            required_splits = {"train", "validation", "test"}
            missing_splits = required_splits - set(excluded_data)
            if missing_splits:
                raise ValueError(
                    "Excluded fallback requires train/validation/test splits; missing "
                    f"{sorted(missing_splits)}"
                )
            # Preserve the earlier held-out rows exactly. This keeps validation
            # comparable and prevents previously trained fallback rows from
            # leaking into validation or test after a fresh random split.
            for split_name in ("validation", "test"):
                rows = excluded_data[split_name]
                missing = set(output_features) - {"sample_origin"} - set(rows.column_names)
                if missing:
                    raise ValueError(
                        f"Excluded fallback split {split_name!r} lacks required columns: {sorted(missing)}"
                    )
                category_counts = {name: rows["category"].count(name) for name in DEFAULT_WEIGHTS}
                for category, count in category_counts.items():
                    build_targets[category] -= count
                    if build_targets[category] < 0:
                        raise ValueError(
                            f"Excluded {split_name} contains more {category} rows than the requested "
                            "mixture can accommodate"
                        )
                reused_heldout[split_name] = rows
    print("Loading held-out benchmark prompts for decontamination...", flush=True)
    blocked = _evaluation_hashes(load)
    print(f"Loaded {len(blocked):,} held-out prompt hashes.", flush=True)
    rng = random.Random(config.seed)
    used: set[str] = set()
    if reused_heldout:
        for rows in reused_heldout.values():
            for row in rows.select_columns(["instruction", "input"]):
                used.update(row_prompt_hashes(row))
    filter_stats: dict[str, int] = {}

    def shuffled(path: str, name: str | None = None, split: str = "train"):
        return load(path, name, split=split).shuffle(seed=config.seed)

    print("Loading source datasets (cached sources should open without downloading)...", flush=True)
    sources: dict[str, list[tuple[str, Iterable[dict[str, Any]], Callable, float]]] = {
        "general": [
            ("general:clean-instruct", shuffled("crumb/Clean-Instruct-3M", split="train"), lambda x: {"instruction": x.get("instruction", ""), "input": x.get("input", ""), "output": x.get("output", "")}, .40),
            ("general:tulu-3", shuffled("allenai/tulu-3-sft-mixture"), format_tulu, .35),
            ("general:alpaca-gpt4", shuffled("vicgalle/alpaca-gpt4"), lambda x: {k: x.get(k, "") for k in ("instruction", "input", "output")}, .20),
            ("general:alpaca", shuffled("tatsu-lab/alpaca"), lambda x: {k: x.get(k, "") for k in ("instruction", "input", "output")}, .05),
        ],
        "reasoning": [
            ("reasoning:mmlu", shuffled("cais/mmlu", "all", "auxiliary_train"), lambda x: format_mc(x["question"], x["choices"], x["answer"]), .73),
            ("reasoning:hellaswag", shuffled("Rowan/hellaswag"), format_hellaswag, .25),
            ("reasoning:arc-easy", shuffled("allenai/ai2_arc", "ARC-Easy"), lambda x: format_mc(x["question"], x["choices"]["text"], x["choices"]["label"].index(x["answerKey"]) if x["answerKey"] in x["choices"]["label"] else -1), .02),
        ],
        "math": [
            ("math:orca", shuffled("microsoft/orca-math-word-problems-200k"), lambda x: {"instruction": "Solve the following math problem step by step.", "input": x.get("question", ""), "output": x.get("answer", "")}, .95),
            ("math:gsm8k", shuffled("openai/gsm8k", "main"), lambda x: {"instruction": "Solve the following math problem step by step.", "input": x.get("question", ""), "output": x.get("answer", "")}, .05),
        ],
        "code": [
            ("code:opencoder", shuffled("OpenCoder-LLM/opc-sft-stage2", "educational_instruct"), lambda x: {"instruction": x.get("instruction", ""), "input": "", "output": x.get("output", "")}, .997),
            ("code:mbpp", shuffled("google-research-datasets/mbpp", "sanitized"), lambda x: {"instruction": "Write Python code to solve the following task.", "input": x.get("prompt", ""), "output": x.get("code", "")}, .003),
        ],
    }
    print("Source datasets loaded; formatting and tokenization are starting.", flush=True)
    fallback_train = None
    if config.allow_excluded_fallback:
        fallback_columns = ["system", "instruction", "input", "output", "category", "source"]
        missing = set(fallback_columns) - set(excluded_data["train"].column_names)
        if missing:
            raise ValueError(f"Excluded fallback train split lacks required columns: {sorted(missing)}")
        fallback_train = excluded_data["train"].select_columns(fallback_columns).shuffle(seed=config.seed)

    groups = []
    category_report = {}
    for category, entries in sources.items():
        wanted = build_targets[category]
        progress = tqdm(total=wanted, desc=f"Preparing {category}", unit="rows")
        allocations = [int(wanted * share) for *_, share in entries]
        allocations[0] += wanted - sum(allocations)
        rows = []
        # Keep iterators alive after the preferred-share pass. This lets a
        # larger source contribute additional unused rows when a smaller source
        # cannot meet its allocation, without rescanning or duplicating rows.
        prepared = [(source, data, formatter, iter(data)) for source, data, formatter, _ in entries]
        for (source, _, formatter, iterator), count in zip(prepared, allocations):
            rows.extend(_take(
                iterator, formatter, count, tokenizer, config, blocked, source, progress,
                excluded=excluded_hashes, used=used, stats=filter_stats,
            ))
        if len(rows) < wanted:
            # Exhaust still-unused rows from the largest sources first. Dataset
            # size is only a priority heuristic; exact token filtering remains
            # authoritative.
            for source, data, formatter, iterator in sorted(prepared, key=lambda item: len(item[1]), reverse=True):
                rows.extend(_take(
                    iterator, formatter, wanted - len(rows), tokenizer, config, blocked, source, progress,
                    excluded=excluded_hashes, used=used, stats=filter_stats,
                ))
                if len(rows) >= wanted:
                    break
        novel_count = len(rows)
        if len(rows) < wanted and fallback_train is not None:
            def fallback_rows():
                for old_row in fallback_train:
                    if old_row.get("category") == category:
                        yield old_row

            def format_fallback(old_row):
                return {
                    "system": old_row.get("system", ""),
                    "instruction": old_row.get("instruction", ""),
                    "input": old_row.get("input", ""),
                    "output": old_row.get("output", ""),
                    "_lad_source": old_row.get("source", f"{category}:excluded-fallback"),
                }

            rows.extend(_take(
                fallback_rows(), format_fallback, wanted - len(rows), tokenizer, config,
                blocked, f"{category}:excluded-fallback", progress,
                used=used, stats=filter_stats, sample_origin="excluded_training_fallback",
            ))
        fallback_unique = len(rows) - novel_count
        unique_count = len(rows)
        if not unique_count:
            hint = (
                "; add new source datasets or pass --allow-excluded-fallback"
                if config.exclude_dataset and not config.allow_excluded_fallback else ""
            )
            raise RuntimeError(f"{category}: no rows survived filtering{hint}")
        rng.shuffle(rows)
        group = Dataset.from_list(rows, features=output_features)
        if unique_count < wanted:
            print(
                f"{category}: {unique_count:,} unique eligible rows; oversampling to {wanted:,} "
                f"({wanted / unique_count:.2f}x exposure)"
            )
            progress.update(wanted - unique_count)
        progress.set_postfix_str(f"{unique_count:,} unique", refresh=True)
        progress.close()
        repeated_group = _repeat_dataset(group, wanted, config.seed, concatenate_datasets)
        category_report[category] = {
            "target_rows": wanted,
            "new_unique_rows": novel_count,
            "excluded_training_fallback_unique_rows": fallback_unique,
            "oversampled_rows": wanted - unique_count,
            "final_sample_origins": {
                origin: repeated_group["sample_origin"].count(origin)
                for origin in sorted(set(repeated_group["sample_origin"]))
            },
        }
        groups.append(repeated_group)
    combined = concatenate_datasets(groups).shuffle(seed=config.seed)
    if reused_heldout:
        normalized_heldout = {}
        for split_name, rows in reused_heldout.items():
            if "sample_origin" in rows.column_names:
                rows = rows.remove_columns("sample_origin")
            rows = rows.map(lambda _: {"sample_origin": "excluded_heldout"})
            normalized_heldout[split_name] = rows.select_columns(list(output_features)).cast(output_features)
        result = DatasetDict(
            train=combined,
            validation=normalized_heldout["validation"],
            test=normalized_heldout["test"],
        )
    else:
        first = combined.train_test_split(test_size=heldout, seed=config.seed)
        second = first["test"].train_test_split(test_size=config.test_fraction / heldout, seed=config.seed)
        result = DatasetDict(train=first["train"], validation=second["train"], test=second["test"])
    config.build_report = {
        "exclude_dataset": config.exclude_dataset,
        "allow_excluded_fallback": config.allow_excluded_fallback,
        "excluded_prompt_hashes": len(excluded_hashes),
        "filter_counts": filter_stats,
        "categories": category_report,
        "reused_heldout_rows": {name: len(rows) for name, rows in reused_heldout.items()},
    }
    return result


def write_manifest(dataset: Any, config: BuildConfig, path: str | Path) -> None:
    counts: dict[str, dict[str, int]] = {}
    truncated: dict[str, int] = {}
    origins: dict[str, dict[str, int]] = {}
    for split, rows in dataset.items():
        counts[split] = {name: rows["category"].count(name) for name in DEFAULT_WEIGHTS}
        truncated[split] = sum(rows["answer_truncated"])
        origins[split] = {
            name: rows["sample_origin"].count(name)
            for name in sorted(set(rows["sample_origin"]))
        }
    serialized_config = {key: value for key, value in config.__dict__.items() if key != "build_report"}
    Path(path).write_text(json.dumps({
        "config": serialized_config,
        "rows": counts,
        "truncated_answers": truncated,
        "sample_origins": origins,
        "novelty": config.build_report,
    }, indent=2, sort_keys=True) + "\n")