commit files to HF hub

ptq_ner.py

@@ -0,0 +1,30 @@
+from functools import partial
+from optimum.intel.openvino import OVQuantizer, OVModelForTokenClassification
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+model_id = "elastic/distilbert-base-uncased-finetuned-conll03-english"
+# model_id = "xlm-roberta-large-finetuned-conll03-english"
+model_id = "dbmdz/bert-large-cased-finetuned-conll03-english"
+model = AutoModelForTokenClassification.from_pretrained(model_id)    
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+# tokenizer.pad_token_id=0
+
+def preprocess_fn(examples, tokenizer):
+    return tokenizer(
+        examples["tokens"], padding="max_length", max_length=128, truncation=True, is_split_into_words=True
+    )
+
+quantizer = OVQuantizer.from_pretrained(model)
+calibration_dataset = quantizer.get_calibration_dataset(
+    "conll2003",
+    preprocess_function=partial(preprocess_fn, tokenizer=tokenizer),
+    num_samples=300,
+    dataset_split="validation",
+    preprocess_batch=True,
+)
+# The directory where the quantized model will be saved
+save_dir = f"{model_id}_ov_int8"
+# Apply static quantization and save the resulting model in the OpenVINO IR format
+quantizer.quantize(calibration_dataset=calibration_dataset, save_directory=save_dir)
+# Load the quantized model
+optimized_model = OVModelForTokenClassification.from_pretrained(save_dir)