Monte Carlo Dropout-based uncertainty estimation for Transformers
Project description
Dropwise
Dropwise is a lightweight PyTorch/HuggingFace wrapper for performing Monte Carlo Dropout–based uncertainty estimation in Transformers. It enables confidence-aware decision making by revealing how certain a model is about its predictions — all with just a few lines of code.
🚀 Features
- ✅ Enable dropout during inference for Bayesian-like uncertainty estimation
- ✅ Compute predictive entropy, confidence, and per-class standard deviation
- ✅ Modular support for classification, QA, token tagging, and regression
- ✅ Works seamlessly with Hugging Face Transformers and PyTorch
- ✅ Supports batch inference, CPU/GPU, and customizable
num_passes - ✅ Cleanly packaged and extensible for research or production
🤖 Supported Tasks
| Task Type | Example Model |
|---|---|
sequence-classification |
distilbert-base-uncased-finetuned-sst-2-english |
token-classification |
dslim/bert-base-NER |
question-answering |
deepset/bert-base-cased-squad2 |
regression |
roberta-base (with custom regression head) |
⚠️ Your model must contain dropout layers for MC sampling to work (most HF models do).
📦 Installation
pip install dropwise
Or install from source:
git clone https://github.com/aryanator01/dropwise.git
cd dropwise
pip install -e .
🧠 Example Usage (Per Task)
📘 Sequence Classification
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from dropwise.predictor import DropwisePredictor
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
predictor = DropwisePredictor(model, tokenizer, task_type="sequence-classification", num_passes=20)
results = predictor(["The movie was fantastic!"])
print(results[0])
🏷️ Token Classification (NER)
from transformers import AutoModelForTokenClassification, AutoTokenizer
from dropwise.predictor import DropwisePredictor
model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
predictor = DropwisePredictor(model, tokenizer, task_type="token-classification", num_passes=15)
results = predictor(["Hugging Face is based in New York City."])
print(results[0]['token_predictions'])
❓ Question Answering
from transformers import AutoModelForQuestionAnswering, AutoTokenizer
from dropwise.predictor import DropwisePredictor
model = AutoModelForQuestionAnswering.from_pretrained("deepset/bert-base-cased-squad2")
tokenizer = AutoTokenizer.from_pretrained("deepset/bert-base-cased-squad2")
question = "Where is Hugging Face based?"
context = "Hugging Face Inc. is a company based in New York City."
qa_input = f"{question} [SEP] {context}"
predictor = DropwisePredictor(model, tokenizer, task_type="question-answering", num_passes=10)
results = predictor([qa_input])
print(results[0]['answer'])
📈 Regression
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from dropwise.predictor import DropwisePredictor
model = AutoModelForSequenceClassification.from_pretrained("roberta-base", num_labels=1)
tokenizer = AutoTokenizer.from_pretrained("roberta-base")
predictor = DropwisePredictor(model, tokenizer, task_type="regression", num_passes=20)
results = predictor(["The child is very young."])
print(results[0]['predicted_score'], "+/-", results[0]['uncertainty'])
📊 Output Dictionary (per sample)
| Field | Description |
|---|---|
predicted_class |
Index of most probable class (classification) |
predicted_score |
Scalar prediction (regression only) |
confidence |
Highest softmax probability |
entropy |
Predictive entropy (higher = less confident) |
std_dev |
Standard deviation across MC passes |
probs |
Class-wise softmax probabilities |
margin |
Confidence gap between top-2 predictions |
answer |
Predicted answer span (QA only) |
token_predictions |
List of per-token dicts (token classification) |
🧪 Run Tests
python tests/test_predictor.py
Covers all task types with preloaded models.
📂 Folder Structure
dropwise/
├── predictor.py
├── tasks/
│ ├── __init__.py
│ ├── sequence_classification.py
│ ├── token_classification.py
│ ├── question_answering.py
│ └── regression.py
tests/
└── test_predictor.py
📝 License
MIT License
Built with ❤️ for robust, explainable, uncertainty-aware AI systems.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file dropwise-0.2.1.tar.gz.
File metadata
- Download URL: dropwise-0.2.1.tar.gz
- Upload date:
- Size: 6.9 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.12.4
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
aa782472d6c056329443717a523771e74f69485f2c3a8e1f4b70448c2767cdf4
|
|
| MD5 |
3a603cd21a9ca5a01a915ba62e4e1692
|
|
| BLAKE2b-256 |
be737d7d08eeaa6c2749a12a7de7c1da819fd9b65af2f7cad7da8e52f6c6cede
|
File details
Details for the file dropwise-0.2.1-py3-none-any.whl.
File metadata
- Download URL: dropwise-0.2.1-py3-none-any.whl
- Upload date:
- Size: 8.1 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.12.4
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
1e5bcb0a65baa0cbda27cab5dba09f53352f55878afc8a3e8310abf9609758c0
|
|
| MD5 |
28996d2691b34d852f43be3ca2ebe004
|
|
| BLAKE2b-256 |
15bd476b59c52eae4f4a2a8e80b054c2c7b1807add8b584f929bfd7f300e6b70
|