DistilBERT IMDB — MLflow Integration
A DistilBERT model fine-tuned on the IMDB dataset for sentiment analysis, with MLflow experiment tracking.
Video walkthrough: Log with MLflow and Hugging Face Transformers
Model Details
| Detail | Value |
|---|---|
| Base model | distilbert-base-cased |
| Task | Binary sentiment classification |
| Dataset | IMDB (50K movie reviews) |
| License | Apache 2.0 |
Training Configuration
| Parameter | Value |
|---|---|
| Learning rate | 5e-05 |
| Batch size | 8 (train & eval) |
| Optimizer | Adam (β₁=0.9, β₂=0.999, ε=1e-08) |
| LR scheduler | Linear |
| Epochs | 1 |
Usage
from transformers import pipeline
classifier = pipeline("sentiment-analysis", model="juliensimon/distilbert-imdb-mlflow")
classifier("This movie was absolutely fantastic!")
Viewing MLflow Logs
MLflow training logs are included in the repository. To view:
pip install mlflow
mlflow ui
Framework Versions
- Transformers 4.20.1
- PyTorch 1.10.0
- Datasets 2.3.2
- Downloads last month
- 11
Model tree for juliensimon/distilbert-imdb-mlflow
Base model
distilbert/distilbert-base-cased