



'''simple docstring'''





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 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


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

lowercase_		       =  logging.getLogger(__name__)

lowercase_		       =  argparse.ArgumentParser()

# Required parameters
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.",
)

# Other parameters

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",
)

lowercase_		       =  parser.parse_args()

if args.tokenizer_name:
  lowercase_		       =  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)

lowercase_		       =  args.per_device_eval_batch_size

lowercase_		       =  (args.eval_batch_size, args.max_seq_length)

# TRT Engine properties
lowercase_		       =  True

lowercase_		       =  "temp_engine/bert-fp32.engine"
if args.fpaa:
  lowercase_		       =  "temp_engine/bert-fp16.engine"
if args.inta:
  lowercase_		       =  "temp_engine/bert-int8.engine"

# import ONNX file
if not os.path.exists("temp_engine"):
  os.makedirs("temp_engine")

lowercase_		       =  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))

    # Query input names and shapes from parsed TensorRT network
  lowercase_		       =  [network.get_input(i) for i in range(network.num_inputs)]
  lowercase_		       =  [_input.name for _input in network_inputs]  # ex: ["actual_input1"]

  with builder.create_builder_config() as config:
    lowercase_		       =  1 << 50
    if STRICT_TYPES:
      config.set_flag(trt.BuilderFlag.STRICT_TYPES)
    if args.fpaa:
      config.set_flag(trt.BuilderFlag.FPaa)
    if args.inta:
      config.set_flag(trt.BuilderFlag.INTa)
    lowercase_		       =  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)
    lowercase_		       =  builder.build_engine(network, config)

    # serialize_engine and store in file (can be directly loaded and deserialized):
    with open(engine_name, "wb") as f:
      f.write(engine.serialize())




def lowerCAmelCase  (__A  ,	__A  ,	__A  ,	__A  ,	__A  ,	__A  ,	__A  ,	__A):



     """simple docstring"""



     _a      =  np.asarray(inputs['''input_ids''']  ,	dtype=np.intaa)
     _a      =  np.asarray(inputs['''attention_mask''']  ,	dtype=np.intaa)
     _a      =  np.asarray(inputs['''token_type_ids''']  ,	dtype=np.intaa)

     # Copy inputs
     cuda.memcpy_htod_async(d_inputs[0]  ,	input_ids.ravel()  ,	__A)
     cuda.memcpy_htod_async(d_inputs[1]  ,	attention_mask.ravel()  ,	__A)
     cuda.memcpy_htod_async(d_inputs[2]  ,	token_type_ids.ravel()  ,	__A)
     # start time
     _a      =  time.time()
     # Run inference
     context.execute_async(
         bindings=[int(__A) for d_inp in d_inputs] + [int(__A), int(__A)]  ,	stream_handle=stream.handle)
     # Transfer predictions back from GPU
     cuda.memcpy_dtoh_async(__A  ,	__A  ,	__A)
     cuda.memcpy_dtoh_async(__A  ,	__A  ,	__A)
     # Synchronize the stream and take time
     stream.synchronize()
     # end time
     _a      =  time.time()
     _a      =  end_time - start_time
     _a      =  (h_outputa, h_outputa)
     # print(outputs)
     return outputs, infer_time


# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
lowercase_		       =  Accelerator()
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)

# 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.
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 passed along, set the training seed now.
if args.seed is not None:
  set_seed(args.seed)

# 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).
if args.dataset_name is not None:
  # Downloading and loading a dataset from the hub.
  lowercase_		       =  load_dataset(args.dataset_name, args.dataset_config_name)
else:
  raise ValueError("Evaluation 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.co/docs/datasets/loading_datasets.html.

# Preprocessing the datasets.
# Preprocessing is slighlty different for training and evaluation.

lowercase_		       =  raw_datasets["validation"].column_names

lowercase_		       =  "question" if "question" in column_names else column_names[0]
lowercase_		       =  "context" if "context" in column_names else column_names[1]
lowercase_		       =  "answers" if "answers" in column_names else column_names[2]

# Padding side determines if we do (question|context) or (context|question).
lowercase_		       =  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}."""
  )

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




def lowerCAmelCase  (__A):



     """simple docstring"""



     _a      =  [q.lstrip() for q in examples[question_column_name]]

     # 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.
     _a      =  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=__A  ,	stride=args.doc_stride  ,	return_overflowing_tokens=__A  ,	return_offsets_mapping=__A  ,	padding='''max_length'''  ,	)

     # 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.
     _a      =  tokenized_examples.pop('''overflow_to_sample_mapping''')

     # 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.
     _a      =  []

     for i in range(len(tokenized_examples['''input_ids'''])):
          # Grab the sequence corresponding to that example (to know what is the context and what is the question).
          _a      =  tokenized_examples.sequence_ids(__A)
          _a      =  1 if pad_on_right else 0

          # One example can give several spans, this is the index of the example containing this span of text.
          _a      =  sample_mapping[i]
          tokenized_examples["example_id"].append(examples['''id'''][sample_index])

          # 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.
          _a      =  [
              (o if sequence_ids[k] == context_index else None)
              for k, o in enumerate(tokenized_examples['''offset_mapping'''][i])
          ]

     return tokenized_examples


lowercase_		       =  raw_datasets["validation"]
# Validation Feature Creation
lowercase_		       =  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",
)

lowercase_		       =  default_data_collator

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




def lowerCAmelCase  (__A  ,	__A  ,	__A  ,	__A="eval"):



     """simple docstring"""



     _a      =  postprocess_qa_predictions(
         examples=__A  ,	features=__A  ,	predictions=__A  ,	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=__A  ,	)
     # Format the result to the format the metric expects.
     if args.version_2_with_negative:
          _a      =  [
              {'''id''': k, '''prediction_text''': v, '''no_answer_probability''': 0.0} for k, v in predictions.items()
          ]
     else:
          _a      =  [{'''id''': k, '''prediction_text''': v} for k, v in predictions.items()]

     _a      =  [{'''id''': ex['''id'''], '''answers''': ex[answer_column_name]} for ex in examples]
     return EvalPrediction(predictions=__A  ,	label_ids=__A)


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

# Evaluation!
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:
  # setup for TRT inferrence
  for i in range(len(input_names)):
    context.set_binding_shape(i, INPUT_SHAPE)
  assert context.all_binding_shapes_specified




  def lowerCAmelCase  (__A):



       """simple docstring"""



       return trt.volume(engine.get_binding_shape(__A)) * engine.get_binding_dtype(__A).itemsize

  # Allocate device memory for inputs and outputs.
  lowercase_		       =  [cuda.mem_alloc(binding_nbytes(binding)) for binding in engine if engine.binding_is_input(binding)]

  # Allocate output buffer
  lowercase_		       =  cuda.pagelocked_empty(tuple(context.get_binding_shape(3)), dtype=np.floataa)
  lowercase_		       =  cuda.pagelocked_empty(tuple(context.get_binding_shape(4)), dtype=np.floataa)
  lowercase_		       =  cuda.mem_alloc(h_outputa.nbytes)
  lowercase_		       =  cuda.mem_alloc(h_outputa.nbytes)

  # Create a stream in which to copy inputs/outputs and run inference.
  lowercase_		       =  cuda.Stream()

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

  lowercase_		       =  0.0
  lowercase_		       =  0
  lowercase_		       =  timeit.default_timer()

  lowercase_		       =  None
  for step, batch in enumerate(eval_dataloader):
    lowercase_       ,					lowercase_		       =  model_infer(batch, context, d_inputs, h_outputa, h_outputa, d_outputa, d_outputa, stream)
    total_time += infer_time
    niter += 1

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

    # necessary to pad predictions and labels for being gathered
    lowercase_		       =  accelerator.pad_across_processes(start_logits, dim=1, pad_index=-100)
    lowercase_		       =  accelerator.pad_across_processes(end_logits, dim=1, pad_index=-100)

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

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

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

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








'''simple docstring'''





# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)






@flax.struct.dataclass
class __A			:



 '''simple docstring'''







 __lowerCamelCase    :       CommonSchedulerState

 # setable values
 __lowerCamelCase    :       jnp.ndarray
 __lowerCamelCase    :       jnp.ndarray
 __lowerCamelCase    :       Optional[int]          =       None



 @classmethod
 def 					a__  (cls ,   A ,   A ,   A       )    ->     str:


      """simple docstring"""



      return cls(common=A ,   init_noise_sigma=A ,   timesteps=A       )






@dataclass
class __A			(  A						):



 '''simple docstring'''







 __lowerCamelCase    :       DDPMSchedulerState






class __A			(  A   ,   A						):



 '''simple docstring'''







 __lowerCamelCase    :       Dict          =       [e.name for e in FlaxKarrasDiffusionSchedulers]

 __lowerCamelCase    :       jnp.dtype



 @property
 def 					a__  (self       )    ->     List[str]:


      """simple docstring"""



      return True



 @register_to_config
 def __init__(self ,   A = 1_000 ,   A = 0.0001 ,   A = 0.02 ,   A = "linear" ,   A = None ,   A = "fixed_small" ,   A = True ,   A = "epsilon" ,   A = jnp.floataa ,   )    ->     Union[str, Any]:


      """simple docstring"""



      _a      =  dtype



 def 					a__  (self ,   A = None       )    ->     DDPMSchedulerState:


      """simple docstring"""



      if common is None:
           _a      =  CommonSchedulerState.create(self       )

      # standard deviation of the initial noise distribution
      _a      =  jnp.array(1.0 ,   dtype=self.dtype       )

      _a      =  jnp.arange(0 ,   self.config.num_train_timesteps       ).round()[::-1]

      return DDPMSchedulerState.create(
          common=A ,   init_noise_sigma=A ,   timesteps=A ,   )



 def 					a__  (self ,   A ,   A ,   A = None       )    ->     jnp.ndarray:


      """simple docstring"""



      return sample



 def 					a__  (self ,   A ,   A ,   A = ()       )    ->     DDPMSchedulerState:


      """simple docstring"""



      _a      =  self.config.num_train_timesteps // num_inference_steps
      # creates integer timesteps by multiplying by ratio
      # rounding to avoid issues when num_inference_step is power of 3
      _a      =  (jnp.arange(0 ,   A       ) * step_ratio).round()[::-1]

      return state.replace(
          num_inference_steps=A ,   timesteps=A ,   )



 def 					a__  (self ,   A ,   A ,   A=None ,   A=None       )    ->     int:


      """simple docstring"""



      _a      =  state.common.alphas_cumprod[t]
      _a      =  jnp.where(t > 0 ,   state.common.alphas_cumprod[t - 1] ,   jnp.array(1.0 ,   dtype=self.dtype       )       )

      # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
      # and sample from it to get previous sample
      # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
      _a      =  (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t]

      if variance_type is None:
           _a      =  self.config.variance_type

      # hacks - were probably added for training stability
      if variance_type == "fixed_small":
           _a      =  jnp.clip(A ,   a_min=1E-20       )
      # for rl-diffuser https://arxiv.org/abs/2205.09991
      elif variance_type == "fixed_small_log":
           _a      =  jnp.log(jnp.clip(A ,   a_min=1E-20       )       )
      elif variance_type == "fixed_large":
           _a      =  state.common.betas[t]
      elif variance_type == "fixed_large_log":
           # Glide max_log
           _a      =  jnp.log(state.common.betas[t]       )
      elif variance_type == "learned":
           return predicted_variance
      elif variance_type == "learned_range":
           _a      =  variance
           _a      =  state.common.betas[t]
           _a      =  (predicted_variance + 1) / 2
           _a      =  frac * max_log + (1 - frac) * min_log

      return variance



 def 					a__  (self ,   A ,   A ,   A ,   A ,   A = None ,   A = True ,   )    ->     Union[FlaxDDPMSchedulerOutput, Tuple]:


      """simple docstring"""



      _a      =  timestep

      if key is None:
           _a      =  jax.random.PRNGKey(0       )

      if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
           _a  ,    _a      =  jnp.split(A ,   sample.shape[1] ,   axis=1       )
      else:
           _a      =  None

      # 1. compute alphas, betas
      _a      =  state.common.alphas_cumprod[t]
      _a      =  jnp.where(t > 0 ,   state.common.alphas_cumprod[t - 1] ,   jnp.array(1.0 ,   dtype=self.dtype       )       )
      _a      =  1 - alpha_prod_t
      _a      =  1 - alpha_prod_t_prev

      # 2. compute predicted original sample from predicted noise also called
      # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
      if self.config.prediction_type == "epsilon":
           _a      =  (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
      elif self.config.prediction_type == "sample":
           _a      =  model_output
      elif self.config.prediction_type == "v_prediction":
           _a      =  (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
      else:
           raise ValueError(
               f'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` '''
               ''' for the FlaxDDPMScheduler.'''       )

      # 3. Clip "predicted x_0"
      if self.config.clip_sample:
           _a      =  jnp.clip(A ,   -1 ,   1       )

      # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
      # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
      _a      =  (alpha_prod_t_prev ** 0.5 * state.common.betas[t]) / beta_prod_t
      _a      =  state.common.alphas[t] ** 0.5 * beta_prod_t_prev / beta_prod_t

      # 5. Compute predicted previous sample µ_t
      # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
      _a      =  pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

      # 6. Add noise
      def random_variance():
           _a      =  jax.random.split(A ,   num=1       )
           _a      =  jax.random.normal(A ,   shape=model_output.shape ,   dtype=self.dtype       )
           return (self._get_variance(A ,   A ,   predicted_variance=A       ) ** 0.5) * noise

      _a      =  jnp.where(t > 0 ,   random_variance() ,   jnp.zeros(model_output.shape ,   dtype=self.dtype       )       )

      _a      =  pred_prev_sample + variance

      if not return_dict:
           return (pred_prev_sample, state)

      return FlaxDDPMSchedulerOutput(prev_sample=A ,   state=A       )



 def 					a__  (self ,   A ,   A ,   A ,   A ,   )    ->     jnp.ndarray:


      """simple docstring"""



      return add_noise_common(state.common ,   A ,   A ,   A       )



 def 					a__  (self ,   A ,   A ,   A ,   A ,   )    ->     jnp.ndarray:


      """simple docstring"""



      return get_velocity_common(state.common ,   A ,   A ,   A       )



 def __len__(self       )    ->     Tuple:


      """simple docstring"""



      return self.config.num_train_timesteps



