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}")