src/train_tokenizer.py

"""
Train a BPE tokenizer on the code corpus using the HuggingFace `tokenizers` library.

Produces a 32,000-token vocabulary optimized for source code across
Python, JS/TS, Rust, Go, C/C++, and other languages.
"""

import os
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder

DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
TOKENIZER_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tokenizer")
CORPUS_PATH = os.path.join(DATA_DIR, "corpus.txt")
TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer.json")

VOCAB_SIZE = 32_000


def train_tokenizer():
    os.makedirs(TOKENIZER_DIR, exist_ok=True)

    tokenizer = Tokenizer(BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=True, use_regex=True)
    tokenizer.decoder = ByteLevelDecoder()

    trainer = BpeTrainer(
        vocab_size=VOCAB_SIZE,
        special_tokens=["<pad>", "<bos>", "<eos>", "<unk>"],
        show_progress=True,
        initial_alphabet=ByteLevel.alphabet(),
    )

    print(f"Training BPE tokenizer (vocab_size={VOCAB_SIZE}) on {CORPUS_PATH}...")
    tokenizer.train([CORPUS_PATH], trainer)

    tokenizer.save(TOKENIZER_PATH)
    print(f"Tokenizer saved to {TOKENIZER_PATH}")

    # Print stats
    vocab = tokenizer.get_vocab()
    print(f"Vocabulary size: {len(vocab)}")

    # Test encoding
    test_code = "def hello_world():\n    print('Hello, World!')"
    encoded = tokenizer.encode(test_code)
    print(f"\nTest encoding:")
    print(f"  Input:    {test_code}")
    print(f"  Tokens:   {encoded.tokens[:20]}")
    print(f"  IDs:      {encoded.ids[:20]}")
    print(f"  # tokens: {len(encoded.ids)}")

    return tokenizer


if __name__ == "__main__":
    train_tokenizer()