File size: 2,710 Bytes
d7aa146 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import json
import numpy as np
from datasets import load_dataset
from sklearn.metrics import accuracy_score, f1_score
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
Trainer,
TrainingArguments,
)
from src.labels import LABELS
# Modelo base en español (BETO)
BASE_MODEL = "dccuchile/bert-base-spanish-wwm-cased"
MODEL_DIR = "out_emotions"
# Mapeos etiqueta <-> id
label2id = {lab: i for i, lab in enumerate(LABELS)}
id2label = {i: lab for i, lab in enumerate(LABELS)}
def compute_metrics(p):
preds = np.argmax(p.predictions, axis=1)
labels = p.label_ids
acc = accuracy_score(labels, preds)
f1 = f1_score(labels, preds, average="weighted")
return {"accuracy": acc, "f1": f1}
def main():
# Carga dataset JSONL
dataset = load_dataset("json", data_files={
"train": "data/train.jsonl",
"validation": "data/valid.jsonl",
"test": "data/test.jsonl",
})
# Tokenizador
tok = AutoTokenizer.from_pretrained(BASE_MODEL)
# Tokenización
def encode(batch):
return tok(batch["text"], truncation=True, padding="max_length", max_length=128)
dataset = dataset.map(encode, batched=True)
# Renombra 'label' -> 'labels' y convierte string -> id (entero)
dataset = dataset.rename_column("label", "labels")
def to_ids(ex):
ex["labels"] = int(label2id[ex["labels"]]) # 'entusiasta' -> 2, etc.
return ex
dataset = dataset.map(to_ids)
# Formato tensores
dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
# Modelo con metadatos de etiquetas
model = AutoModelForSequenceClassification.from_pretrained(
BASE_MODEL,
num_labels=len(LABELS),
id2label=id2label,
label2id=label2id,
)
# Args de entrenamiento (4.56 usa 'eval_strategy'; 'evaluation_strategy' quedará deprecado)
args = TrainingArguments(
output_dir=MODEL_DIR,
eval_strategy="epoch",
save_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=8,
num_train_epochs=5,
weight_decay=0.01,
logging_dir="./logs",
logging_steps=50,
load_best_model_at_end=True,
)
trainer = Trainer(
model=model,
args=args,
train_dataset=dataset["train"],
eval_dataset=dataset["validation"],
tokenizer=tok, # aviso deprecado, pero sigue funcionando
compute_metrics=compute_metrics,
)
trainer.train()
# Guarda mejor modelo + tokenizer en out_emotions/
trainer.save_model(MODEL_DIR)
tok.save_pretrained(MODEL_DIR)
if __name__ == "__main__":
main()
|