src/korbyte/train.py

"""Train and package the KorByte-128K tokenizer."""

from __future__ import annotations

import hashlib
import importlib.metadata
import json
import platform
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from tokenizers import Regex, Tokenizer, decoders, pre_tokenizers, processors
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer

from .config import (
    CORE_VOCAB_SIZE,
    DEFAULT_MODEL_MAX_LENGTH,
    TOTAL_VOCAB_SIZE,
    Paths,
)
from .special_tokens import added_special_tokens, special_token_strings

# Keep a leading ASCII space with the following unit, but do not bind arbitrary
# punctuation to a word. Six-digit number chunks avoid needless fragmentation
# of dates, standards, and common Korean financial notation.
PRETOKENIZER_PATTERN = (
    r" ?[\p{L}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|"
    r" ?\p{N}{1,6}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|"
    r"\s+(?!\S)|\s+"
)


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def build_tokenizer() -> Tokenizer:
    """Construct an untrained, exactly reversible byte-level BPE pipeline."""

    tokenizer = Tokenizer(BPE(unk_token=None, byte_fallback=False, fuse_unk=False))
    tokenizer.normalizer = None
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
        [
            pre_tokenizers.Split(Regex(PRETOKENIZER_PATTERN), behavior="isolated", invert=False),
            pre_tokenizers.ByteLevel(
                add_prefix_space=False,
                use_regex=False,
                trim_offsets=True,
            ),
        ]
    )
    tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
    tokenizer.decoder = decoders.ByteLevel()
    return tokenizer


def _compatibility_files(tokenizer: Tokenizer, root: Path) -> None:
    tokens = special_token_strings()
    vocab = tokenizer.get_vocab(with_added_tokens=True)
    additional = list(tokens[4:])
    tokenizer_config: dict[str, Any] = {
        "add_bos_token": False,
        "add_eos_token": False,
        "additional_special_tokens": additional,
        "bos_token": tokens[0],
        "clean_up_tokenization_spaces": False,
        "eos_token": tokens[1],
        "model_max_length": DEFAULT_MODEL_MAX_LENGTH,
        "pad_token": tokens[2],
        "tokenizer_class": "PreTrainedTokenizerFast",
        "unk_token": tokens[3],
    }
    special_tokens_map = {
        "bos_token": tokens[0],
        "eos_token": tokens[1],
        "pad_token": tokens[2],
        "unk_token": tokens[3],
        "additional_special_tokens": additional,
    }
    added_tokens_map = {token: vocab[token] for token in tokens}

    (root / "tokenizer_config.json").write_text(
        json.dumps(tokenizer_config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    (root / "special_tokens_map.json").write_text(
        json.dumps(special_tokens_map, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    (root / "added_tokens.json").write_text(
        json.dumps(added_tokens_map, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )


def train_tokenizer(root: Path, *, min_frequency: int = 3) -> dict[str, Any]:
    """Train from the prepared corpus and write a complete Hub-compatible snapshot."""

    paths = Paths(root)
    if not paths.corpus.is_file():
        raise FileNotFoundError(f"Prepared corpus not found: {paths.corpus}")
    if not paths.corpus_manifest.is_file():
        raise FileNotFoundError(f"Corpus manifest not found: {paths.corpus_manifest}")

    tokenizer = build_tokenizer()
    trainer = BpeTrainer(
        vocab_size=CORE_VOCAB_SIZE,
        min_frequency=min_frequency,
        show_progress=True,
        special_tokens=[],
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        max_token_length=64,
    )
    tokenizer.train([str(paths.corpus)], trainer)
    core_size = tokenizer.get_vocab_size(with_added_tokens=False)
    if core_size != CORE_VOCAB_SIZE:
        raise RuntimeError(f"Expected {CORE_VOCAB_SIZE} core tokens, trained {core_size}")

    added_count = tokenizer.add_special_tokens(added_special_tokens())
    if added_count != TOTAL_VOCAB_SIZE - CORE_VOCAB_SIZE:
        raise RuntimeError(f"Expected 256 added tokens, got {added_count}")
    total_size = tokenizer.get_vocab_size(with_added_tokens=True)
    if total_size != TOTAL_VOCAB_SIZE:
        raise RuntimeError(f"Expected {TOTAL_VOCAB_SIZE} total tokens, got {total_size}")

    tokenizer.save(str(root / "tokenizer.json"), pretty=True)
    tokenizer.model.save(str(root))
    _compatibility_files(tokenizer, root)

    corpus_manifest = json.loads(paths.corpus_manifest.read_text(encoding="utf-8"))
    build_manifest = {
        "schema_version": 1,
        "created_at": datetime.now(UTC).isoformat(),
        "algorithm": "byte-level-bpe",
        "normalizer": None,
        "pretokenizer_pattern": PRETOKENIZER_PATTERN,
        "core_vocab_size": core_size,
        "special_token_count": added_count,
        "total_vocab_size": total_size,
        "min_frequency": min_frequency,
        "max_token_length": 64,
        "corpus_sha256": corpus_manifest["sha256"],
        "tokenizer_sha256": _sha256(root / "tokenizer.json"),
        "python": platform.python_version(),
        "platform": platform.platform(),
        "packages": {
            package: importlib.metadata.version(package)
            for package in ("datasets", "huggingface-hub", "tokenizers", "transformers")
        },
    }
    paths.build_manifest.write_text(
        json.dumps(build_manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return build_manifest