File size: 3,242 Bytes
58d9159 | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | from datasets import load_dataset
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace, ByteLevel
import time # Para medir o tempo
from tokenizers.normalizers import Sequence, NFD, Lowercase, StripAccents, NFC
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
# 1. Carregar o dataset em modo streaming
dataset_stream = load_dataset("bobboyms/subset-Itau-Unibanco-aroeira-1B-tokens", split="train", streaming=True)
print("Dataset carregado em modo streaming:")
print(dataset_stream)
# Nome da coluna que contém o texto
coluna_texto = "text"
# 2. Criar o gerador para o treinamento do tokenizador
# Esta função irá iterar sobre o dataset streaming e fornecer o texto
def get_training_corpus_streaming():
count = 0
start_time = time.time()
print("Iniciando iteração sobre o dataset streaming para o tokenizador...")
for sample in dataset_stream:
# Certifique-se de que a amostra não é None e a coluna existe
if sample and coluna_texto in sample and isinstance(sample[coluna_texto], str):
yield sample[coluna_texto]
count += 1
if count % 10000 == 0: # Log a cada 10000 amostras
elapsed = time.time() - start_time
print(f" Processadas {count} amostras para o tokenizador em {elapsed:.2f} segundos...")
else: # Opcional: Logar amostras inválidas/puladas
print(f"Aviso: Pulando amostra inválida ou sem coluna '{coluna_texto}': {sample}")
end_time = time.time()
print(
f"Iteração completa. Total de {count} amostras fornecidas ao tokenizador em {end_time - start_time:.2f} segundos.")
special_tokens=[
"[UNK]", "<|endoftext|>",
"<|user_start|>", "<|user_end|>",
"<|assistant_start|>", "<|assistant_end|>",
"<|think_start|>", "<|think_end|>",
"<|command_start|>", "<|command_end|>",
]
if __name__ == "__main__":
print("Inicializando o tokenizador BPE...")
# tokenizer.pre_tokenizer = Whitespace()
tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=True)
tokenizer.normalizer = NFC()
tokenizer.decoder = ByteLevelDecoder(add_prefix_space=True)
# Aqui: merges mais agressivos
trainer = BpeTrainer(
vocab_size=36000 + len(special_tokens),
min_frequency=7,
limit_alphabet=1300,
# continuing_subword_prefix="##",
# end_of_word_suffix="</w>", # baixa frequência mínima para 1
show_progress=True, # barra de progresso
special_tokens=special_tokens,
)
print("Iniciando o treinamento do tokenizador a partir do stream...")
start_train_time = time.time()
tokenizer.train_from_iterator(
get_training_corpus_streaming(),
trainer=trainer
)
end_train_time = time.time()
print(f"Treinamento do tokenizador concluído em {end_train_time - start_train_time:.2f} segundos!")
save_path = "tokens-bpe-36k.json"
tokenizer.save("tokens-bpe-36k.json", pretty=True)
print(f"Tokenizador salvo em {save_path}") |