dank/serverAppConfigs
1
0 Fork 0
Code Issues 4 Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(quant): one corpus plumbing module, and the builders out of main()

Browse Source 
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.
This commit is contained in:
Frank committed 2026-10-29 15:26:08 -07:00
1 parent 1e4fa14772
commit 3e5a37a21f
5 files changed
+405 -229
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

scripts/rechunk-kld-corpus.py
+94 -85
View File
@@ -38,12 +38,19 @@ 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
import tempfile
_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:
@@ -69,100 +76,96 @@ def parse_args() -> argparse.Namespace:
return ap.parse_args()
def load_tokenizer(path: str):
try:
from tokenizers import Tokenizer
except ImportError as exc: # the corpus scripts carry their own venv
raise SystemExit(
f"REFUSED: tokenizers is not importable by {sys.executable} — run the "
f"corpus venv's python (on dankserver: /home/dank/venvs/radcorpus/bin/python) "
f"or set FRANK_QWEN38_TOKENIZER: {exc}") from None
if not os.path.isfile(path):
raise SystemExit(f"REFUSED: no tokenizer at {path}")
return Tokenizer.from_file(path)
@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 refuse(msg: str) -> None:
print(f"REFUSED: {msg}", file=sys.stderr)
sys.exit(2) # 2 is the code argparse uses for unusable arguments
def main() -> None:
a = parse_args()
def validate(a: argparse.Namespace) -> None:
"""Everything that makes this run unusable, named before a token is cut."""
if a.max_doc_tokens <= 0:
refuse(f"--max-doc-tokens must be positive (got {a.max_doc_tokens})")
qc.refuse(f"--max-doc-tokens must be positive (got {a.max_doc_tokens})")
if a.score_tail <= 0:
refuse(f"--score-tail must be positive (got {a.score_tail})")
qc.refuse(f"--score-tail must be positive (got {a.score_tail})")
if a.max_doc_tokens <= a.score_tail:
refuse(f"--max-doc-tokens {a.max_doc_tokens} cannot hold a {a.score_tail}-position "
"score tail: every window would be all tail and no context")
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):
refuse(f"no corpus at {a.corpus}")
qc.refuse(f"no corpus at {a.corpus}")
if os.path.realpath(a.out) == os.path.realpath(a.corpus):
refuse(f"--out {a.out} is the corpus itself; move it to trash first — the trash law "
"means the old segmentation stays recoverable")
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:
refuse(f"{a.out} exists and is not empty; pass --force after moving it to trash")
qc.refuse(f"{a.out} exists and is not empty; pass --force after moving it to trash")
tok = load_tokenizer(a.tokenizer_json)
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 = 0
clamped = scored = 0
longest_child = clamped = scored = 0
sources: dict[str, int] = {}
fd, tmp = tempfile.mkstemp(dir=os.path.dirname(os.path.abspath(a.out)) or ".",
prefix=".rechunk-")
try:
with os.fdopen(fd, "w", encoding="utf-8") as out:
with open(a.corpus, encoding="utf-8") as src:
for lineno, line in enumerate(src, 1):
if not line.strip():
continue
try:
rec = json.loads(line)
except json.JSONDecodeError as exc:
refuse(f"{a.corpus}:{lineno} is not JSON: {exc}")
prompt = rec.get("prompt")
if not isinstance(prompt, str):
refuse(f"{a.corpus}:{lineno} has no string 'prompt' field")
ids = tok.encode(prompt, add_special_tokens=False).ids
docs_in += 1
tokens_in += len(ids)
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
score_from = max(len(window) - a.score_tail, 0)
if len(window) < a.score_tail:
clamped += 1
scored += len(window) - score_from
text = tok.decode(window)
if tok.encode(text, add_special_tokens=False).ids != window:
refuse(f"{a.corpus}:{lineno} window at token {start} does not "
"round-trip through text — the reference would score a "
"different token sequence than the corpus holds")
child = {"prompt": text,
"score_from": score_from}
if "source" in rec:
child["source"] = rec["source"]
out.write(json.dumps(child, ensure_ascii=False) + "\n")
docs_out += 1
tokens_out += len(window)
longest_child = max(longest_child, len(window))
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")
os.replace(tmp, a.out)
except BaseException:
if os.path.exists(tmp):
os.unlink(tmp)
raise
os.chmod(a.out, 0o644)
print(json.dumps({
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),
@@ -176,7 +179,13 @@ def main() -> None:
"scored_positions": scored,
"clamped_score_from": clamped,
"source_documents": sources,
}, indent=1))
}
def main() -> None:
a = parse_args()
validate(a)
print(json.dumps(rechunk(a, qc.load_tokenizer(a.tokenizer_json)), indent=1))
if __name__ == "__main__":