Hugging Face's logo

Spaces:
BeaMorais
/
Seeds
Running

App Files Community
Seeds / src /03_train.py
BeaMorais's picture
BeaMorais
Update src/03_train.py
ce03405 verified
Raw History Blame Contribute Delete
13.3 kB
import json
import os
import random
import subprocess
import sys
import numpy as np
import torch
from datasets import Dataset
from transformers import (
AutoTokenizer,
AutoModelForSeq2SeqLM,
DataCollatorForSeq2Seq,
Seq2SeqTrainer,
Seq2SeqTrainingArguments,
set_seed,
)
# ============================================================
# CONFIGURAÇÕES
# ============================================================
MODEL_NAME = "google/byt5-small"
TRAIN_FILE = "data/train.jsonl"
VALIDATION_FILE = "data/validation.jsonl"
# Pasta diferente para não sobrescrever os testes anteriores
OUTPUT_DIR = "models/seeds-byt5-overfit-v2"
FINAL_MODEL_DIR = "models/seeds-byt5-overfit-v2-final"
# ------------------------------------------------------------
# TESTE CONTROLADO DE MEMORIZAÇÃO
# ------------------------------------------------------------
MAX_TRAIN_EXAMPLES = 20
MAX_VALIDATION_EXAMPLES = 20
# Os exemplos selecionados são curtos, então é possível aumentar
# o limite sem sofrer com truncamento.
MAX_SOURCE_LENGTH = 256
MAX_TARGET_LENGTH = 256
NUM_EPOCHS = 10
LEARNING_RATE = 1e-4
TRAIN_BATCH_SIZE = 1
EVAL_BATCH_SIZE = 1
# Cada exemplo gera uma atualização do otimizador.
GRADIENT_ACCUMULATION_STEPS = 1
SEED = 42
# ============================================================
# FUNÇÕES AUXILIARES
# ============================================================
def carregar_jsonl(caminho):
with open(caminho, "r", encoding="utf-8") as arquivo:
return [json.loads(linha) for linha in arquivo]
def garantir_dataset():
if os.path.exists(TRAIN_FILE) and os.path.exists(VALIDATION_FILE):
return
print("Arquivos de dataset não encontrados.")
print("Executando automaticamente o script 02_prepare_dataset.py...")
resultado = subprocess.run(
[sys.executable, "src/02_prepare_dataset.py"],
capture_output=True,
text=True,
)
print(resultado.stdout)
if resultado.returncode != 0:
print(resultado.stderr)
raise RuntimeError(
"Não foi possível preparar o dataset automaticamente."
)
def preparar_exemplos():
train_data = carregar_jsonl(TRAIN_FILE)
validation_data = carregar_jsonl(VALIDATION_FILE)
print(f"Treino completo: {len(train_data)} exemplos")
print(f"Validação completa: {len(validation_data)} exemplos")
# --------------------------------------------------------
# Selecionado os exemplos mais curtos.
#
# Isso torna o teste de memorização mais limpo:
# o modelo não precisa lidar com frases enormes nem
# truncamento.
# --------------------------------------------------------
train_data = sorted(
train_data,
key=lambda item: len(item["source"]) + len(item["target"])
)[:MAX_TRAIN_EXAMPLES]
validation_data = sorted(
validation_data,
key=lambda item: len(item["source"]) + len(item["target"])
)[:MAX_VALIDATION_EXAMPLES]
print(f"Treino selecionado: {len(train_data)} exemplos")
print(f"Validação selecionada: {len(validation_data)} exemplos")
print("\nExemplos selecionados para treino:")
for indice, exemplo in enumerate(train_data[:5], start=1):
print(
f"{indice}. "
f"source_chars={len(exemplo['source'])}, "
f"target_chars={len(exemplo['target'])}"
)
print(f" EN: {exemplo['source']}")
print(f" MBC: {exemplo['target']}")
print()
return train_data, validation_data
def criar_dataset(exemplos):
return Dataset.from_list(
[
{
"source": exemplo["source"],
"target": exemplo["target"],
}
for exemplo in exemplos
]
)
# ============================================================
# PREPARAÇÃO
# ============================================================
print("=" * 70)
print("SEEDS — TESTE FINAL DE MEMORIZAÇÃO")
print("=" * 70)
garantir_dataset()
set_seed(SEED)
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
train_data, validation_data = preparar_exemplos()
train_dataset = criar_dataset(train_data)
validation_dataset = criar_dataset(validation_data)
# ============================================================
# TOKENIZER
# ============================================================
print("\nCarregando tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME
)
# ============================================================
# MODELO
# ============================================================
print("Carregando modelo...")
model = AutoModelForSeq2SeqLM.from_pretrained(
MODEL_NAME
)
# Evita incompatibilidade entre cache e
# gradient checkpointing.
model.config.use_cache = False
# Evita que o Transformers tente amarrar novamente
# shared.weight e lm_head.weight.
model.config.tie_word_embeddings = False
# ============================================================
# TOKENIZAÇÃO
# ============================================================
def preprocess_function(exemplos):
prompts = [
f"English: {texto}\nMacushi:"
for texto in exemplos["source"]
]
targets = exemplos["target"]
# Entrada
model_inputs = tokenizer(
prompts,
max_length=MAX_SOURCE_LENGTH,
truncation=True,
)
# Saída esperada
labels = tokenizer(
text_target=targets,
max_length=MAX_TARGET_LENGTH,
truncation=True,
)
model_inputs["labels"] = labels["input_ids"]
return model_inputs
print("\nTokenizando dataset de treino...")
tokenized_train = train_dataset.map(
preprocess_function,
batched=True,
remove_columns=train_dataset.column_names,
)
print("Tokenizando dataset de validação...")
tokenized_validation = validation_dataset.map(
preprocess_function,
batched=True,
remove_columns=validation_dataset.column_names,
)
# ============================================================
# VERIFICAÇÃO DOS TAMANHOS
# ============================================================
print("\nComprimentos dos primeiros exemplos tokenizados:")
for indice in range(min(5, len(tokenized_train))):
print(
f"Exemplo {indice + 1}: "
f"source={len(tokenized_train[indice]['input_ids'])}, "
f"target={len(tokenized_train[indice]['labels'])}"
)
# ============================================================
# DATA COLLATOR
# ============================================================
data_collator = DataCollatorForSeq2Seq(
tokenizer=tokenizer,
model=model,
padding=True,
)
# ============================================================
# ARGUMENTOS DE TREINAMENTO
# ============================================================
training_args = Seq2SeqTrainingArguments(
output_dir=OUTPUT_DIR,
# --------------------------------------------------------
# ÉPOCAS
# --------------------------------------------------------
num_train_epochs=NUM_EPOCHS,
# --------------------------------------------------------
# BATCH
# --------------------------------------------------------
per_device_train_batch_size=TRAIN_BATCH_SIZE,
per_device_eval_batch_size=EVAL_BATCH_SIZE,
# --------------------------------------------------------
# GRADIENT ACCUMULATION
# --------------------------------------------------------
# Cada exemplo gera uma atualização.
gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
# --------------------------------------------------------
# LEARNING RATE
# --------------------------------------------------------
learning_rate=LEARNING_RATE,
# Mantem o learning rate constante durante o teste.
lr_scheduler_type="constant",
warmup_steps=0,
# Não queremos regularização adicional nesse teste.
weight_decay=0.0,
# --------------------------------------------------------
# LOGS
# --------------------------------------------------------
logging_steps=1,
# --------------------------------------------------------
# AVALIAÇÃO E CHECKPOINTS
# --------------------------------------------------------
eval_strategy="epoch",
save_strategy="epoch",
save_total_limit=2,
# Não é preciso gerar durante cada avaliação.
# A geração será feita manualmente no final.
predict_with_generate=False,
# --------------------------------------------------------
# MEMÓRIA
# --------------------------------------------------------
gradient_checkpointing=True,
use_cpu=True,
# Não enviar métricas para serviços externos.
report_to="none",
seed=SEED,
)
# ============================================================
# TRAINER
# ============================================================
trainer = Seq2SeqTrainer(
model=model,
args=training_args,
train_dataset=tokenized_train,
eval_dataset=tokenized_validation,
processing_class=tokenizer,
data_collator=data_collator,
)
# ============================================================
# TREINAMENTO
# ============================================================
print("\n" + "=" * 70)
print("INICIANDO TREINAMENTO")
print("=" * 70)
print(f"Exemplos de treino: {len(tokenized_train)}")
print(f"Épocas: {NUM_EPOCHS}")
print(
"Atualizações por época: "
f"aproximadamente {len(tokenized_train)}"
)
print(
"Atualizações totais: "
f"aproximadamente {len(tokenized_train) * NUM_EPOCHS}"
)
print(f"Learning rate: {LEARNING_RATE}")
print(
"Gradient accumulation: "
f"{GRADIENT_ACCUMULATION_STEPS}"
)
train_result = trainer.train()
print("\nTreinamento concluído.")
print("\nMétricas finais de treino:")
print(train_result.metrics)
# ============================================================
# SALVAMENTO
# ============================================================
print("\nSalvando modelo final...")
trainer.save_model(
FINAL_MODEL_DIR
)
tokenizer.save_pretrained(
FINAL_MODEL_DIR
)
print(
f"Modelo salvo em: {FINAL_MODEL_DIR}"
)
# ============================================================
# RECARREGAMENTO
# ============================================================
print("\nRecarregando modelo salvo...")
final_tokenizer = AutoTokenizer.from_pretrained(
FINAL_MODEL_DIR
)
final_model = AutoModelForSeq2SeqLM.from_pretrained(
FINAL_MODEL_DIR
)
final_model.config.use_cache = True
print("Modelo recarregado com sucesso.")
# ============================================================
# GERAÇÃO
# ============================================================
def gerar_traducao(texto, num_beams):
prompt = f"English: {texto}\nMacushi:"
inputs = final_tokenizer(
prompt,
return_tensors="pt",
truncation=True,
max_length=MAX_SOURCE_LENGTH,
)
with torch.no_grad():
output_ids = final_model.generate(
**inputs,
max_length=MAX_TARGET_LENGTH,
num_beams=num_beams,
early_stopping=True,
# Evita repetições muito óbvias.
no_repeat_ngram_size=3,
)
return final_tokenizer.decode(
output_ids[0],
skip_special_tokens=True,
)
# ============================================================
# TESTE DE MEMORIZAÇÃO
# ============================================================
print("\n" + "=" * 70)
print("TESTE DE MEMORIZAÇÃO")
print("=" * 70)
print(
"\nO modelo será testado nos próprios exemplos "
"que foram utilizados durante o treinamento."
)
print(
"Isso NÃO é uma avaliação de generalização."
)
print(
"O objetivo aqui é verificar se o modelo consegue "
"simplesmente aprender os pares apresentados."
)
for indice, exemplo in enumerate(
train_data[:5],
start=1
):
print("\n" + "-" * 70)
print(f"EXEMPLO {indice}")
print(
f"\nEN esperado:\n"
f"{exemplo['source']}"
)
print(
f"\nMBC esperado:\n"
f"{exemplo['target']}"
)
# --------------------------------------------------------
# GREEDY
# --------------------------------------------------------
traducao_gulosa = gerar_traducao(
exemplo["source"],
num_beams=1,
)
# --------------------------------------------------------
# BEAM SEARCH
# --------------------------------------------------------
traducao_beam = gerar_traducao(
exemplo["source"],
num_beams=4,
)
print(
f"\nGerado — greedy:\n"
f"{traducao_gulosa}"
)
print(
f"\nGerado — beam 4:\n"
f"{traducao_beam}"
)
# ============================================================
# FIM
# ============================================================
print("\n" + "=" * 70)
print("TESTE FINALIZADO")
print("=" * 70)
print(
"\nSe o modelo ainda não conseguir reproduzir "
"os próprios exemplos de treinamento, "
"vamos considerar este teste suficiente "
"para abandonar essa abordagem e partir "
"para outra estratégia."
)