




import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


_lowerCamelCase			:	Dict			 =     logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm
  import torch_xla.debug.metrics as met


class        UpperCamelCase_					(				__snake_case  ):

    '''simple docstring'''






    def __init__(    self   :			str      ,				*UpperCAmelCase__   :			Dict      ,				UpperCAmelCase__   :			Any=None      ,				UpperCAmelCase__   :			Optional[Any]=None      ,				UpperCAmelCase__   :			str=None      ,				**UpperCAmelCase__   :			List[str])		->int:





           '''simple docstring'''

           super().__init__(*A_      ,				**A_)
           A__   				=					eval_examples
           A__   				=					post_process_function
           A__   				=					quant_trainer_args
           A__   				=					128  # default number of calibration samples


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[Any]      ,				UpperCAmelCase__   :			Optional[int]=None)		->Tuple:





           '''simple docstring'''

           if calib_dataset is None and self.calib_dataset is None:
                  raise ValueError('''Trainer: calibration requires an calib_dataset.''')
           A__   				=					calib_dataset if calib_dataset is not None else self.calib_dataset

           A__   				=					self._remove_unused_columns(A_      ,				description='''Calibration''')

           return DataLoader(
               A_      ,				batch_size=self.args.eval_batch_size      ,				collate_fn=self.data_collator      ,				drop_last=self.args.dataloader_drop_last      ,				num_workers=self.args.dataloader_num_workers      ,				pin_memory=self.args.dataloader_pin_memory      ,				shuffle=A_      ,				)


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple      ,				UpperCAmelCase__   :			List[str]=None)		->Tuple:





           '''simple docstring'''

           A__   				=					self.train_dataset if calib_dataset is None else calib_dataset
           A__   				=					self.get_calib_dataloader(A_)

           A__   				=					self.model
           quant_trainer.configure_model(A_      ,				self.quant_trainer_args      ,				calib=A_)
           model.eval()
           quant_trainer.enable_calibration(A_)

           logger.info('''***** Running calibration *****''')
           logger.info(f"""  Num examples = {self.calib_num}""")
           logger.info(f"""  Batch size = {calib_dataloader.batch_size}""")

           for step, inputs in enumerate(A_):
                  # Prediction step
                  A__   				=					self.prediction_step(A_      ,				A_      ,				prediction_loss_only=A_)
                  if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                         break

           quant_trainer.finish_calibration(A_      ,				self.quant_trainer_args)
           A__   				=					model


    def  SCREAMING_SNAKE_CASE							(    self   :			Dict      ,				UpperCAmelCase__   :			Dict=None      ,				UpperCAmelCase__   :			List[str]=None      ,				UpperCAmelCase__   :			int=None      ,				UpperCAmelCase__   :			Union[str, Any] = "eval")		->Any:





           '''simple docstring'''

           A__   				=					self.eval_dataset if eval_dataset is None else eval_dataset
           A__   				=					self.get_eval_dataloader(A_)
           A__   				=					self.eval_examples if eval_examples is None else eval_examples

           # Temporarily disable metric computation, we will do it in the loop here.
           A__   				=					self.compute_metrics
           A__   				=					None
           A__   				=					self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
           try:
                  A__   				=					eval_loop(
                      A_      ,				description='''Evaluation'''      ,				prediction_loss_only=True if compute_metrics is None else None      ,				ignore_keys=A_      ,				)
           finally:
                  A__   				=					compute_metrics

           if self.post_process_function is not None and self.compute_metrics is not None:
                  A__   				=					self.post_process_function(A_      ,				A_      ,				output.predictions)
                  A__   				=					self.compute_metrics(A_)

                  # Prefix all keys with metric_key_prefix + '_'
                  for key in list(metrics.keys()):
                         if not key.startswith(f"""{metric_key_prefix}_"""):
                                A__   				=					metrics.pop(A_)

                  self.log(A_)
           else:
                  A__   				=					{}

           if self.args.tpu_metrics_debug or self.args.debug:
                  # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
                  xm.master_print(met.metrics_report())

           A__   				=					self.callback_handler.on_evaluate(self.args      ,				self.state      ,				self.control      ,				A_)
           return metrics


    def  SCREAMING_SNAKE_CASE							(    self   :			List[str]      ,				UpperCAmelCase__   :			Any      ,				UpperCAmelCase__   :			int      ,				UpperCAmelCase__   :			Tuple=None      ,				UpperCAmelCase__   :			Any = "test")		->str:





           '''simple docstring'''

           A__   				=					self.get_test_dataloader(A_)

           # Temporarily disable metric computation, we will do it in the loop here.
           A__   				=					self.compute_metrics
           A__   				=					None
           A__   				=					self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
           try:
                  A__   				=					eval_loop(
                      A_      ,				description='''Prediction'''      ,				prediction_loss_only=True if compute_metrics is None else None      ,				ignore_keys=A_      ,				)
           finally:
                  A__   				=					compute_metrics

           if self.post_process_function is None or self.compute_metrics is None:
                  return output

           A__   				=					self.post_process_function(A_      ,				A_      ,				output.predictions      ,				'''predict''')
           A__   				=					self.compute_metrics(A_)

           # Prefix all keys with metric_key_prefix + '_'
           for key in list(metrics.keys()):
                  if not key.startswith(f"""{metric_key_prefix}_"""):
                         A__   				=					metrics.pop(A_)

           return PredictionOutput(predictions=predictions.predictions      ,				label_ids=predictions.label_ids      ,				metrics=A_)





    def  SCREAMING_SNAKE_CASE							(    self   :			Any      ,				UpperCAmelCase__   :			List[Any]="./")		->List[Any]:





           '''simple docstring'''

           A__   				=					self.eval_dataset
           A__   				=					self.get_eval_dataloader(A_)

           A__   				=					next(iter(A_))

           # saving device - to make it consistent
           A__   				=					torch.device('''cuda''' if torch.cuda.is_available() else '''cpu''')

           # convert to tuple
           A__   				=					tuple(v.to(A_) for k, v in batch.items())

           logger.info('''Converting model to be onnx compatible''')
           from pytorch_quantization.nn import TensorQuantizer

           A__   				=					True

           A__   				=					self.model.to(A_)

           model.eval()
           model.float()

           A__   				=					model.module if hasattr(A_      ,				'''module''') else model
           quant_trainer.configure_model(A_      ,				self.quant_trainer_args)

           A__   				=					os.path.join(A_      ,				'''model.onnx''')
           logger.info(f"""exporting model to {output_model_file}""")

           A__   				=					{0: "batch_size", 1: "seq_len"}

           torch.onnx.export(
               A_      ,				A_      ,				A_      ,				export_params=A_      ,				opset_version=13      ,				do_constant_folding=A_      ,				input_names=['''input_ids''', '''attention_mask''', '''token_type_ids''']      ,				output_names=['''output_start_logits''', '''output_end_logits''']      ,				dynamic_axes={
                   '''input_ids''': axes,
                   '''attention_mask''': axes,
                   '''token_type_ids''': axes,
                   '''output_start_logits''': axes,
                   '''output_end_logits''': axes,
               }      ,				verbose=A_      ,				)
           logger.info('''onnx export finished''')








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCamelCase			:	Tuple							 =     logging.get_logger(__name__)

_lowerCamelCase			:	Tuple							 =     {
    """tiiuae/falcon-40b""": """https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json""",
    """tiiuae/falcon-7b""": """https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json""",
}


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    UpperCAmelCase__													=   '''falcon'''
    UpperCAmelCase__													=   ['''past_key_values''']


    def __init__(    self   :			Optional[Any]      ,				UpperCAmelCase__   :			str=65_024      ,				UpperCAmelCase__   :			List[Any]=4_544      ,				UpperCAmelCase__   :			Optional[Any]=32      ,				UpperCAmelCase__   :			Dict=71      ,				UpperCAmelCase__   :			Any=1e-5      ,				UpperCAmelCase__   :			Optional[int]=0.02      ,				UpperCAmelCase__   :			Optional[int]=True      ,				UpperCAmelCase__   :			Union[str, Any]=0.0      ,				UpperCAmelCase__   :			List[str]=0.0      ,				UpperCAmelCase__   :			str=None      ,				UpperCAmelCase__   :			Dict=False      ,				UpperCAmelCase__   :			List[str]=False      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			Optional[Any]=False      ,				UpperCAmelCase__   :			Dict=11      ,				UpperCAmelCase__   :			str=11      ,				**UpperCAmelCase__   :			Dict      ,				)		->Any:





           '''simple docstring'''

           A__   				=					vocab_size
           # Backward compatibility with n_embed kwarg
           A__   				=					kwargs.pop('''n_embed'''      ,				UpperCAmelCase__)
           A__   				=					hidden_size if n_embed is None else n_embed
           A__   				=					num_hidden_layers
           A__   				=					num_attention_heads
           A__   				=					layer_norm_epsilon
           A__   				=					initializer_range
           A__   				=					use_cache
           A__   				=					hidden_dropout
           A__   				=					attention_dropout

           A__   				=					bos_token_id
           A__   				=					eos_token_id
           A__   				=					num_attention_heads if num_kv_heads is None else num_kv_heads
           A__   				=					alibi
           A__   				=					new_decoder_architecture
           A__   				=					multi_query  # Ignored when new_decoder_architecture is True
           A__   				=					parallel_attn
           A__   				=					bias

           super().__init__(bos_token_id=UpperCAmelCase__      ,				eos_token_id=UpperCAmelCase__      ,				**UpperCAmelCase__)


    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->int:





           '''simple docstring'''

           return self.hidden_size // self.num_attention_heads





    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Dict:





           '''simple docstring'''

           return not self.alibi


