pngwn/distilbert-emotion

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, DataCollatorWithPadding
from sklearn.metrics import accuracy_score, f1_score
import numpy as np

dataset = load_dataset("dair-ai/emotion", "split")
tokenizer = AutoTokenizer.from_pretrained("pngwn/distilbert-emotion")
model = AutoModelForSequenceClassification.from_pretrained("pngwn/distilbert-emotion")

def preprocess(examples):
    return tokenizer(examples["text"], truncation=True)

tokenized = dataset.map(preprocess, batched=True)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return {
        "accuracy": accuracy_score(labels, preds),
        "macro_f1": f1_score(labels, preds, average="macro"),
    }

trainer = Trainer(
    model=model,
    eval_dataset=tokenized["test"],
    tokenizer=tokenizer,
    data_collator=DataCollatorWithPadding(tokenizer),
    compute_metrics=compute_metrics,
)
results = trainer.evaluate()
print(f"Test accuracy: {results['eval_accuracy']:.4f}")
print(f"Test macro_f1: {results['eval_macro_f1']:.4f}")

from transformers import pipeline

classifier = pipeline("text-classification", model="pngwn/distilbert-emotion", top_k=None)
classifier("I am so happy today!")