serverAppConfigs/scripts/rechunk-kld-corpus.py

#!/usr/bin/env python3
"""Re-chunk a KLD corpus so every document fits the container that scores it.

A KLD reference is recorded by booting one container and prefilling every corpus
document through it. radiance refuses a document the KV pool cannot hold on its
own (core/mem/kv.cpp `can_ever_fit`: a paged group holds `n_blocks *
block_size` tokens), and the number that matters is printed by the boot ledger,
not by a flag:

    W kv group 'kv_attn' holds 7436 tokens but the context limit is 32768, so
    a prompt longer than 7436 tokens is refused (400) ... or lower
    --max-model-len to 7436

The BF16 qwen4exp container at the campaign's shape gets 270.00 MiB of cache --
the expert plane owns both cards -- so its real window is 7,436 tokens however
`--max-model-len` is spelled. The corpus the campaign inherited from the swift15
era is 197 documents of 8,193-11,970 tokens, because make-calib-corpus.py's
`--kld-min-prompt 8192` sets a floor and no ceiling: every one of them is
unscoreable there, and three drained windows have proved it.

This script moves the document boundaries and nothing else. Each record's token
ids -- counted with the same tokenizer the generator used -- are cut into
windows of at most `--max-doc-tokens`, so a record's windows concatenate back to
that record's own token sequence, unchanged: no token dropped, rewritten or
reworded. Each window gets its own `score_from = len - --score-tail`, the same
relation the generator wrote for the parent, so the reference scores 256
positions per window instead of per document: 100,864 measured positions over
the same text against the 50,432 the inherited corpus scored, at a shorter
context. The calibration corpus is untouched, so the quantization itself does
not move.

Usage:
  rechunk-kld-corpus.py --corpus kld.jsonl --out kld.jsonl.new \
      [--tokenizer-json PATH] [--max-doc-tokens 7168] [--score-tail 256] [--force]

--out must not exist and be non-empty without --force: the same
one-artifact-one-name law run-qwen38-bf16.sh applies to a score report. The file
is written through a temporary name and renamed, so an interrupted run leaves no
half corpus behind. A JSON summary of what was written goes to stdout.
"""
import argparse
import dataclasses
import json
import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

import qwen38_common as qc

MISSING = object()  # "the parent record carried no source field"


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--corpus", required=True,
                    help="the kld.jsonl to re-chunk (records: prompt, score_from, source)")
    ap.add_argument("--out", required=True,
                    help="file to write; refuse if it exists and is non-empty without --force")
    ap.add_argument("--tokenizer-json",
                    default=os.environ.get("FRANK_QWEN38_TOKENIZER",
                                           "/home/dank/models/Qwen3.8-Flash-Next/tokenizer.json"),
                    help="the tokenizer the corpus was authored with (the same one the "
                         "manifest in the corpus output dir names)")
    ap.add_argument("--max-doc-tokens", type=int, default=7168,
                    help="the widest document the scoring container can hold; 7,168 is "
                         "7x1024, the largest whole kilopage under the 7,436 the bf16 "
                         "container's ledger reports")
    ap.add_argument("--score-tail", type=int, default=256,
                    help="positions scored at the end of every document, as the generator "
                         "recorded it")
    ap.add_argument("--force", action="store_true",
                    help="overwrite a non-empty --out (after moving the old file to trash)")
    return ap.parse_args()


@dataclasses.dataclass(frozen=True)
class Window:
    """One scoring document cut out of a parent record."""
    prompt: str
    score_from: int
    n_tokens: int
    clamped: bool
    source: object = MISSING

    def record(self) -> dict:
        rec = {"prompt": self.prompt, "score_from": self.score_from}
        if self.source is not MISSING:
            rec["source"] = self.source
        return rec


def validate(a: argparse.Namespace) -> None:
    """Everything that makes this run unusable, named before a token is cut."""
    if a.max_doc_tokens <= 0:
        qc.refuse(f"--max-doc-tokens must be positive (got {a.max_doc_tokens})")
    if a.score_tail <= 0:
        qc.refuse(f"--score-tail must be positive (got {a.score_tail})")
    if a.max_doc_tokens <= a.score_tail:
        qc.refuse(f"--max-doc-tokens {a.max_doc_tokens} cannot hold a "
                  f"{a.score_tail}-position score tail: every window would be all "
                  "tail and no context")
    if not os.path.isfile(a.corpus):
        qc.refuse(f"no corpus at {a.corpus}")
    if os.path.realpath(a.out) == os.path.realpath(a.corpus):
        qc.refuse(f"--out {a.out} is the corpus itself; move it to trash first — the "
                  "trash law means the old segmentation stays recoverable")
    if os.path.exists(a.out) and os.path.getsize(a.out) > 0 and not a.force:
        qc.refuse(f"{a.out} exists and is not empty; pass --force after moving it to trash")


def cut(rec: dict, tok, a: argparse.Namespace,
        where: str) -> tuple[list[Window], int]:
    """Cut one record's token ids into windows, proving each one round-trips.

    The round-trip encode is the whole reason this tool may claim the reference
    scores the corpus's own token sequence: decoding a window to text and
    encoding it back must return the identical ids, or the score is of a
    different document than the one on disk.
    """
    ids = tok.encode(rec["prompt"], add_special_tokens=False).ids
    windows: list[Window] = []
    for start in range(0, max(len(ids), 1), a.max_doc_tokens):
        window = ids[start:start + a.max_doc_tokens]
        if not window:
            continue
        text = tok.decode(window)
        if tok.encode(text, add_special_tokens=False).ids != window:
            qc.refuse(f"{where} window at token {start} does not round-trip through "
                      "text — the reference would score a different token sequence "
                      "than the corpus holds")
        windows.append(Window(prompt=text,
                              score_from=max(len(window) - a.score_tail, 0),
                              n_tokens=len(window),
                              clamped=len(window) < a.score_tail,
                              source=rec["source"] if "source" in rec else MISSING))
    return windows, len(ids)


def rechunk(a: argparse.Namespace, tok) -> dict:
    """Write the re-chunked corpus and return the summary it printed."""
    docs_in = docs_out = tokens_in = tokens_out = 0
    longest_child = clamped = scored = 0
    sources: dict[str, int] = {}
    with qc.atomic_writer(a.out) as out:
        for lineno, rec in qc.iter_jsonl(a.corpus):
            where = f"{a.corpus}:{lineno}"
            if not isinstance(rec.get("prompt"), str):
                qc.refuse(f"{where} has no string 'prompt' field")
            windows, n_parent = cut(rec, tok, a, where)
            docs_in += 1
            tokens_in += n_parent
            for w in windows:
                out.write(json.dumps(w.record(), ensure_ascii=False) + "\n")
                docs_out += 1
                tokens_out += w.n_tokens
                longest_child = max(longest_child, w.n_tokens)
                clamped += 1 if w.clamped else 0
                scored += w.n_tokens - w.score_from
                if "source" in rec:
                    for cls in str(rec["source"]).split("+"):
                        sources[cls] = sources.get(cls, 0) + 1
    if tokens_out != tokens_in:
        raise SystemExit(f"REFUSED: the windows hold {tokens_out} tokens and the corpus "
                         f"has {tokens_in} — re-chunking lost or invented tokens")
    return {
        "corpus": a.corpus,
        "out": a.out,
        "tokenizer": os.path.basename(a.tokenizer_json),
        "documents_in": docs_in,
        "documents_out": docs_out,
        "tokens_in": tokens_in,
        "tokens_out": tokens_out,
        "longest_document": longest_child,
        "max_doc_tokens": a.max_doc_tokens,
        "score_tail": a.score_tail,
        "scored_positions": scored,
        "clamped_score_from": clamped,
        "source_documents": sources,
    }


def main() -> None:
    a = parse_args()
    validate(a)
    print(json.dumps(rechunk(a, qc.load_tokenizer(a.tokenizer_json)), indent=1))


if __name__ == "__main__":
    main()