The three corpus tools now share scripts/qwen38_common.py: tokenizer load with the venv guidance, one n_tokens (so a KLD count cannot drift from the corpus it scores), JSONL read/write, an atomic writer, and refuse(). make-calib-corpus.py's 170-line main() with three nested closures became classify / split_pools / fragment_text / build_frags / pack_class / emit_calib / emit_kld / manifest; its magic numbers are named constants. rechunk-kld-corpus.py's main became validate / cut / rechunk around a Window record, so the round-trip proof sits in one function that can be tested without a tokenizer. calib-push.py grew parse_args / load_prompts / push_round, and a malformed corpus now refuses with a line number instead of raising KeyError. Parity, measured on dankserver against the pre-refactor scripts at the same seed: calib.jsonl and kld.jsonl carry identical records in identical order, the manifest is identical, and the re-chunked corpus is byte-identical (md5 e824e67d...). The one byte-level change is JSON escaping — write_jsonl emits literal UTF-8, which is what the re-chunker already wrote; the parsed records do not move.
193 lines
8.6 KiB
Python
Executable File
193 lines
8.6 KiB
Python
Executable File
#!/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()
|