Categories
Machine Learning

Create Custom Metric Accuracy for Trainer Classification

Here are a simple custom metric to measure accuracy, which you can attach it into Trainer parameter

from transformers import Trainer
from evaluate import load

metric = load('accuracy')

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    accuracy = metric.compute(predictions=preds, references=labels)

    return {'accuracy': accuracy}
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation'],
    data_collator=data_collator,
    compute_metrics=compute_metrics, <----- HERE
    tokenizer=tokenizer
)

trainer.train()

Leave a Reply

Your email address will not be published. Required fields are marked *