evidence/source/audit_corpus.py

"""
Audit the tokenized shards before committing a multi-day run to them.

A corrupted or mis-tokenized corpus does not announce itself. Training proceeds,
the loss falls, and the model is simply worse than it should be, which is
indistinguishable from "small model, hard problem" unless somebody looked at the
tokens first. Every check here is cheap and answers a question that would
otherwise be answered on day four.

Checks:

* **Id range.** Any id at or above the tokenizer's vocab is a corrupt shard or a
  vocab mismatch, and would index out of the embedding table.
* **Compression.** Chars per token on decoded samples. A code tokenizer should
  land near 3.5. Materially below that means the effective corpus is smaller
  than the token count suggests, since the same code costs more tokens.
* **FIM rate.** The fraction of documents carrying the sentinels should match
  `fim_rate` in the index. This is the headline claim of the model, and it is
  applied at prepare time, so if it is wrong it is wrong in the data and no
  amount of training fixes it.
* **Document length.** Mean tokens between EOS. Pathologically short documents
  mean the quality gate or the source is shredding files.
* **Repetition.** The fraction of the sample made up of its single most common
  token. A shard that is 40 percent one token is padding or a broken decode.

Usage:
    python scripts/audit_corpus.py --index data/shards/index.json \\
        --tokenizer tokenizer/code32k.json
"""

import argparse
import json
import os
import sys
from collections import Counter

import numpy as np
from tokenizers import Tokenizer

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data import validate_data_contract


def audit_shard(path, tok, sentinels, eos, sample_tokens, rng):
    arr = np.memmap(path, dtype=np.uint16, mode="r")
    n = arr.shape[0]
    take = min(sample_tokens, n)
    start = int(rng.integers(0, max(1, n - take)))
    sample = np.asarray(arr[start:start + take])

    counts = Counter(sample.tolist())
    top_id, top_n = counts.most_common(1)[0]

    eos_positions = np.flatnonzero(sample == eos)
    if eos_positions.size > 1:
        doc_lens = np.diff(eos_positions)
        mean_doc = float(doc_lens.mean())
        n_docs = int(eos_positions.size - 1)
    else:
        mean_doc, n_docs = float("nan"), 0

    # A document is FIM-transformed if it carries the prefix sentinel. Count per
    # document rather than per token so the rate is comparable to fim_rate.
    fim_docs = 0
    if n_docs:
        pref = sentinels["prefix"]
        for a, b in zip(eos_positions[:-1], eos_positions[1:]):
            if np.any(sample[a:b] == pref):
                fim_docs += 1

    text = tok.decode([int(x) for x in sample[:20000]])
    chars_per_token = len(text) / max(min(20000, sample.size), 1)

    return {
        "tokens": int(n),
        "max_id": int(sample.max()),
        "min_id": int(sample.min()),
        "chars_per_token": chars_per_token,
        "mean_doc_tokens": mean_doc,
        "docs_sampled": n_docs,
        "fim_doc_rate": (fim_docs / n_docs) if n_docs else float("nan"),
        "top_token_share": top_n / sample.size,
        "top_token_id": int(top_id),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--index", default="data/shards/index.json")
    ap.add_argument("--config", default="config/run1.json")
    ap.add_argument("--tokenizer", default="tokenizer/code32k.json")
    ap.add_argument("--sample-tokens", type=int, default=400_000)
    ap.add_argument("--max-shards", type=int, default=8)
    cli = ap.parse_args()

    with open(cli.index) as f:
        index = json.load(f)
    with open(cli.config) as f:
        config = json.load(f)
    validate_data_contract(config, cli.index)
    tok = Tokenizer.from_file(cli.tokenizer)
    vocab = tok.get_vocab_size()
    eos = tok.token_to_id("<|endoftext|>")
    sentinels = {k: tok.token_to_id(f"<|fim_{k}|>") for k in ("prefix", "middle", "suffix")}
    rng = np.random.default_rng(0)
    root = os.path.dirname(os.path.abspath(cli.index))

    print(f"tokenizer vocab {vocab}, index vocab {index['vocab_size']}, "
          f"declared fim_rate {index['fim_rate']}")
    problems = []

    # Compression measured by encoding a fixed reference, which is deterministic
    # and comparable across tokenizers. Decoding shard windows is not: FIM
    # sentinels and partial documents distort the ratio.
    reference = (
        "def binary_search(items: list[int], target: int) -> int:\n"
        "    low, high = 0, len(items) - 1\n"
        "    while low <= high:\n"
        "        mid = (low + high) // 2\n"
        "        if items[mid] == target:\n"
        "            return mid\n"
        "        if items[mid] < target:\n"
        "            low = mid + 1\n"
        "        else:\n"
        "            high = mid - 1\n"
        "    return -1\n"
    )
    ref_tokens = len(tok.encode(reference).ids)
    ref_cpt = len(reference) / max(ref_tokens, 1)
    print(f"reference compression: {len(reference)} chars -> {ref_tokens} tokens "
          f"({ref_cpt:.2f} chars/token)")
    if ref_cpt < 3.0:
        problems.append(f"reference compression {ref_cpt:.2f} chars/token is poor for "
                        f"code. The same code costs more tokens, so the effective "
                        f"corpus is smaller than 5B tokens suggests.")

    if vocab != index["vocab_size"]:
        problems.append(f"vocab mismatch: tokenizer {vocab} vs index {index['vocab_size']}")

    for split, shards in index["splits"].items():
        picks = shards[:: max(1, len(shards) // cli.max_shards)][:cli.max_shards]
        print(f"\n=== {split}: {len(shards)} shards, auditing {len(picks)} ===")
        print(f"{'shard':<20} {'tokens':>12} {'ch/tok':>7} {'doc len':>8} "
              f"{'fim':>6} {'top tok':>8}")
        agg_cpt, agg_fim = [], []
        for entry in picks:
            path = os.path.join(root, entry["path"])
            if not os.path.exists(path):
                problems.append(f"missing shard {entry['path']}")
                continue
            r = audit_shard(path, tok, sentinels, eos, cli.sample_tokens, rng)
            print(f"{entry['path']:<20} {r['tokens']:>12,} {r['chars_per_token']:>7.2f} "
                  f"{r['mean_doc_tokens']:>8.0f} {r['fim_doc_rate']:>6.2f} "
                  f"{r['top_token_share']:>7.1%}")
            agg_cpt.append(r["chars_per_token"])
            if r["docs_sampled"]:
                agg_fim.append(r["fim_doc_rate"])
            if r["max_id"] >= vocab:
                problems.append(f"{entry['path']}: id {r['max_id']} >= vocab {vocab}")
            if r["top_token_share"] > 0.25:
                problems.append(f"{entry['path']}: token {r['top_token_id']} is "
                                f"{r['top_token_share']:.1%} of the sample")

        if agg_cpt:
            cpt = float(np.mean(agg_cpt))
            print(f"\n  mean chars/token on decoded shard windows {cpt:.2f}")

        # Documents shorter than 16 tokens are passed through untransformed by
        # apply_fim, so the realised rate sits a little under the declared one by
        # construction. Only flag a real shortfall, and only when enough
        # documents were sampled for the rate to mean anything.
        docs_seen = sum(1 for _ in agg_fim)
        if agg_fim and docs_seen >= 3:
            fim = float(np.mean(agg_fim))
            declared = index["fim_rate"]
            print(f"  mean FIM document rate {fim:.2f} against declared {declared}")
            if fim < declared - 0.2:
                problems.append(f"{split}: FIM rate {fim:.2f} is far below the "
                                f"declared {declared}, the transform is not being applied")
            elif fim > declared + 0.15:
                problems.append(f"{split}: FIM rate {fim:.2f} exceeds the declared "
                                f"{declared}, which should be impossible")

    print("\n" + "=" * 60)
    if problems:
        print("PROBLEMS FOUND:")
        for p in problems:
            print(f"  - {p}")
    else:
        print("no problems found. the corpus is fit to train on.")
    raise SystemExit(1 if problems else 0)


if __name__ == "__main__":
    main()