byte-sized-brain-distilbert-imdb

DistilBERT movie-review sentiment, FP32 and INT8 ONNX

Two ONNX versions of one DistilBERT model that reads an English movie review and says whether it is positive or negative. They were made by Byte-Sized Brain, a project that measures what quantization does to the size, accuracy, speed and memory of neural networks.

• distilbert_imdb_fp32.onnx is the full-precision model, fine-tuned on 3,000 IMDB reviews.
• distilbert_imdb_int8.onnx is the same model after ONNX Runtime dynamic INT8 quantization, which stores the weights as 8-bit numbers.

Try both side by side in the hosted demo.

Benchmark (500 IMDB test reviews, CPU)

Use it

The model reads at most 128 tokens. Label 0 is negative and label 1 is positive. Both inputs must be int64, which the tokenizer does not guarantee on every platform.

import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

repo = "vardhjain20/byte-sized-brain-distilbert-imdb"
tokenizer = AutoTokenizer.from_pretrained(repo)
session = ort.InferenceSession(hf_hub_download(repo, "distilbert_imdb_int8.onnx"))
enc = tokenizer("A moving, beautifully acted film.", truncation=True, max_length=128,
                return_tensors="np")
feed = {name: enc[name].astype(np.int64) for name in ("input_ids", "attention_mask")}
logits = session.run(None, feed)[0]
print(["negative", "positive"][int(np.argmax(logits))])