distilbert-imdb-sentiment
This is distilbert-base-uncased fine-tuned for binary sentiment
classification (positive/negative) on the IMDB movie reviews dataset.
Training details
- Base model:
distilbert-base-uncased - Dataset:
stanfordnlp/imdb, full train split (25,000 examples) - Epochs: 3
- Max sequence length: 256 (padding + truncation)
- Evaluation: full test split (25,000 examples), evaluated after every epoch
Results
| Epoch | Train loss | Test loss | Test accuracy |
|---|---|---|---|
| 1 | 0.302 | 0.283 | 88.5% |
| 2 | 0.157 | 0.366 | 89.2% |
| 3 | 0.068 | 0.385 | 91.2% |
Final test accuracy: 91.2%
Usage
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
repo_id = "Niophy/distilbert-imdb-sentiment"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForSequenceClassification.from_pretrained(repo_id)
model.eval()
id2label = {0: "negative", 1: "positive"}
sentence = "This movie was absolutely fantastic"
inputs = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=256)
with torch.no_grad():
logits = model(**inputs).logits
predicted_id = torch.argmax(logits, dim=-1).item()
print(id2label[predicted_id])
Limitations
Test loss rises after epoch 1 even as accuracy keeps improving, indicating mild overfitting by epoch 3. Training for more epochs without regularization (e.g. weight decay, early stopping) is unlikely to help much beyond this point.