Cómo afinar Gemma 4: guía completa con un dataset de emociones humanas

Google acaba de presentar Gemma 4, su familia de modelos abiertos más inteligente hasta la fecha, pensada para un razonamiento sólido y flujos de trabajo con agentes. Los modelos Gemma están diseñados para ser flexibles en distintos entornos, con soporte y herramientas oficiales para desarrollo local, despliegue en la nube y personalización de modelos, lo que los convierte en una gran opción para proyectos de fine-tuning.

En este tutorial, vamos a afinar Gemma 4 E4B-it con un dataset de clasificación de emociones humanas de Hugging Face. Configuraremos un entorno con GPU 3090, cargaremos e inspeccionaremos el dataset, prepararemos y daremos formato a los datos para entrenamiento supervisado, cargaremos el modelo base, haremos una evaluación inicial antes del entrenamiento, afinaremos el modelo y, por último, volveremos a evaluar su rendimiento tras entrenar.

1. Configuración del entorno

Empieza lanzando una nueva instancia de Runpod y asegúrate de que tu cuenta tenga al menos 5 $ de saldo antes de comenzar. Para este tutorial, elige un pod con GPU 3090 y selecciona la última plantilla de PyTorch.

Antes de desplegar, abre los ajustes de la plantilla y realiza algunos cambios. Aumenta tanto el disco del contenedor como el disco del volumen a 40 GB para tener espacio suficiente para el modelo, el dataset, los archivos en caché y los checkpoints de entrenamiento. 

También deberías añadir tu token de Hugging Face como variable de entorno. Puedes generarlo en Settings > Access Tokens de tu cuenta de Hugging Face.

Cuando tengas todo listo, despliega el pod. La instancia puede tardar uno o dos minutos en iniciarse. Cuando esté preparada, abre la interfaz de JupyterLab para empezar a trabajar en el entorno.

Lo primero en JupyterLab es crear un nuevo cuaderno de Python e instalar los paquetes necesarios. Ejecuta este comando en una celda:

%%capture
!pip install -U transformers accelerate datasets trl peft bitsandbytes scikit-learn huggingface_hub

Estos paquetes cubren todo el flujo: cargar el dataset, preparar el modelo, afinarlo y evaluar.

El último paso es iniciar sesión en Hugging Face Hub con tu token guardado. Esto te da acceso al modelo con puerta y facilita subir archivos, crear repositorios y publicar tu modelo afinado más adelante.

import os
from huggingface_hub import login

hf_token = os.environ.get("HF_TOKEN")
if not hf_token:
    raise ValueError("Set HF_TOKEN in the RunPod environment before running this notebook.")

login(token=hf_token)
print("Logged in to Hugging Face.")

2. Cargar y preparar el dataset de emociones

Ahora que el entorno está listo, el siguiente paso es cargar el dataset de emociones desde Hugging Face y preparar subconjuntos más pequeños para entrenamiento y evaluación. 

En este tutorial no usamos el dataset completo. En su lugar, creamos divisiones limitadas de train, validation y test para que el proceso de fine-tuning sea más rápido y fácil de ejecutar en una sola GPU.

from datasets import load_dataset, DatasetDict

TRAIN_LIMIT = 4000    
VALIDATION_LIMIT = 400  
TEST_LIMIT = 400          
EVAL_LIMIT = 400        

raw_dataset = load_dataset("dair-ai/emotion")

def maybe_limit(split, limit):
    split = split.shuffle(seed=42)
    if limit is None:
        return split
    return split.select(range(min(limit, len(split))))

dataset = DatasetDict({
    "train": maybe_limit(raw_dataset["train"], TRAIN_LIMIT),
    "validation": maybe_limit(raw_dataset["validation"], VALIDATION_LIMIT),
    "test": maybe_limit(raw_dataset["test"], TEST_LIMIT),
})

dataset

El dataset final contiene 4.000 ejemplos de entrenamiento, 400 de validación y 400 de prueba.

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 4000
    })
    validation: Dataset({
        features: ['text', 'label'],
        num_rows: 400
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 400
    })
})

A continuación, consultamos los nombres de las etiquetas almacenadas en el dataset. Estas son las clases de emoción que el modelo aprenderá a predecir.

label_names = dataset["train"].features["label"].names
label_names

Vemos que la tarea tiene seis categorías de emociones: sadness, joy, love, anger, fear y surprise.

['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']

También podemos inspeccionar un ejemplo del split de entrenamiento para ver la estructura de los datos.

dataset["train"][0]

Cada ejemplo contiene un texto y una etiqueta numérica. En este caso, la etiqueta 4 corresponde a fear según la lista de etiquetas anterior.

{'text': 'while cycling in the country', 'label': 4}

3. Dar formato a los datos para el fine-tuning de Gemma 4

Antes de poder afinar el modelo, hay que convertir el dataset al formato que Gemma 4 usará durante el entrenamiento. 

En lugar de pasar solo texto y etiquetas, estructuramos cada ejemplo como una breve interacción tipo chat con un mensaje del sistema, un mensaje del usuario y la respuesta esperada del asistente.

El prompt del sistema indica exactamente la tarea a realizar. En este caso, queremos que el modelo actúe como un asistente de clasificación de emociones y devuelva solo una de las seis etiquetas permitidas.

SYSTEM_PROMPT = """You are an emotion classification assistant.
Read the user's text and answer with exactly one label.
Only choose from: sadness, joy, love, anger, fear, surprise.
Return only the label and nothing else."""

En este esquema, el mensaje del usuario contiene el texto a clasificar y el mensaje del asistente contiene la etiqueta correcta. Es el formato típico para entrenamiento supervisado, donde el modelo aprende a generar la respuesta adecuada para cada ejemplo.

def to_prompt_completion(example):
    text = example["text"]
    label = label_names[example["label"]]
    return {
        "prompt": [
            {
                "role": "system",
                "content": SYSTEM_PROMPT,
            },
            {
                "role": "user",
                "content": f"Classify the emotion of this text:\n\n{text}",
            },
        ],
        "completion": [
            {
                "role": "assistant",
                "content": label,
            }
        ],
    }

sft_dataset = dataset.map(to_prompt_completion, remove_columns=dataset["train"].column_names)

Tras aplicar esta función de formateo, las columnas originales text y label se sustituyen por los campos estructurados prompt y completion.

Podemos inspeccionar un ejemplo para confirmar que el dataset se ha formateado correctamente.

sft_dataset["train"][0]

La salida muestra claramente toda la estructura de entrenamiento. El modelo ve la instrucción, lee el texto de entrada y aprende a producir la etiqueta correcta como respuesta.

{'prompt': [{'content': "You are an emotion classification assistant.\nRead the user's text and answer with exactly one label.\nOnly choose from: sadness, joy, love, anger, fear, surprise.\nReturn only the label and nothing else.",
   'role': 'system'},
  {'content': 'Classify the emotion of this text:\n\nwhile cycling in the country',
   'role': 'user'}],
 'completion': [{'content': 'fear', 'role': 'assistant'}]}

4. Cargar Gemma E4B-it con cuantización de 4 bits

Ahora podemos cargar Gemma 4 E4B-it y prepararlo para el fine-tuning. Como es un modelo relativamente grande, lo cargamos con cuantización de 4 bits para reducir el uso de memoria y facilitar su ejecución en una GPU 3090. También usamos bfloat16 como tipo de cómputo para mantener una configuración eficiente.

Empezamos importando las librerías necesarias y definiendo los ajustes principales del modelo.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

MODEL_ID = "google/gemma-4-E4B-it"
MODEL_DTYPE = torch.bfloat16
USE_4BIT = True

Después, activamos algunas optimizaciones de CUDA y cargamos el tokenizer.

if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

processor = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
if processor.pad_token is None:
    processor.pad_token = processor.eos_token

Ahora preparamos los ajustes de cuantización y los argumentos de carga del modelo.

bnb_config = None
model_kwargs = {
    "device_map": "auto",
}
if USE_4BIT:
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=MODEL_DTYPE,
    )
    model_kwargs["quantization_config"] = bnb_config
else:
    model_kwargs["torch_dtype"] = MODEL_DTYPE

Por último, cargamos el modelo y alineamos su configuración con el tokenizer.

base_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **model_kwargs)
base_model.config.use_cache = False
base_model.config.pad_token_id = processor.pad_token_id
base_model.config.bos_token_id = processor.bos_token_id
base_model.config.eos_token_id = processor.eos_token_id
base_model.generation_config.pad_token_id = processor.pad_token_id
base_model.generation_config.bos_token_id = processor.bos_token_id
base_model.generation_config.eos_token_id = processor.eos_token_id

print(f"Base model loaded with 4-bit={USE_4BIT} and dtype={MODEL_DTYPE}.")

Esto carga el modelo base en el dispositivo disponible, desactiva la caché para el entrenamiento y asegura que los IDs de los tokens especiales estén configurados correctamente tanto en la config del modelo como en la de generación.

Base model loaded with 4-bit=True and dtype=torch.bfloat16.

5. Evaluar el modelo base

Antes de afinar, conviene evaluar primero el modelo base para contar con una referencia clara con la que comparar después. 

En esta sección definimos algunas funciones auxiliares que generan predicciones, extraen etiquetas de emoción válidas y ejecutan la evaluación sobre el split de test.

Empezamos creando un patrón sencillo para extraer etiquetas y funciones auxiliares para la predicción.

Estas funciones gestionan todo el flujo de predicción. El modelo recibe la entrada en formato chat, genera una respuesta corta y luego extraemos la etiqueta predicha. Si el modelo devuelve texto adicional, la función intenta recuperar la primera etiqueta válida.

import re

LABEL_PATTERN = re.compile(r"\b(sadness|joy|love|anger|fear|surprise)\b", re.IGNORECASE)

def extract_label(raw_text: str) -> str:
    raw_text = raw_text.strip().lower()
    match = LABEL_PATTERN.search(raw_text)
    if match:
        return match.group(1)

    first_token = raw_text.split()[0].strip(".,!?:;\"'()[]{}") if raw_text.split() else ""
    return first_token

def generate_label(model, processor, user_text, system_prompt, max_new_tokens=4):
    messages = [
        {
            "role": "system",
            "content": system_prompt,
        },
        {
            "role": "user",
            "content": f"Classify the emotion of this text:\n\n{user_text}",
        },
    ]

    device = next(model.parameters()).device
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(device)

    input_len = inputs["input_ids"].shape[-1]

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=processor.pad_token_id,
            eos_token_id=processor.eos_token_id,
        )

    raw_pred = processor.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
    return extract_label(raw_pred)


def predict_emotion(user_text: str, model=None, proc=None) -> str:
    model = model or base_model
    proc = proc or processor
    return generate_label(model, proc, user_text, SYSTEM_PROMPT)

Ahora podemos probar la configuración con un solo ejemplo antes de correr la evaluación completa.

predict_emotion("I feel so happy and excited today!")

La predicción de muestra parece correcta, así que pasamos a evaluar el modelo en el split de test.

'joy'

Este código evalúa el modelo en el split de test y recoge varias salidas útiles. Guarda las etiquetas reales y predichas, registra si cada predicción es correcta y devuelve métricas de resumen, un informe de clasificación y un dataframe con todas las predicciones.

from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score
import pandas as pd
from tqdm.auto import tqdm

VALID_LABELS = set(label_names)
ALL_EVAL_LABELS = label_names + ["INVALID"]

def evaluate_model(model, processor, split="test", limit=EVAL_LIMIT):
    y_true, y_pred, rows = [], [], []
    raw_source = dataset[split]
    if limit is not None:
        raw_source = raw_source.select(range(min(limit, len(raw_source))))

    model.eval()

    for ex in tqdm(raw_source, desc=f"Evaluating {split}", leave=False):
        true_label = label_names[ex["label"]]
        raw_pred_label = generate_label(model, processor, ex["text"], SYSTEM_PROMPT)
        pred_label = raw_pred_label if raw_pred_label in VALID_LABELS else "INVALID"

        y_true.append(true_label)
        y_pred.append(pred_label)
        rows.append({
            "text": ex["text"],
            "true_label": true_label,
            "pred_label": pred_label,
            "raw_pred_label": raw_pred_label,
            "correct": true_label == pred_label,
        })

    metrics = {
        "accuracy": accuracy_score(y_true, y_pred),
        "macro_f1": f1_score(y_true, y_pred, labels=label_names, average="macro", zero_division=0),
        "invalid_predictions": sum(1 for p in y_pred if p == "INVALID"),
        "evaluated_examples": len(y_true),
    }

    report = classification_report(
        y_true,
        y_pred,
        labels=label_names,
        output_dict=True,
        zero_division=0,
    )

    df = pd.DataFrame(rows)
    return metrics, report, df

def confusion_matrix_df(pred_df):
    return pd.DataFrame(
        confusion_matrix(pred_df["true_label"], pred_df["pred_label"], labels=ALL_EVAL_LABELS),
        index=ALL_EVAL_LABELS,
        columns=ALL_EVAL_LABELS,
    )

Ya podemos ejecutar la evaluación de referencia completa sobre el modelo base.

pre_metrics, pre_report, pre_preds = evaluate_model(base_model, processor, "test")
pre_metrics

Estos resultados iniciales muestran que el modelo sin afinar ya rinde razonablemente bien, pero aún hay margen de mejora. 

La accuracy ronda el 58,25%, el F1 macro está en torno a 0,42 y el modelo produjo 33 predicciones no válidas, lo que significa que a veces devolvió algo fuera del conjunto esperado.

{'accuracy': 0.5825,
 'macro_f1': 0.42112912841373906,
 'invalid_predictions': 33,
 'evaluated_examples': 400}

Ahora podemos ver el informe de clasificación completo por categoría de emoción.

pd.DataFrame(pre_report).transpose()

Esto nos da la precisión, el recall, el F1 y el soporte de cada clase. Ayuda a ver qué emociones maneja mejor el modelo y cuáles le cuestan más antes del fine-tuning.

Por último, podemos inspeccionar la matriz de confusión.

confusion_matrix_df(pre_preds)

La matriz de confusión muestra cómo se distribuyen las predicciones entre las distintas clases. 

En el cuaderno aparece como una tabla, lo que facilita detectar qué emociones se confunden entre sí y dónde más falla el modelo base.

6. Afinar Gemma 4 con LoRA

Con los resultados de referencia listos, podemos afinar Gemma 4 usando LoRA. 

LoRA es un método de fine-tuning eficiente en parámetros: no actualizamos todo el modelo, sino que añadimos un número pequeño de pesos adaptadores entrenables sobre el modelo base. Esto hace el entrenamiento mucho más ligero y práctico en una sola GPU.

Empezamos definiendo la configuración de LoRA.

Estos ajustes controlan cómo se acoplan los adaptadores LoRA al modelo. Aquí usamos rango 16, dropout 0,05 y aplicamos LoRA a todas las capas lineales, una configuración habitual para un fine-tuning eficiente.

from peft import LoraConfig

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules="all-linear"
)

Después definimos la configuración de entrenamiento y preparamos el trainer.

Esta configuración busca contener el uso de memoria dando al modelo margen suficiente para aprender del dataset. Entrenamos una época, usamos acumulación de gradientes para simular un batch mayor y activamos opciones como gradient checkpointing y optimización en 8 bits para hacer el entrenamiento más eficiente.

from trl import SFTConfig, SFTTrainer

training_args = SFTConfig(
    output_dir="./gemma4-emotion-lora",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=2,
    learning_rate=1e-4,
    weight_decay=0.01,
    lr_scheduler_type="linear",
    warmup_steps=50,
    num_train_epochs=1,
    logging_steps=50,
    eval_strategy="steps",
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    gradient_checkpointing=True,
    bf16=True,
    fp16=False,
    tf32=True,
    max_length=256,
    packing=False,
    completion_only_loss=True,
    remove_unused_columns=False,
    dataloader_num_workers=2,
    optim="paged_adamw_8bit",
    report_to="none",
)

Ahora comprobamos que el modelo base está listo e inicializamos el trainer. Este paso acopla los adaptadores LoRA al modelo base y prepara el entrenamiento supervisado con nuestros splits de train y validation ya formateados.

from peft import PeftModel

if isinstance(base_model, PeftModel):
    base_model = base_model.unload()
    base_model.config.use_cache = False

trainer = SFTTrainer(
    model=base_model,
    train_dataset=sft_dataset["train"],
    eval_dataset=sft_dataset["validation"],
    peft_config=lora_config,
    args=training_args,
    processing_class=processor,
)

Antes de empezar a entrenar, conviene confirmar que los parámetros LoRA se hayan acoplado correctamente.

Esto cuenta el número de parámetros entrenables y lanza un error si no se añadieron capas LoRA. 

Después, comienza el entrenamiento. 

trainable_params = 0
for param in trainer.model.parameters():
    if param.requires_grad:
        trainable_params += param.numel()

if trainable_params == 0:
    raise RuntimeError("No trainable LoRA parameters were attached. Check target_modules before training.")

print(f"Trainable LoRA parameters: {trainable_params:,}")
train_result = trainer.train()
trainer.model.eval()
trainer.model.config.use_cache = True
train_result

En esta ejecución, el entrenamiento duró casi 9 minutos y tanto la pérdida de entrenamiento como la de validación fueron disminuyendo, buena señal de que el modelo estaba aprendiendo del dataset.

Cuando termine el entrenamiento, guarda el adaptador y el tokenizer en local.

trainer.model.save_pretrained("./gemma4-emotion-lora")
processor.save_pretrained("./gemma4-emotion-lora")

Por último, podemos subir el modelo a Hugging Face Hub.

Esto sube el adaptador afinado y el tokenizer al Hub para que puedas acceder desde cualquier lugar, compartirlo o cargarlo directamente en otro cuaderno o aplicación.

repo_id = "kingabzpro/gemma4-emotion-lora"

# Push adapter + processor to the Hub
trainer.model.push_to_hub(
    repo_id,
    private=False,
)

processor.push_to_hub(
    repo_id,
    private=False,
)

Ahora puedes ver kingabzpro/gemma4-emotion-lora en Hugging Face y probarlo tú mismo. El repositorio incluye los archivos del modelo, instrucciones de uso y los resultados del fine-tuning.

Fuente: kingabzpro/gemma4-emotion-lora · Hugging Face

7. Evaluar el modelo afinado

Con el entrenamiento completo, el paso final es evaluar el modelo afinado en el mismo split de test y comparar los resultados con el modelo base. Así vemos si el fine-tuning con LoRA ha mejorado su capacidad de clasificar emociones con más acierto.

Empezamos cargando el modelo afinado desde el trainer y ejecutando la evaluación.

ft_model = trainer.model
ft_model.eval()
ft_model.config.use_cache = True
post_metrics, post_report, post_preds = evaluate_model(ft_model, processor, "test")
post_metrics

Esto nos da las métricas principales de evaluación del modelo afinado.

Los resultados son claramente mejores que la referencia. Tras el fine-tuning, el modelo alcanza un 77,25% de accuracy y un F1 macro de 0,698. El número de predicciones no válidas también baja de 33 a 20, lo que indica que el modelo afinado no solo es más preciso, sino también más consistente al devolver etiquetas válidas.

{'accuracy': 0.7725,
 'macro_f1': 0.697702361480462,
 'invalid_predictions': 20,
 'evaluated_examples': 400}

Después, podemos ver el informe de clasificación completo.

Esto muestra el informe como un DataFrame de pandas directamente en el cuaderno. Incluye precisión, recall, F1 y soporte por clase, lo que facilita ver qué categorías más han mejorado tras afinar.

pd.DataFrame(post_report).transpose()

Esto también se muestra como tabla en el cuaderno. Te ayuda a ver dónde el modelo afinado sigue fallando y qué categorías de emoción se confunden con más frecuencia.

confusion_matrix_df(post_preds)

Para que la comparación sea más clara, podemos poner las métricas de antes y después del fine-tuning lado a lado.

comparison_df = pd.DataFrame([
    {"stage": "pre_finetuning", **pre_metrics},
    {"stage": "post_finetuning", **post_metrics},
])
comparison_df

Ofrece un resumen rápido de cuánto ha mejorado el modelo tras el entrenamiento.

Nota: si tienes algún problema al ejecutar el código, puedes consultar el cuaderno completo de Jupyter aquí: fine-tune-gemma-4-on-emotions_final.ipynb

Reflexiones finales

El fine-tuning de Gemma 4 es muy sensible a la configuración, especialmente a la estructura del prompt y a los argumentos de entrenamiento. Si el formato del prompt es incorrecto o no usas el template adecuado de forma consistente, el modelo puede entrenarse sin aprender bien la tarea. Lo mismo ocurre con los ajustes de entrenamiento. Suelen ser las principales razones por las que la pérdida no baja, o baja pero apenas mejora la evaluación.

Otra lección importante es max_length. Si lo reduces demasiado, especialmente por debajo de 125, el modelo puede no aprender el patrón en absoluto. Me encontré varios problemas, pero se fueron resolviendo uno a uno, y casi todos tenían que ver con lo mismo: el formateo del prompt y la configuración de entrenamiento.

Para mejorar aún más, un buen siguiente paso sería afinar con el dataset completo y entrenar al menos durante 3 épocas en lugar de solo una. Así el modelo vería más ejemplos y tendría más tiempo para adaptarse, lo que debería traducirse en mejores valores de accuracy y F1.