codingutf8 2021 nvidia corporation licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license finetuning the library models for questionanswering on squad distilbert bert xlm xlnet import argparse import logging import os import time import timeit import datasets import numpy as np import pycuda autoinit noqa f401 import pycuda driver as cuda import tensorrt as trt import torch from absl import logging as absllogging from accelerate import accelerator from datasets import loaddataset loadmetric from torch utils data import dataloader from utilsqa import postprocessqapredictions import transformers from transformers import autotokenizer evalprediction defaultdatacollator setseed from transformers trainerptutils import nestedconcat nestedtruncate trtlogger trt loggertrt logger warning absllogger absllogging getabsllogger absllogger setlevellogging warning logger logging getloggername parser argparse argumentparser required parameters parser addargument onnxmodelpath defaultnone typestr requiredtrue helppath to onnx model parser addargument outputdir defaultnone typestr requiredtrue helpthe output directory where the model checkpoints and predictions will be written other parameters parser addargument tokenizername default typestr requiredtrue helppretrained tokenizer name or path if not the same as modelname parser addargument version2withnegative actionstoretrue helpif true the squad examples contain some that do not have an answer parser addargument nullscorediffthreshold typefloat default0 0 helpif nullscore bestnonnull is greater than the threshold predict null parser addargument maxseqlength default384 typeint help the maximum total input sequence length after wordpiece tokenization sequences longer than this will be truncated and sequences shorter than this will be padded parser addargument docstride default128 typeint helpwhen splitting up a long document into chunks how much stride to take between chunks parser addargumentperdeviceevalbatchsize default8 typeint helpbatch size per gpucpu for evaluation parser addargument nbestsize default20 typeint helpthe total number of nbest predictions to generate in the nbestpredictions json output file parser addargument maxanswerlength default30 typeint help the maximum length of an answer that can be generated this is needed because the start and end predictions are not conditioned on one another parser addargumentseed typeint default42 helprandom seed for initialization parser addargument datasetname typestr defaultnone requiredtrue helpthe name of the dataset to use via the datasets library parser addargument datasetconfigname typestr defaultnone helpthe configuration name of the dataset to use via the datasets library parser addargument preprocessingnumworkers typeint default4 helpa csv or a json file containing the training data parser addargumentoverwritecache actionstoretrue helpoverwrite the cached training and evaluation sets parser addargument fp16 actionstoretrue helpwhether to use 16bit mixed precision instead of 32bit parser addargument int8 actionstoretrue helpwhether to use int8 args parser parseargs if args tokenizername tokenizer autotokenizer frompretrainedargs tokenizername usefasttrue else raise valueerror you are instantiating a new tokenizer from scratch this is not supported by this script you can do it from another script save it and load it from here using tokenizername logger infotrainingevaluation parameters s args args evalbatchsize args perdeviceevalbatchsize inputshape args evalbatchsize args maxseqlength trt engine properties stricttypes true enginename tempenginebertfp32 engine if args fp16 enginename tempenginebertfp16 engine if args int8 enginename tempenginebertint8 engine import onnx file if not os path existstempengine os makedirstempengine explicitbatch 1 inttrt networkdefinitioncreationflag explicitbatch with trt buildertrtlogger as builder builder createnetworkexplicitbatch as network trt onnxparser network trtlogger as parser with openargs onnxmodelpath rb as model if not parser parsemodel read for error in rangeparser numerrors printparser geterrorerror query input names and shapes from parsed tensorrt network networkinputs network getinputi for i in rangenetwork numinputs inputnames input name for input in networkinputs ex actualinput1 with builder createbuilderconfig as config config maxworkspacesize 1 50 if stricttypes config setflagtrt builderflag stricttypes if args fp16 config setflagtrt builderflag fp16 if args int8 config setflagtrt builderflag int8 profile builder createoptimizationprofile config addoptimizationprofileprofile for i in rangeleninputnames profile setshapeinputnamesi inputshape inputshape inputshape engine builder buildenginenetwork config serializeengine and store in file can be directly loaded and deserialized with openenginename wb as f f writeengine serialize run inference with trt def modelinferinputs context dinputs houtput0 houtput1 doutput0 doutput1 stream inputids np asarrayinputsinputids dtypenp int32 attentionmask np asarrayinputsattentionmask dtypenp int32 tokentypeids np asarrayinputstokentypeids dtypenp int32 copy inputs cuda memcpyhtodasyncdinputs0 inputids ravel stream cuda memcpyhtodasyncdinputs1 attentionmask ravel stream cuda memcpyhtodasyncdinputs2 tokentypeids ravel stream start time starttime time time run inference context executeasync bindingsintdinp for dinp in dinputs intdoutput0 intdoutput1 streamhandlestream handle transfer predictions back from gpu cuda memcpydtohasynchoutput0 doutput0 stream cuda memcpydtohasynchoutput1 doutput1 stream synchronize the stream and take time stream synchronize end time endtime time time infertime endtime starttime outputs houtput0 houtput1 printoutputs return outputs infertime initialize the accelerator we will let the accelerator handle device placement for us in this example accelerator accelerator make one log on every process with the configuration for debugging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s levellogging info setup logging we only want one process per machine to log things on the screen accelerator islocalmainprocess is only true for one process per machine logger setlevellogging info if accelerator islocalmainprocess else logging error if accelerator islocalmainprocess datasets utils logging setverbositywarning transformers utils logging setverbosityinfo else datasets utils logging setverbosityerror transformers utils logging setverbosityerror if passed along set the training seed now if args seed is not none setseedargs seed get the datasets you can either provide your own csvjsontxt training and evaluation files see below or just provide the name of one of the public datasets available on the hub at https huggingface codatasets the dataset will be downloaded automatically from the datasets hub for csvjson files this script will use the column called text or the first column if no column called text is found you can easily tweak this behavior see below if args datasetname is not none downloading and loading a dataset from the hub rawdatasets loaddatasetargs datasetname args datasetconfigname else raise valueerrorevaluation requires a dataset name see more about loading any type of standard or custom dataset from files python dict pandas dataframe etc at https huggingface codocsdatasetsloadingdatasets preprocessing the datasets preprocessing is slighlty different for training and evaluation columnnames rawdatasetsvalidation columnnames questioncolumnname question if question in columnnames else columnnames0 contextcolumnname context if context in columnnames else columnnames1 answercolumnname answers if answers in columnnames else columnnames2 padding side determines if we do questioncontext or contextquestion padonright tokenizer paddingside right if args maxseqlength tokenizer modelmaxlength logger warning fthe maxseqlength passed args maxseqlength is larger than the maximum length for the fmodel tokenizer modelmaxlength using maxseqlengthtokenizer modelmaxlength maxseqlength minargs maxseqlength tokenizer modelmaxlength validation preprocessing def preparevalidationfeaturesexamples some of the questions have lots of whitespace on the left which is not useful and will make the truncation of the context fail the tokenized question will take a lots of space so we remove that left whitespace examplesquestioncolumnname q lstrip for q in examplesquestioncolumnname tokenize our examples with truncation and maybe padding but keep the overflows using a stride this results in one example possible giving several features when a context is long each of those features having a context that overlaps a bit the context of the previous feature tokenizedexamples tokenizer examplesquestioncolumnname if padonright else contextcolumnname examplescontextcolumnname if padonright else questioncolumnname truncationonlysecond if padonright else onlyfirst maxlengthmaxseqlength strideargs docstride returnoverflowingtokenstrue returnoffsetsmappingtrue paddingmaxlength since one example might give us several features if it has a long context we need a map from a feature to its corresponding example this key gives us just that samplemapping tokenizedexamples popoverflowtosamplemapping for evaluation we will need to convert our predictions to substrings of the context so we keep the corresponding exampleid and we will store the offset mappings tokenizedexamplesexampleid for i in rangelentokenizedexamplesinputids grab the sequence corresponding to that example to know what is the context and what is the question sequenceids tokenizedexamples sequenceidsi contextindex 1 if padonright else 0 one example can give several spans this is the index of the example containing this span of text sampleindex samplemappingi tokenizedexamplesexampleid appendexamplesidsampleindex set to none the offsetmapping that are not part of the context so it s easy to determine if a token position is part of the context or not tokenizedexamplesoffsetmappingi o if sequenceidsk contextindex else none for k o in enumeratetokenizedexamplesoffsetmappingi return tokenizedexamples evalexamples rawdatasetsvalidation validation feature creation evaldataset evalexamples map preparevalidationfeatures batchedtrue numprocargs preprocessingnumworkers removecolumnscolumnnames loadfromcachefilenot args overwritecache descrunning tokenizer on validation dataset datacollator defaultdatacollator evaldatasetformodel evaldataset removecolumnsexampleid offsetmapping evaldataloader dataloader evaldatasetformodel collatefndatacollator batchsizeargs perdeviceevalbatchsize postprocessing def postprocessingfunctionexamples features predictions stageeval postprocessing we match the start logits and end logits to answers in the original context predictions postprocessqapredictions examplesexamples featuresfeatures predictionspredictions version2withnegativeargs version2withnegative nbestsizeargs nbestsize maxanswerlengthargs maxanswerlength nullscorediffthresholdargs nullscorediffthreshold outputdirargs outputdir prefixstage format the result to the format the metric expects if args version2withnegative formattedpredictions id k predictiontext v noanswerprobability 0 0 for k v in predictions items else formattedpredictions id k predictiontext v for k v in predictions items references id exid answers exanswercolumnname for ex in examples return evalpredictionpredictionsformattedpredictions labelidsreferences metric loadmetricsquadv2 if args version2withnegative else squad evaluation logger infoloading onnx model s for evaluation args onnxmodelpath with openenginename rb as f trt runtimetrtlogger as runtime runtime deserializecudaengine f read as engine engine createexecutioncontext as context setup for trt inferrence for i in rangeleninputnames context setbindingshapei inputshape assert context allbindingshapesspecified def bindingnbytesbinding return trt volumeengine getbindingshapebinding engine getbindingdtypebinding itemsize allocate device memory for inputs and outputs dinputs cuda memallocbindingnbytesbinding for binding in engine if engine bindingisinputbinding allocate output buffer houtput0 cuda pagelockedemptytuplecontext getbindingshape3 dtypenp float32 houtput1 cuda pagelockedemptytuplecontext getbindingshape4 dtypenp float32 doutput0 cuda memallochoutput0 nbytes doutput1 cuda memallochoutput1 nbytes create a stream in which to copy inputsoutputs and run inference stream cuda stream evaluation logger info running evaluation logger infof num examples lenevaldataset logger infof batch size args perdeviceevalbatchsize totaltime 0 0 niter 0 starttime timeit defaulttimer allpreds none for step batch in enumerateevaldataloader outputs infertime modelinferbatch context dinputs houtput0 houtput1 doutput0 doutput1 stream totaltime infertime niter 1 startlogits endlogits outputs startlogits torch tensorstartlogits endlogits torch tensorendlogits necessary to pad predictions and labels for being gathered startlogits accelerator padacrossprocessesstartlogits dim1 padindex100 endlogits accelerator padacrossprocessesendlogits dim1 padindex100 logits accelerator gatherstartlogits cpu numpy accelerator gatherendlogits cpu numpy allpreds logits if allpreds is none else nestedconcatallpreds logits paddingindex100 if allpreds is not none allpreds nestedtruncateallpreds lenevaldataset evaltime timeit defaulttimer starttime logger info evaluation done in total f secs f sec per example evaltime evaltime lenevaldataset inference time from trt logger infoaverage inference time 3f ms formattotaltime 1000 niter logger infototal inference time 3f ms formattotaltime 1000 logger infototal number of inference d niter prediction postprocessingfunctionevalexamples evaldataset allpreds evalmetric metric computepredictionsprediction predictions referencesprediction labelids logger infofevaluation metrics evalmetric coding utf 8 2021 nvidia corporation licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license finetuning the library models for question answering on squad distilbert bert xlm xlnet noqa f401 required parameters other parameters trt engine properties import onnx file query input names and shapes from parsed tensorrt network ex actual_input1 serialize_engine and store in file can be directly loaded and deserialized run inference with trt copy inputs start time run inference transfer predictions back from gpu synchronize the stream and take time end time print outputs initialize the accelerator we will let the accelerator handle device placement for us in this example make one log on every process with the configuration for debugging setup logging we only want one process per machine to log things on the screen accelerator is_local_main_process is only true for one process per machine if passed along set the training seed now get the datasets you can either provide your own csv json txt training and evaluation files see below or just provide the name of one of the public datasets available on the hub at https huggingface co datasets the dataset will be downloaded automatically from the datasets hub for csv json files this script will use the column called text or the first column if no column called text is found you can easily tweak this behavior see below downloading and loading a dataset from the hub see more about loading any type of standard or custom dataset from files python dict pandas dataframe etc at https huggingface co docs datasets loading_datasets preprocessing the datasets preprocessing is slighlty different for training and evaluation padding side determines if we do question context or context question validation preprocessing some of the questions have lots of whitespace on the left which is not useful and will make the truncation of the context fail the tokenized question will take a lots of space so we remove that left whitespace tokenize our examples with truncation and maybe padding but keep the overflows using a stride this results in one example possible giving several features when a context is long each of those features having a context that overlaps a bit the context of the previous feature since one example might give us several features if it has a long context we need a map from a feature to its corresponding example this key gives us just that for evaluation we will need to convert our predictions to substrings of the context so we keep the corresponding example_id and we will store the offset mappings grab the sequence corresponding to that example to know what is the context and what is the question one example can give several spans this is the index of the example containing this span of text set to none the offset_mapping that are not part of the context so it s easy to determine if a token position is part of the context or not validation feature creation post processing post processing we match the start logits and end logits to answers in the original context format the result to the format the metric expects evaluation setup for trt inferrence allocate device memory for inputs and outputs allocate output buffer create a stream in which to copy inputs outputs and run inference evaluation necessary to pad predictions and labels for being gathered inference time from trt
import argparse
import logging
import os
import time
import timeit

import datasets
import numpy as np
import pycuda.autoinit  
import pycuda.driver as cuda
import tensorrt as trt
import torch
from absl import logging as absl_logging
from accelerate import Accelerator
from datasets import load_dataset, load_metric
from torch.utils.data import DataLoader
from utils_qa import postprocess_qa_predictions

import transformers
from transformers import AutoTokenizer, EvalPrediction, default_data_collator, set_seed
from transformers.trainer_pt_utils import nested_concat, nested_truncate


TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
absl_logger = absl_logging.get_absl_logger()
absl_logger.setLevel(logging.WARNING)

logger = logging.getLogger(__name__)

parser = argparse.ArgumentParser()


parser.add_argument(
    "--onnx_model_path",
    default=None,
    type=str,
    required=True,
    help="Path to ONNX model: ",
)

parser.add_argument(
    "--output_dir",
    default=None,
    type=str,
    required=True,
    help="The output directory where the model checkpoints and predictions will be written.",
)



parser.add_argument(
    "--tokenizer_name",
    default="",
    type=str,
    required=True,
    help="Pretrained tokenizer name or path if not the same as model_name",
)

parser.add_argument(
    "--version_2_with_negative",
    action="store_true",
    help="If true, the SQuAD examples contain some that do not have an answer.",
)
parser.add_argument(
    "--null_score_diff_threshold",
    type=float,
    default=0.0,
    help="If null_score - best_non_null is greater than the threshold predict null.",
)

parser.add_argument(
    "--max_seq_length",
    default=384,
    type=int,
    help=(
        "The maximum total input sequence length after WordPiece tokenization. Sequences "
        "longer than this will be truncated, and sequences shorter than this will be padded."
    ),
)
parser.add_argument(
    "--doc_stride",
    default=128,
    type=int,
    help="When splitting up a long document into chunks, how much stride to take between chunks.",
)

parser.add_argument("--per_device_eval_batch_size", default=8, type=int, help="Batch size per GPU/CPU for evaluation.")

parser.add_argument(
    "--n_best_size",
    default=20,
    type=int,
    help="The total number of n-best predictions to generate in the nbest_predictions.json output file.",
)
parser.add_argument(
    "--max_answer_length",
    default=30,
    type=int,
    help=(
        "The maximum length of an answer that can be generated. This is needed because the start "
        "and end predictions are not conditioned on one another."
    ),
)

parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")

parser.add_argument(
    "--dataset_name",
    type=str,
    default=None,
    required=True,
    help="The name of the dataset to use (via the datasets library).",
)
parser.add_argument(
    "--dataset_config_name",
    type=str,
    default=None,
    help="The configuration name of the dataset to use (via the datasets library).",
)
parser.add_argument(
    "--preprocessing_num_workers", type=int, default=4, help="A csv or a json file containing the training data."
)
parser.add_argument("--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets")
parser.add_argument(
    "--fp16",
    action="store_true",
    help="Whether to use 16-bit (mixed) precision instead of 32-bit",
)
parser.add_argument(
    "--int8",
    action="store_true",
    help="Whether to use INT8",
)

args = parser.parse_args()

if args.tokenizer_name:
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=True)
else:
    raise ValueError(
        "You are instantiating a new tokenizer from scratch. This is not supported by this script. "
        "You can do it from another script, save it, and load it from here, using --tokenizer_name."
    )

logger.info("Training/evaluation parameters %s", args)

args.eval_batch_size = args.per_device_eval_batch_size

INPUT_SHAPE = (args.eval_batch_size, args.max_seq_length)


STRICT_TYPES = True

engine_name = "temp_engine/bert-fp32.engine"
if args.fp16:
    engine_name = "temp_engine/bert-fp16.engine"
if args.int8:
    engine_name = "temp_engine/bert-int8.engine"


if not os.path.exists("temp_engine"):
    os.makedirs("temp_engine")

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(
    network, TRT_LOGGER
) as parser:
    with open(args.onnx_model_path, "rb") as model:
        if not parser.parse(model.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))

    
    network_inputs = [network.get_input(i) for i in range(network.num_inputs)]
    input_names = [_input.name for _input in network_inputs]  

    with builder.create_builder_config() as config:
        config.max_workspace_size = 1 << 50
        if STRICT_TYPES:
            config.set_flag(trt.BuilderFlag.STRICT_TYPES)
        if args.fp16:
            config.set_flag(trt.BuilderFlag.FP16)
        if args.int8:
            config.set_flag(trt.BuilderFlag.INT8)
        profile = builder.create_optimization_profile()
        config.add_optimization_profile(profile)
        for i in range(len(input_names)):
            profile.set_shape(input_names[i], INPUT_SHAPE, INPUT_SHAPE, INPUT_SHAPE)
        engine = builder.build_engine(network, config)

        
        with open(engine_name, "wb") as f:
            f.write(engine.serialize())



def model_infer(inputs, context, d_inputs, h_output0, h_output1, d_output0, d_output1, stream):
    input_ids = np.asarray(inputs["input_ids"], dtype=np.int32)
    attention_mask = np.asarray(inputs["attention_mask"], dtype=np.int32)
    token_type_ids = np.asarray(inputs["token_type_ids"], dtype=np.int32)

    
    cuda.memcpy_htod_async(d_inputs[0], input_ids.ravel(), stream)
    cuda.memcpy_htod_async(d_inputs[1], attention_mask.ravel(), stream)
    cuda.memcpy_htod_async(d_inputs[2], token_type_ids.ravel(), stream)
    
    start_time = time.time()
    
    context.execute_async(
        bindings=[int(d_inp) for d_inp in d_inputs] + [int(d_output0), int(d_output1)], stream_handle=stream.handle
    )
    
    cuda.memcpy_dtoh_async(h_output0, d_output0, stream)
    cuda.memcpy_dtoh_async(h_output1, d_output1, stream)
    
    stream.synchronize()
    
    end_time = time.time()
    infer_time = end_time - start_time
    outputs = (h_output0, h_output1)
    
    return outputs, infer_time



accelerator = Accelerator()

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)



logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
if accelerator.is_local_main_process:
    datasets.utils.logging.set_verbosity_warning()
    transformers.utils.logging.set_verbosity_info()
else:
    datasets.utils.logging.set_verbosity_error()
    transformers.utils.logging.set_verbosity_error()


if args.seed is not None:
    set_seed(args.seed)







if args.dataset_name is not None:
    
    raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
else:
    raise ValueError("Evaluation requires a dataset name")






column_names = raw_datasets["validation"].column_names

question_column_name = "question" if "question" in column_names else column_names[0]
context_column_name = "context" if "context" in column_names else column_names[1]
answer_column_name = "answers" if "answers" in column_names else column_names[2]


pad_on_right = tokenizer.padding_side == "right"

if args.max_seq_length > tokenizer.model_max_length:
    logger.warning(
        f"The max_seq_length passed ({args.max_seq_length}) is larger than the maximum length for the "
        f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
    )

max_seq_length = min(args.max_seq_length, tokenizer.model_max_length)



def prepare_validation_features(examples):
    
    
    
    examples[question_column_name] = [q.lstrip() for q in examples[question_column_name]]

    
    
    
    tokenized_examples = tokenizer(
        examples[question_column_name if pad_on_right else context_column_name],
        examples[context_column_name if pad_on_right else question_column_name],
        truncation="only_second" if pad_on_right else "only_first",
        max_length=max_seq_length,
        stride=args.doc_stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

    
    
    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")

    
    
    tokenized_examples["example_id"] = []

    for i in range(len(tokenized_examples["input_ids"])):
        
        sequence_ids = tokenized_examples.sequence_ids(i)
        context_index = 1 if pad_on_right else 0

        
        sample_index = sample_mapping[i]
        tokenized_examples["example_id"].append(examples["id"][sample_index])

        
        
        tokenized_examples["offset_mapping"][i] = [
            (o if sequence_ids[k] == context_index else None)
            for k, o in enumerate(tokenized_examples["offset_mapping"][i])
        ]

    return tokenized_examples


eval_examples = raw_datasets["validation"]

eval_dataset = eval_examples.map(
    prepare_validation_features,
    batched=True,
    num_proc=args.preprocessing_num_workers,
    remove_columns=column_names,
    load_from_cache_file=not args.overwrite_cache,
    desc="Running tokenizer on validation dataset",
)

data_collator = default_data_collator

eval_dataset_for_model = eval_dataset.remove_columns(["example_id", "offset_mapping"])
eval_dataloader = DataLoader(
    eval_dataset_for_model, collate_fn=data_collator, batch_size=args.per_device_eval_batch_size
)



def post_processing_function(examples, features, predictions, stage="eval"):
    
    predictions = postprocess_qa_predictions(
        examples=examples,
        features=features,
        predictions=predictions,
        version_2_with_negative=args.version_2_with_negative,
        n_best_size=args.n_best_size,
        max_answer_length=args.max_answer_length,
        null_score_diff_threshold=args.null_score_diff_threshold,
        output_dir=args.output_dir,
        prefix=stage,
    )
    
    if args.version_2_with_negative:
        formatted_predictions = [
            {"id": k, "prediction_text": v, "no_answer_probability": 0.0} for k, v in predictions.items()
        ]
    else:
        formatted_predictions = [{"id": k, "prediction_text": v} for k, v in predictions.items()]

    references = [{"id": ex["id"], "answers": ex[answer_column_name]} for ex in examples]
    return EvalPrediction(predictions=formatted_predictions, label_ids=references)


metric = load_metric("squad_v2" if args.version_2_with_negative else "squad")


logger.info("Loading ONNX model %s for evaluation", args.onnx_model_path)
with open(engine_name, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(
    f.read()
) as engine, engine.create_execution_context() as context:
    
    for i in range(len(input_names)):
        context.set_binding_shape(i, INPUT_SHAPE)
    assert context.all_binding_shapes_specified

    def binding_nbytes(binding):
        return trt.volume(engine.get_binding_shape(binding)) * engine.get_binding_dtype(binding).itemsize

    
    d_inputs = [cuda.mem_alloc(binding_nbytes(binding)) for binding in engine if engine.binding_is_input(binding)]

    
    h_output0 = cuda.pagelocked_empty(tuple(context.get_binding_shape(3)), dtype=np.float32)
    h_output1 = cuda.pagelocked_empty(tuple(context.get_binding_shape(4)), dtype=np.float32)
    d_output0 = cuda.mem_alloc(h_output0.nbytes)
    d_output1 = cuda.mem_alloc(h_output1.nbytes)

    
    stream = cuda.Stream()

    
    logger.info("***** Running Evaluation *****")
    logger.info(f"  Num examples = {len(eval_dataset)}")
    logger.info(f"  Batch size = {args.per_device_eval_batch_size}")

    total_time = 0.0
    niter = 0
    start_time = timeit.default_timer()

    all_preds = None
    for step, batch in enumerate(eval_dataloader):
        outputs, infer_time = model_infer(batch, context, d_inputs, h_output0, h_output1, d_output0, d_output1, stream)
        total_time += infer_time
        niter += 1

        start_logits, end_logits = outputs
        start_logits = torch.tensor(start_logits)
        end_logits = torch.tensor(end_logits)

        
        start_logits = accelerator.pad_across_processes(start_logits, dim=1, pad_index=-100)
        end_logits = accelerator.pad_across_processes(end_logits, dim=1, pad_index=-100)

        logits = (accelerator.gather(start_logits).cpu().numpy(), accelerator.gather(end_logits).cpu().numpy())
        all_preds = logits if all_preds is None else nested_concat(all_preds, logits, padding_index=-100)

    if all_preds is not None:
        all_preds = nested_truncate(all_preds, len(eval_dataset))

    evalTime = timeit.default_timer() - start_time
    logger.info("  Evaluation done in total %f secs (%f sec per example)", evalTime, evalTime / len(eval_dataset))
    
    logger.info("Average Inference Time = {:.3f} ms".format(total_time * 1000 / niter))
    logger.info("Total Inference Time =  {:.3f} ms".format(total_time * 1000))
    logger.info("Total Number of Inference =  %d", niter)

prediction = post_processing_function(eval_examples, eval_dataset, all_preds)
eval_metric = metric.compute(predictions=prediction.predictions, references=prediction.label_ids)
logger.info(f"Evaluation metrics: {eval_metric}")