This model is a fine-tuned version of
distilbert-base-uncased on the
dair-ai/emotion dataset for 6-class emotion classification.
1from datasets import load_dataset
2from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, DataCollatorWithPadding
3from sklearn.metrics import accuracy_score, f1_score
4import numpy as np
5
6dataset = load_dataset("dair-ai/emotion", "split")
7tokenizer = AutoTokenizer.from_pretrained("pngwn/distilbert-emotion")
8model = AutoModelForSequenceClassification.from_pretrained("pngwn/distilbert-emotion")
9
10def preprocess(examples):
11 return tokenizer(examples["text"], truncation=True)
12
13tokenized = dataset.map(preprocess, batched=True)
14
15def compute_metrics(eval_pred):
16 logits, labels = eval_pred
17 preds = np.argmax(logits, axis=-1)
18 return {
19 "accuracy": accuracy_score(labels, preds),
20 "macro_f1": f1_score(labels, preds, average="macro"),
21 }
22
23trainer = Trainer(
24 model=model,
25 eval_dataset=tokenized["test"],
26 tokenizer=tokenizer,
27 data_collator=DataCollatorWithPadding(tokenizer),
28 compute_metrics=compute_metrics,
29)
30results = trainer.evaluate()
31print(f"Test accuracy: {results['eval_accuracy']:.4f}")
32print(f"Test macro_f1: {results['eval_macro_f1']:.4f}")
1from transformers import pipeline
2
3classifier = pipeline("text-classification", model="pngwn/distilbert-emotion", top_k=None)
4classifier("I am so happy today!")
