<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 <h1>Table of Contents<span class="tocSkip"></span></h1>
<div class="toc"><ul class="toc-item"><li><span><a href="#Finetuning-Pre-trained-BERT-Model-on-Text-Classification-Task-And-Inferencing-with-ONNX-Runtime" data-toc-modified-id="Finetuning-Pre-trained-BERT-Model-on-Text-Classification-Task-And-Inferencing-with-ONNX-Runtime-1"><span class="toc-item-num">1&nbsp;&nbsp;</span>Finetuning Pre-trained BERT Model on Text Classification Task And Inferencing with ONNX Runtime</a></span><ul class="toc-item"><li><span><a href="#Tokenizer" data-toc-modified-id="Tokenizer-1.1"><span class="toc-item-num">1.1&nbsp;&nbsp;</span>Tokenizer</a></span></li><li><span><a href="#Model-FineTuning" data-toc-modified-id="Model-FineTuning-1.2"><span class="toc-item-num">1.2&nbsp;&nbsp;</span>Model FineTuning</a></span></li><li><span><a href="#Onnx-Runtime" data-toc-modified-id="Onnx-Runtime-1.3"><span class="toc-item-num">1.3&nbsp;&nbsp;</span>Onnx Runtime</a></span></li></ul></li><li><span><a href="#Reference" data-toc-modified-id="Reference-2"><span class="toc-item-num">2&nbsp;&nbsp;</span>Reference</a></span></li></ul></div>
Step1:  Finetuning Pre-trained BERT Model on Text Classification Task And Inferencing with ONNX Runtime
In this article, we'll be going over two main things
Step2:  Tokenizer
We won't be going over the details of the pre-trained tokenizer or model and only load a pre-trained one available from the huggingface model repository.
Step3:  We can feed our tokenizer directly with a pair of sentences.
Step4:  Decoding the tokenized inputs, this model's tokenizer adds some special tokens such as, [SEP], that is used to indicate which token belongs to which segment/pair.
Step5:  The proprocessing step will be task specific, if we happen to be using another dataset, this function needs to be modified accordingly.
Step6:  Model FineTuning
Having preprocessed our raw dataset, for our text classification task, we use AutoModelForSequenceClassification class to load the pre-trained model, the only other argument we need to specify is the number of class/label our text classification task has. Upon instantiating the model for the first time, we'll see some warnings generated, telling us we should fine tune this model on our down stream tasks before using it.
Step7:  We can perform all sorts of hyper parameter tuning on the fine tuning step, here we'll pick some default parameters for illustration purposes.
Step8:  The next couple of code chunks performs batch inferencing on our dataset, and reports standard binary classification evaluation metrics.
Step9:  Onnx Runtime
This section walks through the process of serializing our Pytorch model into ONNX format, and using ONNX runtime for inferencing.
Exporting the model can be done via the torch.onnx.export function, which requires a sample input.
Step10:  For our neural network computational graph, we expect a single input with a dynamic batch size and sequence length, as well as a single output, which a dynamic batch size. We need to specify dynamic_axes, so during inferencing, onnx won't be limited to the sample input size we've provided when exporting the model into onnx format.
Step11:  We'll also experiment dynamic quantization on our onnx model. Quantization here perfers to converting our weights or activations in our models from a floating point representation to an integer reresentation. Based on the documentation, this is the simplest form of applying quantization, where our weights are quantized ahead of time but the activations are dynamically quantized during inference time.
Given an already trained model, quantization provides a quick way for us to make a tradeoff between model performance versus model size and latency without the need of re-training our models with different parameter settings.
Step12:  One of the most important parameters to tune when using onnxruntime is intra_op_num_threads.
Step13:  We test the inferencing API by feeding the original PyTorch model and ONNX model with a sample input.
Step16:  The next few section runs the benchmark on different batch sizes and inferencing options.
<ASSISTANT_TASK:>
Python Code:
# code for loading the format for the notebook
import os
# path : store the current path to convert back to it later
path = os.getcwd()
os.chdir(os.path.join('..', '..', 'notebook_format'))
from formats import load_style
load_style(css_style='custom2.css', plot_style=False)
os.chdir(path)
# 1. magic for inline plot
# 2. magic to print version
# 3. magic so that the notebook will reload external python modules
# 4. magic to enable retina (high resolution) plots
# https://gist.github.com/minrk/3301035
%matplotlib inline
%load_ext watermark
%load_ext autoreload
%autoreload 2
import os
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%watermark -a 'Ethen' -d -t -v -p datasets,transformers,torch,tokenizers,numpy,pandas,matplotlib,onnxruntime
Explanation: <h1>Table of Contents<span class="tocSkip"></span></h1>
<div class="toc"><ul class="toc-item"><li><span><a href="#Finetuning-Pre-trained-BERT-Model-on-Text-Classification-Task-And-Inferencing-with-ONNX-Runtime" data-toc-modified-id="Finetuning-Pre-trained-BERT-Model-on-Text-Classification-Task-And-Inferencing-with-ONNX-Runtime-1"><span class="toc-item-num">1&nbsp;&nbsp;</span>Finetuning Pre-trained BERT Model on Text Classification Task And Inferencing with ONNX Runtime</a></span><ul class="toc-item"><li><span><a href="#Tokenizer" data-toc-modified-id="Tokenizer-1.1"><span class="toc-item-num">1.1&nbsp;&nbsp;</span>Tokenizer</a></span></li><li><span><a href="#Model-FineTuning" data-toc-modified-id="Model-FineTuning-1.2"><span class="toc-item-num">1.2&nbsp;&nbsp;</span>Model FineTuning</a></span></li><li><span><a href="#Onnx-Runtime" data-toc-modified-id="Onnx-Runtime-1.3"><span class="toc-item-num">1.3&nbsp;&nbsp;</span>Onnx Runtime</a></span></li></ul></li><li><span><a href="#Reference" data-toc-modified-id="Reference-2"><span class="toc-item-num">2&nbsp;&nbsp;</span>Reference</a></span></li></ul></div>
End of explanation
from datasets import load_dataset, DatasetDict, Dataset
dataset_dict = load_dataset("quora")
dataset_dict
dataset_dict['train'][0]
test_size = 0.1
val_size = 0.1
dataset_dict_test = dataset_dict['train'].train_test_split(test_size=test_size)
dataset_dict_train_val = dataset_dict_test['train'].train_test_split(test_size=val_size)
dataset_dict = DatasetDict({
    "train": dataset_dict_train_val["train"],
    "val": dataset_dict_train_val["test"],
    "test": dataset_dict_test["test"]
})
dataset_dict
Explanation: Finetuning Pre-trained BERT Model on Text Classification Task And Inferencing with ONNX Runtime
In this article, we'll be going over two main things:
Process of finetuning a pre-trained BERT model towards a text classification task, more specificially, the Quora Question Pairs challenge.
Process of converting our model into ONNX format, and perform inferencing benchmark with ONNX runtime.
End of explanation
from transformers import AutoTokenizer
# https://huggingface.co/transformers/model_doc/distilbert.html
pretrained_model_name_or_path = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
tokenizer
Explanation: Tokenizer
We won't be going over the details of the pre-trained tokenizer or model and only load a pre-trained one available from the huggingface model repository.
End of explanation
encoded_input = tokenizer(
    'What is the step by step guide to invest in share market in india?',
    'What is the step by step guide to invest in share market?'
)
encoded_input
Explanation: We can feed our tokenizer directly with a pair of sentences.
End of explanation
tokenizer.decode(encoded_input["input_ids"])
Explanation: Decoding the tokenized inputs, this model's tokenizer adds some special tokens such as, [SEP], that is used to indicate which token belongs to which segment/pair.
End of explanation
def tokenize_fn(examples):
    labels = [int(label) for label in examples['is_duplicate']]
    texts = [question['text'] for question in examples['questions']]
    texts1 = [text[0] for text in texts]
    texts2 = [text[1] for text in texts]
    tokenized_examples = tokenizer(texts1, texts2)
    tokenized_examples['labels'] = labels
    return tokenized_examples
dataset_dict_tokenized = dataset_dict.map(
    tokenize_fn,
    batched=True,
    num_proc=8,
    remove_columns=['is_duplicate', 'questions']
)
dataset_dict_tokenized
dataset_dict_tokenized['train'][0]
Explanation: The proprocessing step will be task specific, if we happen to be using another dataset, this function needs to be modified accordingly.
End of explanation
model_checkpoint = 'text_classification'
num_labels = 2
from transformers import (
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding
)
# we'll save the model after fine tuning it once, so we can skip the fine tuning part during
# the second round if we detect that we already have one available
if os.path.isdir(model_checkpoint):
    model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
else:
    model = AutoModelForSequenceClassification.from_pretrained(pretrained_model_name_or_path, num_labels=num_labels)
print('# of parameters: ', model.num_parameters())
model
data_collator = DataCollatorWithPadding(tokenizer, padding=True)
data_collator
Explanation: Model FineTuning
Having preprocessed our raw dataset, for our text classification task, we use AutoModelForSequenceClassification class to load the pre-trained model, the only other argument we need to specify is the number of class/label our text classification task has. Upon instantiating the model for the first time, we'll see some warnings generated, telling us we should fine tune this model on our down stream tasks before using it.
End of explanation
batch_size = 128
args = TrainingArguments(
    "quora",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=2,
    weight_decay=0.01,
    load_best_model_at_end=True
)
trainer = Trainer(
    model,
    args,
    data_collator=data_collator,
    train_dataset=dataset_dict_tokenized["train"],
    eval_dataset=dataset_dict_tokenized['val']
)
if not os.path.isdir(model_checkpoint):
    trainer.train()
    model.save_pretrained(model_checkpoint)
Explanation: We can perform all sorts of hyper parameter tuning on the fine tuning step, here we'll pick some default parameters for illustration purposes.
End of explanation
import torch
import torch.nn.functional as F
def predict(model, example, round_digits: int = 5):
    input_ids = example['input_ids'].to(model.device)
    attention_mask = example['attention_mask'].to(model.device)
    batch_labels = example['labels'].detach().cpu().numpy().tolist()
    model.eval()
    with torch.no_grad():
        batch_output = model(input_ids, attention_mask)
    batch_scores = F.softmax(batch_output.logits, dim=-1)[:, 1]
    batch_scores = np.round(batch_scores.detach().cpu().numpy(), round_digits).tolist()
    return batch_scores, batch_labels
from torch.utils.data import DataLoader
def predict_data_loader(model, data_loader: DataLoader) -> pd.DataFrame:
    scores = []
    labels = []
    for example in data_loader:
        batch_scores, batch_labels = predict(model, example)
        scores += batch_scores
        labels += batch_labels
    df_predictions = pd.DataFrame.from_dict({'scores': scores, 'labels': labels})
    return df_predictions
data_loader = DataLoader(dataset_dict_tokenized['test'], collate_fn=data_collator, batch_size=64)
start = time.time()
df_predictions = predict_data_loader(model, data_loader)
end = time.time()
print('elapsed: ', end - start)
print(df_predictions.shape)
df_predictions.head()
import sklearn.metrics as metrics
def compute_binary_classification_metrics(y_true, y_score, round_digits: int = 3):
    auc = metrics.roc_auc_score(y_true, y_score)
    log_loss = metrics.log_loss(y_true, y_score)
    precision, recall, threshold = metrics.precision_recall_curve(y_true, y_score)
    f1 = 2 * (precision * recall) / (precision + recall)
    mask = ~np.isnan(f1)
    f1 = f1[mask]
    precision = precision[mask]
    recall = recall[mask]
    best_index = np.argmax(f1)
    precision = precision[best_index]
    recall = recall[best_index]
    f1 = f1[best_index]
    return {
        'auc': round(auc, round_digits),
        'precision': round(precision, round_digits),
        'recall': round(recall, round_digits),
        'f1': round(f1, round_digits),
        'log_loss': round(log_loss, round_digits)
    }
compute_binary_classification_metrics(df_predictions['labels'], df_predictions['scores'])
Explanation: The next couple of code chunks performs batch inferencing on our dataset, and reports standard binary classification evaluation metrics.
End of explanation
data_loader = DataLoader(dataset_dict_tokenized['test'], collate_fn=data_collator, batch_size=8)
example = next(iter(data_loader))
model = model.to('cuda')
input_ids = example['input_ids'].to(model.device)
example['input_ids'][:2]
Explanation: Onnx Runtime
This section walks through the process of serializing our Pytorch model into ONNX format, and using ONNX runtime for inferencing.
Exporting the model can be done via the torch.onnx.export function, which requires a sample input.
End of explanation
opset_version = 12
onnx_model_path = 'text_classification.onnx'
torch.onnx.export(
    model,
    input_ids,
    onnx_model_path,
    opset_version=opset_version,
    input_names=['input_ids'],
    output_names=['output'],
    dynamic_axes={
        'input_ids': {0: 'batch_size', 1: 'sequence_len'},
        'output': {0: 'batch_size'},
    }
)
Explanation: For our neural network computational graph, we expect a single input with a dynamic batch size and sequence length, as well as a single output, which a dynamic batch size. We need to specify dynamic_axes, so during inferencing, onnx won't be limited to the sample input size we've provided when exporting the model into onnx format.
End of explanation
from onnxruntime.quantization import quantize_dynamic, QuantType
quantized_onnx_model_path = 'quantized_text_classification.onnx'
quantize_dynamic(
    onnx_model_path,
    quantized_onnx_model_path,
    activation_type=QuantType.QUInt8,
    weight_type=QuantType.QUInt8
)
print('ONNX full precision model size (MB):', os.path.getsize(onnx_model_path) / (1024 * 1024))
print('ONNX quantized model size (MB):', os.path.getsize(quantized_onnx_model_path) / (1024 * 1024))
Explanation: We'll also experiment dynamic quantization on our onnx model. Quantization here perfers to converting our weights or activations in our models from a floating point representation to an integer reresentation. Based on the documentation, this is the simplest form of applying quantization, where our weights are quantized ahead of time but the activations are dynamically quantized during inference time.
Given an already trained model, quantization provides a quick way for us to make a tradeoff between model performance versus model size and latency without the need of re-training our models with different parameter settings.
End of explanation
from onnxruntime import InferenceSession, SessionOptions
def create_inference_session(
    model_path: str,
    intra_op_num_threads: int = 24,
    provider: str = 'CPUExecutionProvider'
) -> InferenceSession: 
    options = SessionOptions()
    options.intra_op_num_threads = intra_op_num_threads
    # load the model as a onnx graph
    session = InferenceSession(model_path, options, providers=[provider])
    session.disable_fallback()
    return session
Explanation: One of the most important parameters to tune when using onnxruntime is intra_op_num_threads.
End of explanation
model = model.to('cuda')
# sample data
input_ids = dataset_dict_tokenized['test']['input_ids'][:1]
# gpu inferencing
model.eval()
with torch.no_grad():
    torch_input_ids = torch.LongTensor(input_ids).to(model.device)
    torch_output = model(torch_input_ids).logits.detach().cpu().numpy()
# onnx runtime inferencing
input_feed = {'input_ids': input_ids}
session = create_inference_session(onnx_model_path)
onnx_output = session.run(['output'], input_feed)[0]
quantized_session = create_inference_session(quantized_onnx_model_path)
quantized_onnx_output = quantized_session.run(['output'], input_feed)[0]
if np.allclose(torch_output, onnx_output, atol=1e-5):
    print("Exported model has been tested with ONNXRuntime, and the result looks good!")
print('input_ids: ', input_ids)
print('onnx output: ', onnx_output)
print('quantized onnx output: ', quantized_onnx_output)
Explanation: We test the inferencing API by feeding the original PyTorch model and ONNX model with a sample input.
End of explanation
from contextlib import contextmanager
@contextmanager
def track_infer_time(buffer):
    start = time.time()
    yield
    end = time.time()
    elasped_ms = (end - start) * 1000
    buffer.append(elasped_ms)
from tqdm import trange
def benchmark_pytorch(model, input_ids, n_rounds: int, n_warmup_rounds: int = 0):
    for _ in range(n_warmup_rounds):
        model(input_ids)
    time_buffer = []
    for _ in trange(n_rounds, desc="Benchmarking"):
        with track_infer_time(time_buffer):
            model(input_ids).logits.detach().cpu().numpy()
 
    return time_buffer
def benchmark_onnx(session, input_ids, n_rounds: int, n_warmup_rounds: int = 0):
    Expects the input_ids to be padded to the same sequence length, e.g. coming from DataLoader
    time.sleep(10)
    input_feed = {'input_ids': input_ids.detach().cpu().numpy()}
    for _ in range(n_warmup_rounds):
        session.run(['output'], input_feed)[0]
    time_buffer = []
    for _ in trange(n_rounds, desc="Benchmarking"):
        with track_infer_time(time_buffer):
            onnx_output = session.run(['output'], input_feed)[0]
 
    return time_buffer
def benchmark_dynamic_onnx(session, input_ids, n_rounds: int, n_warmup_rounds: int = 0):
    
    Expects the input ids to be of dynamic size (i.e. non-padded), e.g. coming from a slice
    of dataset.
    
    time.sleep(10)
    input_feed = {'input_ids': [input_ids[0]]}
    for _ in range(n_warmup_rounds):
        session.run(['output'], input_feed)[0]
    time_buffer = []
    for _ in trange(n_rounds, desc="Benchmarking"):
        with track_infer_time(time_buffer):
            batch_onnx_output = []
            for input_id in input_ids:
                input_feed = {'input_ids': [input_id]}
                onnx_output = session.run(['output'], input_feed)[0]
                batch_onnx_output.append(onnx_output)
            outputs = np.concatenate(batch_onnx_output)
    return time_buffer
n_rounds = 50
batch_sizes = [1, 4, 8, 16, 32, 64]
model_options = [
    'pytorch_gpu',  # pytorch cpu was really slow, not including it in the benchmark
    'onnx_cpu',
    'dynamic_onnx_cpu',
    'quantized_dynamic_onnx_cpu'
]
batch_time_results = {}
for batch_size in batch_sizes:
    results = {}
    for model_option in model_options:
        if model_option == 'pytorch_gpu':
            model = model.to('cuda')
            model = model.half()
            data_loader = DataLoader(
                dataset_dict_tokenized['test'],
                collate_fn=data_collator,
                batch_size=batch_size
            )
            example = next(iter(data_loader))
            input_ids = example['input_ids'].to(model.device)
            time_buffer = benchmark_pytorch(model, input_ids, n_rounds)
        elif model_option == 'onnx_cpu':
            data_loader = DataLoader(
                dataset_dict_tokenized['test'],
                collate_fn=data_collator,
                batch_size=batch_size
            )
            example = next(iter(data_loader))
            input_ids = example['input_ids'].to('cpu')
            time_buffer = benchmark_onnx(session, input_ids, n_rounds)
        elif model_option == 'dynamic_onnx_cpu':
            input_ids = dataset_dict_tokenized['test'][:batch_size]['input_ids']
            time_buffer = benchmark_dynamic_onnx(session, input_ids, n_rounds)
        elif model_option == 'quantized_dynamic_onnx_cpu':
            input_ids = dataset_dict_tokenized['test'][:batch_size]['input_ids']
            time_buffer = benchmark_dynamic_onnx(quantized_session, input_ids, n_rounds)
        else:
            raise ValueError(f'unknown model option {model_option}')
        results[model_option] = time_buffer
    time_results = {k: np.mean(v) for k, v in results.items()}
    for model_option, timing in time_results.items():
        if model_option in batch_time_results:
            batch_time_results[model_option].append(timing)
        else:
            batch_time_results[model_option] = [timing]
pd.DataFrame(batch_time_results)
# change default style figure and font size
plt.rcParams['figure.figsize'] = 10, 8
plt.rcParams['font.size'] = 12
for model_name, time_list in batch_time_results.items():
    plt.plot(time_list, label=model_name)
plt.xticks(range(len(batch_sizes)), batch_sizes)
plt.xlabel('batch size')
plt.ylabel('time (ms)')
plt.title('Inference time Comparison')
plt.legend()
plt.show()
Explanation: The next few section runs the benchmark on different batch sizes and inferencing options.
End of explanation
<END_TASK>