src/tokenizer/train_tokenizer.py

"""Entraînement du tokenizer BPE de Cortex-0 — 100 % from scratch.

Aucun tokenizer pré-entraîné n'est utilisé : la lib `tokenizers` ne sert qu'à
*apprendre* les fusions BPE sur notre corpus français.

Choix techniques (détaillés dans docs/architecture.md §1) :

* **BPE** : algorithme simple, déterministe, bien adapté au français.
* **pre_tokenizer = ByteLevel** : découpe au niveau des octets, gère nativement
  les accents, les emojis et tout script.
* **initial_alphabet = 256 octets** : garantit que *chaque* octet possède un
  token, donc que `encode` → `decode` est un **round-trip exact** pour n'importe
  quelle entrée (emoji, japonais, arabe…). Sans cela, les caractères hors corpus
  seraient perdus.
* **byte_fallback = True** : repli sur les octets pour les caractères inconnus.
* **min_frequency** : seuil de fusion (évite les tokens vus une seule fois).
* Les tokens spéciaux sont **réservés en premier** (ids 0..N-1) et dans l'ordre
  canonique, ce qui rend leurs identifiants stables et connus à l'avance.

Exécution locale (petit échantillon — pour les tests) :

    python -m src.tokenizer.train_tokenizer --local-corpus data/sample.txt \\
        --vocab-size 2000 --out-dir /tmp/tok_toy

Exécution complète (Colab) :

    python -m src.tokenizer.train_tokenizer --max-mb 200 --vocab-size 32000
"""

from __future__ import annotations

import argparse
import json
import time
from collections.abc import Iterable, Iterator
from pathlib import Path
from typing import Any

from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers

from src.tokenizer.special_tokens import (
    ALL_SPECIAL_TOKENS,
    SPECIAL_TOKENS_MAP_FILE,
    TOKENIZER_CONFIG_FILE,
    TOKENIZER_FILE,
    VOCAB_SIZE,
)

ROOT = Path(__file__).resolve().parents[2]
DEFAULT_OUT_DIR = ROOT / "tokenizer"


def build_tokenizer(byte_fallback: bool = True) -> Tokenizer:
    """Crée un tokenizer BPE vide, configuré au niveau des octets."""
    tokenizer = Tokenizer(models.BPE(byte_fallback=byte_fallback))
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    return tokenizer


def train_tokenizer(
    corpus: Iterable[str],
    vocab_size: int = VOCAB_SIZE,
    min_frequency: int = 2,
    byte_fallback: bool = True,
    show_progress: bool = True,
    length: int | None = None,
) -> Tokenizer:
    """Entraîne un BPE sur ``corpus``.

    ``vocab_size`` est une **borne supérieure** : sur un corpus pauvre, la taille
    finale peut être plus petite (il n'y a pas assez de fusions possibles).
    """
    tokenizer = build_tokenizer(byte_fallback=byte_fallback)
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=list(ALL_SPECIAL_TOKENS),
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        show_progress=show_progress,
    )
    tokenizer.train_from_iterator(corpus, trainer=trainer, length=length)
    return tokenizer


def assert_special_token_ids(tokenizer: Tokenizer) -> None:
    """Vérifie que les tokens spéciaux occupent bien les ids 0..N-1 attendus."""
    from src.tokenizer.special_tokens import SPECIAL_TOKEN_IDS

    for token, expected_id in SPECIAL_TOKEN_IDS.items():
        actual = tokenizer.token_to_id(token)
        if actual != expected_id:
            raise ValueError(
                f"id de {token} = {actual}, attendu {expected_id}. "
                "L'ordre des tokens spéciaux a changé."
            )


def save_artifacts(tokenizer: Tokenizer, out_dir: Path) -> dict[str, Path]:
    """Écrit tokenizer.json, special_tokens_map.json et tokenizer_config.json."""
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    tok_path = out_dir / TOKENIZER_FILE
    tokenizer.save(str(tok_path))

    special_map = {
        "pad_token": "<|pad|>",
        "bos_token": "<|bos|>",
        "eos_token": "<|eos|>",
        "unk_token": "<|unk|>",
        "additional_special_tokens": [
            t for t in ALL_SPECIAL_TOKENS if t not in ("<|pad|>", "<|bos|>", "<|eos|>", "<|unk|>")
        ],
    }
    map_path = out_dir / SPECIAL_TOKENS_MAP_FILE
    map_path.write_text(json.dumps(special_map, ensure_ascii=False, indent=2), encoding="utf-8")

    config = {
        "tokenizer_class": "Cortex0BPETokenizer",
        "model_type": "bpe",
        "vocab_size": tokenizer.get_vocab_size(),
        "byte_fallback": True,
        "pre_tokenizer": "ByteLevel",
        "add_prefix_space": False,
        "special_tokens": list(ALL_SPECIAL_TOKENS),
        "pad_token": "<|pad|>",
        "bos_token": "<|bos|>",
        "eos_token": "<|eos|>",
        "unk_token": "<|unk|>",
    }
    cfg_path = out_dir / TOKENIZER_CONFIG_FILE
    cfg_path.write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8")

    return {TOKENIZER_FILE: tok_path, SPECIAL_TOKENS_MAP_FILE: map_path, TOKENIZER_CONFIG_FILE: cfg_path}


# ---------------------------------------------------------------------------
# Chargement du corpus
# ---------------------------------------------------------------------------

#: Corpus français publics, par ordre de préférence (nom -> config datasets).
CORPUS_PRESETS: dict[str, dict[str, Any]] = {
    "wikipedia_fr": {"path": "wikimedia/wikipedia", "name": "20231101.fr", "text_field": "text"},
    "oscar_fr": {"path": "oscar-corpus/OSCAR-2301", "name": "fr", "text_field": "text"},
}


def iter_local_corpus(path: Path, max_bytes: int | None = None) -> Iterator[str]:
    """Itère les lignes d'un fichier texte local (corpus de secours / tests)."""
    path = Path(path)
    consumed = 0
    with path.open("r", encoding="utf-8", errors="ignore") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            consumed += len(line.encode("utf-8"))
            yield line
            if max_bytes is not None and consumed >= max_bytes:
                break


def iter_hf_corpus(
    preset: str,
    max_bytes: int,
    text_field: str | None = None,
    streaming: bool = True,
) -> Iterator[str]:
    """Itère un corpus Hugging Face en streaming, plafonné à ``max_bytes`` octets.

    Aucun poids ni modèle pré-entraîné n'est utilisé : on ne lit que du **texte**.
    """
    from datasets import load_dataset  # import paresseux : dépendance lourde

    if preset not in CORPUS_PRESETS:
        raise ValueError(f"preset inconnu : {preset!r}. Choix : {sorted(CORPUS_PRESETS)}")
    spec = CORPUS_PRESETS[preset]
    field = text_field or spec["text_field"]

    dataset = load_dataset(spec["path"], spec["name"], split="train", streaming=streaming)
    consumed = 0
    for example in dataset:
        text = example.get(field)
        if not text:
            continue
        consumed += len(text.encode("utf-8"))
        yield text
        if consumed >= max_bytes:
            break


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _log(msg: str) -> None:
    print(f"[train_tokenizer] {msg}", flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Entraîner le tokenizer BPE de Cortex-0.")
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument("--local-corpus", type=Path, help="Fichier texte local")
    src.add_argument(
        "--preset",
        choices=sorted(CORPUS_PRESETS),
        help="Corpus Hugging Face public (streaming)",
    )
    parser.add_argument("--max-mb", type=float, default=200.0, help="Plafond de corpus en Mo")
    parser.add_argument("--vocab-size", type=int, default=VOCAB_SIZE)
    parser.add_argument("--min-frequency", type=int, default=2)
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR)
    parser.add_argument(
        "--allow-smaller-vocab",
        action="store_true",
        help="Ne pas échouer si le vocabulaire final est < vocab-size (petit corpus).",
    )
    args = parser.parse_args()

    max_bytes = int(args.max_mb * 1024 * 1024)

    if args.local_corpus:
        _log(f"corpus local : {args.local_corpus}")
        corpus: Iterable[str] = iter_local_corpus(args.local_corpus, max_bytes=max_bytes)
    else:
        _log(f"corpus Hugging Face : {args.preset} (streaming, plafond {args.max_mb} Mo)")
        corpus = iter_hf_corpus(args.preset, max_bytes=max_bytes)

    _log("entraînement BPE…")
    start = time.time()
    tokenizer = train_tokenizer(
        corpus,
        vocab_size=args.vocab_size,
        min_frequency=args.min_frequency,
        show_progress=True,
    )
    elapsed = time.time() - start

    final_size = tokenizer.get_vocab_size()
    _log(f"vocab final : {final_size} (cible {args.vocab_size}) en {elapsed:.1f}s")
    if final_size > args.vocab_size:
        raise SystemExit(f"ERREUR : vocab final {final_size} > cible {args.vocab_size}")
    if final_size < args.vocab_size and not args.allow_smaller_vocab:
        raise SystemExit(
            f"ERREUR : vocab final {final_size} < cible {args.vocab_size}. "
            "Corpus trop petit. Relance avec --allow-smaller-vocab ou augmente --max-mb."
        )

    assert_special_token_ids(tokenizer)
    paths = save_artifacts(tokenizer, args.out_dir)
    for name, path in paths.items():
        _log(f"écrit : {path} ({name})")
    _log("terminé.")


if __name__ == "__main__":
    main()