spb/wpai0002_uqo

1#!/usr/bin/env python32"""Train the SentencePiece tokenizer on TRAIN-split text only (originals in placeholder form + rewrites)."""3from __future__ import annotations45import argparse6import sys7from pathlib import Path89ROOT = Path(__file__).resolve().parents[1]10sys.path.insert(0, str(ROOT / "src"))11from htr.models.tokenizer import HTRTokenizer, train_tokenizer  # noqa: E40212from htr.utils.io import read_jsonl, write_json  # noqa: E40213from htr.utils.provenance import run_header  # noqa: E402141516def main() -> None:17    ap = argparse.ArgumentParser()18    ap.add_argument("--pairs", default=None, help="pairs.jsonl (originals + rewrites of the train split)")19    ap.add_argument("--corpus", default=None, help="parsed corpus dir: eligible train-split passages in placeholder form")20    ap.add_argument("--splits", required=True)21    ap.add_argument("--out", required=True)22    ap.add_argument("--vocab", type=int, default=32000)23    ap.add_argument("--model-type", default="unigram")24    ap.add_argument("--max-lines", type=int, default=600000)25    a = ap.parse_args()26    if not (a.pairs or a.corpus):27        ap.error("--pairs or --corpus required")28    train_docs = {s["doc_id"] for s in read_jsonl(a.splits) if s["split"] == "train"}29    out = Path(a.out)30    out.mkdir(parents=True, exist_ok=True)31    txt = out / "tokenizer_train_text.txt"32    n = 033    with open(txt, "w", encoding="utf-8") as f:34        if a.pairs:35            for p in read_jsonl(a.pairs):36                if p["doc_id"] in train_docs:37                    f.write(p["target_text"].replace("\n", " ") + "\n")38                    f.write(p["source_text"].replace("\n", " ") + "\n")39                    n += 240        if a.corpus:41            from htr.parse.protect import protect4243            for p in read_jsonl(Path(a.corpus) / "passages.jsonl"):44                if p["eligible"] and p["doc_id"] in train_docs:45                    f.write(protect(p["text"]).text.replace("\n", " ") + "\n")46                    n += 147                    if n >= a.max_lines:48                        break49    vocab = a.vocab50    if n < 5000:51        vocab = min(vocab, 8000)52    path = train_tokenizer(txt, out, vocab_size=vocab, model_type=a.model_type)53    tok = HTRTokenizer(path)54    write_json(out / "tokenizer_train_summary.json", {**run_header(), "n_lines": n, "vocab_size": tok.vocab_size, "train_docs": len(train_docs)})55    print(f"tokenizer {path} vocab={tok.vocab_size} lines={n}")56    txt.unlink()575859if __name__ == "__main__":60    main()61