wehub-resource-sync
92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
"""
|
|
Train a tokenizer using our own BPE Tokenizer library.
|
|
In the style of GPT-4 tokenizer.
|
|
"""
|
|
import os
|
|
import time
|
|
import argparse
|
|
import torch
|
|
from nanochat.tokenizer import RustBPETokenizer
|
|
from nanochat.common import get_base_dir
|
|
from nanochat.dataset import parquets_iter_batched
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Parse command line arguments
|
|
|
|
parser = argparse.ArgumentParser(description='Train a BPE tokenizer')
|
|
parser.add_argument('--max-chars', type=int, default=2_000_000_000, help='Maximum characters to train on (default: 2B)')
|
|
parser.add_argument('--doc-cap', type=int, default=10_000, help='Maximum characters per document (default: 10,000)')
|
|
parser.add_argument('--vocab-size', type=int, default=32768, help='Vocabulary size (default: 32768 = 2^15)')
|
|
args = parser.parse_args()
|
|
print(f"max_chars: {args.max_chars:,}")
|
|
print(f"doc_cap: {args.doc_cap:,}")
|
|
print(f"vocab_size: {args.vocab_size:,}")
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Text iterator
|
|
|
|
def text_iterator():
|
|
"""
|
|
1) Flatten the batches into a single iterator
|
|
2) Crop every document to args.doc_cap characters
|
|
3) Break when we've seen args.max_chars characters
|
|
"""
|
|
nchars = 0
|
|
for batch in parquets_iter_batched(split="train"):
|
|
for doc in batch:
|
|
doc_text = doc
|
|
if len(doc_text) > args.doc_cap:
|
|
doc_text = doc_text[:args.doc_cap]
|
|
nchars += len(doc_text)
|
|
yield doc_text
|
|
if nchars > args.max_chars:
|
|
return
|
|
text_iter = text_iterator()
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Train the tokenizer
|
|
t0 = time.time()
|
|
tokenizer = RustBPETokenizer.train_from_iterator(text_iter, args.vocab_size)
|
|
t1 = time.time()
|
|
train_time = t1 - t0
|
|
print(f"Training time: {train_time:.2f}s")
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Save the tokenizer to disk
|
|
base_dir = get_base_dir()
|
|
tokenizer_dir = os.path.join(base_dir, "tokenizer")
|
|
tokenizer.save(tokenizer_dir)
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Quick inline sanity check
|
|
test_text = """Hello world! This is a test.
|
|
Numbers: 123, 4567, 89
|
|
Contractions: I'm, you're, it's
|
|
Special chars: @#$%^&*()
|
|
Unicode: 你好世界 🌍"""
|
|
encoded = tokenizer.encode(test_text)
|
|
decoded = tokenizer.decode(encoded)
|
|
assert decoded == test_text
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# One more thing: we wish to cache a mapping from token id to number of bytes of that token
|
|
# for efficient evaluation of bits per byte. Unlike the typical mean loss, this
|
|
# allows us to report a loss that is invariant to the vocab size of the tokenizer.
|
|
# The bits per byte on the validation set is then one of the primary metrics we care about.
|
|
vocab_size = tokenizer.get_vocab_size()
|
|
special_ids = set(tokenizer.encode_special(s) for s in tokenizer.get_special_tokens())
|
|
token_bytes = []
|
|
for token_id in range(vocab_size):
|
|
if token_id in special_ids:
|
|
token_bytes.append(0) # special tokens are not counted
|
|
else:
|
|
# use the raw bytes of the token: decoding to a string first corrupts
|
|
# tokens that are not valid standalone UTF-8 (e.g. the raw bytes >= 0x80)
|
|
num_bytes = len(tokenizer.decode_single_token_bytes(token_id))
|
|
token_bytes.append(num_bytes)
|
|
token_bytes = torch.tensor(token_bytes, dtype=torch.int32, device='cpu')
|
|
token_bytes_path = os.path.join(tokenizer_dir, "token_bytes.pt")
|
|
with open(token_bytes_path, "wb") as f:
|
|
torch.save(token_bytes, f)
|
|
print(f"Saved token_bytes to {token_bytes_path}")
|