A quantized version of
dslim/distilbert-NER optimized for efficient named entity recognition. This model uses Quantization-Aware Training (QAT) with INT8 dynamic activation and INT4 weight quantization, exported to ONNX format for production deployment.
This model is a quantized version of the DistilBERT NER model, fine-tuned on the CoNLL-2003 dataset for named entity recognition. The quantization preserves accuracy while significantly reducing model size and inference latency.
1import onnxruntime as ort
2import numpy as np
3from transformers import AutoTokenizer
4import json
5
6class NEROnnxRunner:
7 def __init__(self, model_dir: str):
8 with open(f"{model_dir}/config.json", "r") as f:
9 self.id2label = {int(k): v for k, v in json.load(f).items()}
10 self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
11 self.session = ort.InferenceSession(
12 f"{model_dir}/model.onnx", providers=["CPUExecutionProvider"]
13 )
14
15 def predict(self, text: str) -> list[dict]:
16 encoding = self.tokenizer(
17 text,
18 return_tensors="np",
19 padding="max_length",
20 truncation=True,
21 max_length=128,
22 return_offsets_mapping=True,
23 )
24
25 ort_inputs = {
26 "input_ids": encoding["input_ids"].astype(np.int64),
27 "attention_mask": encoding["attention_mask"].astype(np.int64),
28 }
29 logits = self.session.run(None, ort_inputs)[0]
30 predictions = np.argmax(logits, axis=2)[0]
31
32 batch_encoding = self.tokenizer(
33 text, truncation=True, max_length=128, return_offsets_mapping=True
34 )
35 offsets = batch_encoding["offset_mapping"]
36
37 entities = []
38 for idx, (start, end) in enumerate(offsets):
39 if start == end:
40 continue
41 label = self.id2label.get(predictions[idx], "O")
42 if label != "O":
43 entities.append({
44 "word": text[start:end],
45 "label": label,
46 "start": start,
47 "end": end,
48 })
49 return entities
50
51runner = NEROnnxRunner("./distilbert-ner-qat-int4")
52entities = runner.predict("Apple Inc. is based in Cupertino, California.")
53# [{'word': 'Apple Inc.', 'label': 'B-ORG'}, {'word': 'Cupertino', 'label': 'B-LOC'}, {'word': 'California', 'label': 'B-LOC'}]
The model was calibrated using 300 samples from the AG News dataset for QAT observer statistics.
