Skip to main content

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

dropwise-0.2.1.tar.gz (6.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

dropwise-0.2.1-py3-none-any.whl (8.1 kB view details)

Uploaded Python 3

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

Hashes for dropwise-0.2.1.tar.gz
Algorithm Hash digest
SHA256 aa782472d6c056329443717a523771e74f69485f2c3a8e1f4b70448c2767cdf4
MD5 3a603cd21a9ca5a01a915ba62e4e1692
BLAKE2b-256 be737d7d08eeaa6c2749a12a7de7c1da819fd9b65af2f7cad7da8e52f6c6cede

See more details on using hashes here.

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

Hashes for dropwise-0.2.1-py3-none-any.whl
Algorithm Hash digest
SHA256 1e5bcb0a65baa0cbda27cab5dba09f53352f55878afc8a3e8310abf9609758c0
MD5 28996d2691b34d852f43be3ca2ebe004
BLAKE2b-256 15bd476b59c52eae4f4a2a8e80b054c2c7b1807add8b584f929bfd7f300e6b70

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page