File size: 2,001 Bytes
3738348 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """
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()
|