spb/wpai0002_uqo

1"""SentencePiece tokenizer trained on the project's own training-split text (originals + rewrites), with2reserved special tokens for conditioning (discipline, era, section role) and protected-span placeholders.3The tokenizer is part of the release (``tokenizer.model`` + ``tokenizer_config.json``)."""4from __future__ import annotations56import json7from pathlib import Path89import sentencepiece as spm1011from htr.data.taxonomy import DISCIPLINES, ERAS1213PAD, BOS, EOS, UNK = "<pad>", "<s>", "</s>", "<unk>"14ROLES = ["abstract", "introduction", "background", "methods", "results", "discussion", "conclusion", "related_work", "theory", "data", "limitations", "other", "caption", "chapter", "preface"]15SOURCES = ["pmc", "arxiv", "crossref", "bookshelf", "openstax", "gutenberg", "local", "mock"]161718def control_tokens() -> list[str]:19    toks = [f"<disc:{d}>" for d in DISCIPLINES + ["unknown"]]20    toks += [f"<era:{e[0]}>" for e in ERAS] + ["<era:unknown>", "<era:post2017>"]21    toks += [f"<role:{r}>" for r in ROLES]22    toks += [f"<src:{s}>" for s in SOURCES]23    toks += [f"⟦{k}{i}⟧" for k in "CMU" for i in range(1, 33)]  # protected placeholders as atomic tokens24    toks += ["<sep>", "<meta>", "</meta>"]25    return toks262728def train_tokenizer(texts_path: Path, out_dir: Path, vocab_size: int = 32000, model_type: str = "unigram", input_sentence_size: int = 2_000_000) -> Path:29    out_dir.mkdir(parents=True, exist_ok=True)30    min_vocab = len(control_tokens()) + 256 + 4 + 64  # control + byte fallback + specials + headroom31    if vocab_size < min_vocab:32        raise ValueError(f"vocab_size {vocab_size} too small; need ≥ {min_vocab}")33    prefix = out_dir / "tokenizer"34    spm.SentencePieceTrainer.train(35        input=str(texts_path),36        model_prefix=str(prefix),37        vocab_size=vocab_size,38        model_type=model_type,39        character_coverage=0.9999,40        pad_id=0, bos_id=1, eos_id=2, unk_id=3,41        pad_piece=PAD, bos_piece=BOS, eos_piece=EOS, unk_piece=UNK,42        user_defined_symbols=control_tokens(),43        input_sentence_size=input_sentence_size,44        shuffle_input_sentence=True,45        num_threads=8,46        byte_fallback=True,47        normalization_rule_name="nfkc",48        train_extremely_large_corpus=False,49        hard_vocab_limit=False,  # small corpora (tests, pilots) may not reach the requested size50    )51    cfg = {"type": "sentencepiece", "model_type": model_type, "vocab_size": vocab_size, "pad_id": 0, "bos_id": 1, "eos_id": 2, "unk_id": 3, "control_tokens": control_tokens()}52    (out_dir / "tokenizer_config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False))53    return prefix.with_suffix(".model")545556class HTRTokenizer:57    def __init__(self, model_path: Path | str):58        self.sp = spm.SentencePieceProcessor(model_file=str(model_path))59        self.pad_id, self.bos_id, self.eos_id, self.unk_id = 0, 1, 2, 36061    @property62    def vocab_size(self) -> int:63        return self.sp.get_piece_size()6465    def encode(self, text: str, add_bos: bool = False, add_eos: bool = True, max_len: int | None = None) -> list[int]:66        ids = self.sp.encode(text, out_type=int)67        if max_len is not None:68            ids = ids[: max_len - int(add_bos) - int(add_eos)]69        return ([self.bos_id] if add_bos else []) + ids + ([self.eos_id] if add_eos else [])7071    def decode(self, ids: list[int]) -> str:72        ids = [i for i in ids if i not in (self.pad_id, self.bos_id, self.eos_id)]73        return self.sp.decode(ids)7475    def piece_id(self, piece: str) -> int:76        return self.sp.piece_to_id(piece)77