from transformers import LlamaTokenizer
from datasets import load_dataset

# Initialize blank tokenizer
tokenizer = LlamaTokenizer()

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

def get_training_corpus():
    batch = 1000
    for i in range(0, len(dataset), batch):
        yield dataset[i : i + batch]["text"]

trained_tokenizer = tokenizer.train_new_from_iterator(
    text_iterator=get_training_corpus(),
    vocab_size=32000,
    length=len(dataset),
    show_progress=True,
)

trained_tokenizer.push_to_hub("my_custom_tokenizer")

tokenizer = LlamaTokenizer.from_pretrained("my_custom_tokenizer")