[Dúvida] Não seria melhor nesse caso de avaliação usar um método de regressão para maior acuracia?

Só gostaria de um feedback dessa questão, pois se o modelo errar uma label de 10->9 ou mesmo 10->7 é um erro absoluto de categoria por regressão não seria melhor aplicado neste caso?

import torch
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, root_mean_squared_error
from datasets import load_dataset, DatasetDict, Value
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments
)

# ----- 1. Setup e Dataset (Igual ao anterior) -----
checkpoint_model = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint_model)
model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint_model,
    num_labels=1,
    problem_type="regression"
)

essays_data = load_dataset("csv", data_files="/content/redacoes.csv")
data_split = essays_data["train"].train_test_split(test_size=0.2, shuffle=True, seed=42)

essays_data = DatasetDict({
    "training": data_split["train"],
    "test": data_split["test"]
})

# Tokenização (Mantendo a otimização de padding dinâmico)
def tokenizer_fn(batch):
    return tokenizer(batch["essay"], truncation=True, max_length=384)

tokenized_dataset = essays_data.map(tokenizer_fn, batched=True, remove_columns=["essay"])

# Tratamento de Labels
tokenized_dataset = tokenized_dataset.rename_column("score", "label")
tokenized_dataset = tokenized_dataset.cast_column("label", Value("float32"))
tokenized_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

# Data Collator
data_collator = DataCollatorWithPadding(tokenizer)

# ----- 2. Função de Métricas -----
# O Trainer nos dá as predições e os labels reais.
# Aqui calculamos o quão "longe" a nota predita está da real.
def compute_metrics(eval_pred):
    predictions, labels = eval_pred

    # O Trainer retorna as predições como numpy arrays.
    # Se o modelo retornar logits, às vezes vem dentro de uma tupla, mas em regressão simples geralmente é direto.
    # O .squeeze() garante que o array tenha o shape correto (ex: [8] em vez de [8,1])
    predictions = predictions.squeeze() if isinstance(predictions, np.ndarray) else predictions[0].squeeze()

    mse = mean_squared_error(labels, predictions)
    rmse = root_mean_squared_error(labels, predictions)
    mae = mean_absolute_error(labels, predictions)

    return {"mse": mse, "rmse": rmse, "mae": mae}

# ----- 3. Modelo -----
model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint_model,
    num_labels=1,
    problem_type="regression"
)

# ----- 4. Configuração do Treino (TrainingArguments) -----
training_args = TrainingArguments(
    output_dir="./resultado_trainer",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
    eval_strategy="epoch",       # Avalia ao final de cada época
    save_strategy="epoch",       # Salva o modelo ao final de cada época
    load_best_model_at_end=True, # No final, carrega o melhor checkpoint (menor loss)
    metric_for_best_model="mae", # Usa o MAE para decidir qual é o "melhor" modelo
    greater_is_better=False,     # Para erro (MAE/RMSE), menor é melhor
    logging_dir='./logs',        # Onde salvar os logs
    logging_steps=10,
    fp16=True,                   # [Otimização] Usa precisão mista (treina mais rápido em GPU T4/P100)
    push_to_hub=False,           # Mude para True se quiser subir pro HF automaticamente
)

# ----- 5. Trainer -----
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["training"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics, # Injetamos nossa lógica de avaliação aqui
)

# ----- 6. Executar -----
print("Iniciando treinamento com Trainer API...")
trainer.train()

# ----- 7. Avaliação Final e Salvamento -----
print("\n--- Resultados Finais no Test Set ---")
metrics = trainer.evaluate()
print(metrics)

trainer.save_model("./meu_modelo_final_trainer")
print("Modelo salvo com sucesso.")