src/diffusion_lm/tokenizer.py

"""Train byte-level BPE tokenizers and encode packed token files."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Iterable, Literal, Protocol, TypeAlias

import numpy as np
from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, trainers
from tqdm import tqdm


SpecialTokenRole: TypeAlias = Literal["pad", "unk", "bos", "eos", "mask"]
SPECIAL_TOKEN_ROLES: tuple[SpecialTokenRole, ...] = ("pad", "unk", "bos", "eos", "mask")

# These sentinels deliberately include a project-specific random namespace. Common names such as
# ``[MASK]`` occur in documentation and source code often enough to stop a web-scale encoder.
_SPECIAL_TOKEN_NAMESPACE = "mdlm-v1-8f4e2a6c-917b-4d31-b3f7-6a2c8d5e0f19"
SPECIAL_TOKENS = tuple(
    f"<|{_SPECIAL_TOKEN_NAMESPACE}:{role}|>" for role in SPECIAL_TOKEN_ROLES
)
LEGACY_SPECIAL_TOKENS = ("[PAD]", "[UNK]", "[BOS]", "[EOS]", "[MASK]")
# Native spellings honored on pretrained-backbone tokenizers (e.g. Qwen3). Only roles with a
# trained native token map here; the rest are added as namespaced sentinels on free vocab rows.
PRETRAINED_SPECIAL_TOKENS: tuple[str | None, ...] = (None, None, None, "<|endoftext|>", None)


def _role_token_candidates(index: int) -> tuple[str, ...]:
    candidates = [SPECIAL_TOKENS[index], LEGACY_SPECIAL_TOKENS[index]]
    pretrained = PRETRAINED_SPECIAL_TOKENS[index]
    if pretrained is not None:
        candidates.append(pretrained)
    return tuple(candidates)


class _TokenizerLike(Protocol):
    def token_to_id(self, token: str) -> int | None: ...


def _raw_tokenizer(tokenizer: _TokenizerLike) -> _TokenizerLike:
    return getattr(tokenizer, "raw_tokenizer", tokenizer)


class RoleAwareTokenizer:
    """Proxy a Tokenizer while keeping legacy role spellings usable by old callers.

    Encoding is always delegated unchanged, so strings such as ``[MASK]`` remain ordinary text in
    newly trained tokenizers. Only explicit ``token_to_id`` lookups receive alias compatibility.
    """

    def __init__(self, tokenizer: Tokenizer) -> None:
        self.raw_tokenizer = tokenizer

    def token_to_id(self, token: str) -> int | None:
        token_id = self.raw_tokenizer.token_to_id(token)
        if token_id is not None:
            return token_id
        for index in range(len(SPECIAL_TOKEN_ROLES)):
            candidates = _role_token_candidates(index)
            if token in candidates:
                for concrete in candidates:
                    token_id = self.raw_tokenizer.token_to_id(concrete)
                    if token_id is not None:
                        return token_id
        return None

    def __getattr__(self, name: str):
        return getattr(self.raw_tokenizer, name)


def _role_index(role: SpecialTokenRole | str) -> int:
    try:
        return SPECIAL_TOKEN_ROLES.index(role)  # type: ignore[arg-type]
    except ValueError as exc:
        choices = ", ".join(SPECIAL_TOKEN_ROLES)
        raise ValueError(f"unknown special-token role {role!r}; expected one of {choices}") from exc


def special_token_string(tokenizer: _TokenizerLike, role: SpecialTokenRole | str) -> str:
    """Return the concrete token string used for ``role`` by a new or legacy tokenizer."""

    index = _role_index(role)
    raw = _raw_tokenizer(tokenizer)
    for token in _role_token_candidates(index):
        if raw.token_to_id(token) is not None:
            return token
    raise ValueError(f"tokenizer is missing the {role!r} special token")


def special_token_id(tokenizer: _TokenizerLike, role: SpecialTokenRole | str) -> int:
    """Resolve a special-token ID by semantic role, independent of its serialized spelling."""

    token = special_token_string(tokenizer, role)
    token_id = _raw_tokenizer(tokenizer).token_to_id(token)
    if token_id is None:  # Kept defensive for non-Tokenizer protocol implementations.
        raise ValueError(f"tokenizer is missing the {role!r} special token")
    return token_id


def special_token_ids(tokenizer: _TokenizerLike) -> dict[SpecialTokenRole, int]:
    """Return all semantic special-token IDs for a new or legacy tokenizer."""

    return {role: special_token_id(tokenizer, role) for role in SPECIAL_TOKEN_ROLES}


def _build_tokenizer() -> Tokenizer:
    tokenizer = Tokenizer(models.BPE(unk_token=SPECIAL_TOKENS[1]))
    tokenizer.normalizer = normalizers.NFC()
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    return tokenizer


def _build_trainer(
    *,
    vocab_size: int,
    min_frequency: int,
    max_token_length: int,
    extra_special_tokens: tuple[str, ...] = (),
) -> trainers.BpeTrainer:
    special_tokens = list(SPECIAL_TOKENS) + list(extra_special_tokens)
    if len(set(special_tokens)) != len(special_tokens):
        raise ValueError('extra_special_tokens must not duplicate the role tokens or each other')
    if vocab_size <= len(special_tokens) + 256:
        raise ValueError("vocab_size must leave room for the byte alphabet and special tokens")
    if min_frequency <= 0:
        raise ValueError("min_frequency must be positive")
    if max_token_length <= 0:
        raise ValueError("max_token_length must be positive")
    return trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=special_tokens,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        max_token_length=max_token_length,
        show_progress=True,
    )


def _validate_special_token_layout(tokenizer: Tokenizer) -> None:
    for expected_id, role in enumerate(SPECIAL_TOKEN_ROLES):
        actual_id = special_token_id(tokenizer, role)
        if actual_id != expected_id:
            raise ValueError(
                f"expected the {role!r} special token at id {expected_id}, found {actual_id}"
            )


def _save_tokenizer(tokenizer: Tokenizer, output_path: str | Path) -> None:
    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_name(f".{output.name}.tmp")
    temporary.unlink(missing_ok=True)
    tokenizer.save(str(temporary), pretty=True)
    temporary.replace(output)


TokenizerTrainingItem: TypeAlias = str | list[str] | tuple[str, ...]


def train_tokenizer_from_iterator(
    texts: Iterable[TokenizerTrainingItem],
    output_path: str | Path,
    *,
    vocab_size: int = 32_768,
    min_frequency: int = 10,
    max_token_length: int = 64,
    length: int | None = None,
    extra_special_tokens: tuple[str, ...] = (),
) -> RoleAwareTokenizer:
    """Train a collision-resistant byte-level BPE directly from a text iterator.

    Iterator items may be individual strings or batches of strings. Batched iteration avoids
    materializing a large text export and is the intended entry point for Parquet corpora.
    ``extra_special_tokens`` are assigned the ids directly after the role tokens.
    """

    if length is not None and length < 0:
        raise ValueError("length must be non-negative")
    tokenizer = _build_tokenizer()
    trainer = _build_trainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        max_token_length=max_token_length,
        extra_special_tokens=extra_special_tokens,
    )
    tokenizer.train_from_iterator(texts, trainer=trainer, length=length)
    _validate_special_token_layout(tokenizer)
    _save_tokenizer(tokenizer, output_path)
    return RoleAwareTokenizer(tokenizer)


def train_tokenizer(
    input_paths: Iterable[str | Path],
    output_path: str | Path,
    *,
    vocab_size: int = 32_768,
    min_frequency: int = 10,
    max_token_length: int = 64,
) -> RoleAwareTokenizer:
    """Train a byte-level BPE from local text files."""

    paths = [str(Path(path)) for path in input_paths]
    if not paths:
        raise ValueError("at least one input text file is required")
    missing = [path for path in paths if not Path(path).is_file()]
    if missing:
        raise FileNotFoundError(f"missing tokenizer inputs: {missing}")

    tokenizer = _build_tokenizer()
    trainer = _build_trainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        max_token_length=max_token_length,
    )
    tokenizer.train(paths, trainer)
    _validate_special_token_layout(tokenizer)
    _save_tokenizer(tokenizer, output_path)
    return RoleAwareTokenizer(tokenizer)


def load_tokenizer(path: str | Path) -> RoleAwareTokenizer:
    """Load and validate a project, legacy, or pretrained-backbone token layout.

    Project-trained tokenizers place the role tokens at ids 0-4 and keep the strict layout
    check. Pretrained-backbone tokenizers (role tokens appended on free vocab rows, ``pad``
    far from 0) only need every role to resolve to some id.
    """

    tokenizer_path = Path(path)
    if not tokenizer_path.is_file():
        raise FileNotFoundError(f"tokenizer not found: {tokenizer_path}")
    tokenizer = Tokenizer.from_file(str(tokenizer_path))
    if special_token_id(tokenizer, "pad") == 0:
        _validate_special_token_layout(tokenizer)
    else:
        special_token_ids(tokenizer)
    return RoleAwareTokenizer(tokenizer)


def token_metadata_path(token_path: str | Path) -> Path:
    path = Path(token_path)
    return path.with_suffix(path.suffix + ".json")


def encode_files(
    tokenizer_path: str | Path,
    input_paths: Iterable[str | Path],
    output_path: str | Path,
) -> dict[str, object]:
    """Encode newline-delimited documents to a compact, memory-mappable file.

    This compatibility path remains useful for small corpora. Web-scale Parquet preparation lives
    in :mod:`diffusion_lm.corpus` and preserves embedded newlines within each document.
    """

    tokenizer = load_tokenizer(tokenizer_path)
    inputs = [Path(path) for path in input_paths]
    if not inputs:
        raise ValueError("at least one input text file is required")
    missing = [str(path) for path in inputs if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"missing corpus inputs: {missing}")

    vocab_size = tokenizer.get_vocab_size(with_added_tokens=True)
    dtype = np.dtype("uint16" if vocab_size <= np.iinfo(np.uint16).max else "uint32")
    role_ids = special_token_ids(tokenizer)
    eos_id = role_ids["eos"]
    reserved_ids = set(role_ids.values())

    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    token_count = 0
    document_count = 0
    with output.open("wb") as destination:
        for input_path in inputs:
            with input_path.open("r", encoding="utf-8") as source:
                for line_number, line in enumerate(
                    tqdm(source, desc=f"encoding {input_path.name}", unit="docs"), start=1
                ):
                    text = line.rstrip("\r\n")
                    if not text:
                        continue
                    token_ids = tokenizer.encode(text, add_special_tokens=False).ids
                    encountered = reserved_ids.intersection(token_ids)
                    if encountered:
                        raise ValueError(
                            f"{input_path}:{line_number} encodes reserved special-token ids "
                            f"{sorted(encountered)}; remove literal special tokens from the corpus"
                        )
                    token_ids.append(eos_id)
                    np.asarray(token_ids, dtype=dtype).tofile(destination)
                    token_count += len(token_ids)
                    document_count += 1

    tokenizer_bytes = Path(tokenizer_path).read_bytes()
    metadata: dict[str, object] = {
        "format": "mini-diffusion-lm-packed-tokens-v1",
        "dtype": dtype.name,
        "token_count": token_count,
        "document_count": document_count,
        "vocab_size": vocab_size,
        "mask_token_id": role_ids["mask"],
        "eos_token_id": eos_id,
        "special_token_ids": role_ids,
        "tokenizer_sha256": hashlib.sha256(tokenizer_bytes).hexdigest(),
        "source_files": [str(path) for path in inputs],
    }
    metadata_path = token_metadata_path(output)
    with metadata_path.open("w", encoding="utf-8") as handle:
        json.dump(metadata, handle, indent=2)
        handle.write("\n")
    return metadata


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)

    train_parser = subparsers.add_parser("train", help="train a byte-level BPE tokenizer")
    train_parser.add_argument("--input", type=Path, nargs="+", required=True)
    train_parser.add_argument("--output", type=Path, required=True)
    train_parser.add_argument("--vocab-size", type=int, default=32_768)
    train_parser.add_argument("--min-frequency", type=int, default=10)
    train_parser.add_argument("--max-token-length", type=int, default=64)

    encode_parser = subparsers.add_parser("encode", help="encode text into packed tokens")
    encode_parser.add_argument("--tokenizer", type=Path, required=True)
    encode_parser.add_argument("--input", type=Path, nargs="+", required=True)
    encode_parser.add_argument("--output", type=Path, required=True)
    return parser


def main() -> None:
    args = _build_parser().parse_args()
    if args.command == "train":
        tokenizer = train_tokenizer(
            args.input,
            args.output,
            vocab_size=args.vocab_size,
            min_frequency=args.min_frequency,
            max_token_length=args.max_token_length,
        )
        print(f"saved {tokenizer.get_vocab_size():,}-token tokenizer to {args.output}")
    elif args.command == "encode":
        metadata = encode_files(args.tokenizer, args.input, args.output)
        print(
            f"wrote {metadata['token_count']:,} tokens from "
            f"{metadata['document_count']:,} documents to {args.output}"
        )


if __name__ == "__main__":
    main()