



'''simple docstring'''





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


a_						:     List[str]	     =		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    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     def __init__(self				,			*__a				,			__a=None				,			__a=None				,			__a=None				,			**__a    ):






           '''simple docstring'''



           super().__init__(*__a				,			**__a    )
           lowerCamelCase								=		eval_examples
           lowerCamelCase								=		post_process_function
           lowerCamelCase								=		quant_trainer_args
           lowerCamelCase								=		1_28  # default number of calibration samples







     def 						_a      (self				,			__a=None    ):






           '''simple docstring'''



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

           lowerCamelCase								=		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 						_a      (self				,			__a=None    ):






           '''simple docstring'''



           lowerCamelCase								=		self.train_dataset if calib_dataset is None else calib_dataset
           lowerCamelCase								=		self.get_calib_dataloader(__a    )

           lowerCamelCase								=		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
                 lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase								=		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    )
           lowerCamelCase								=		model







     def 						_a      (self				,			__a=None				,			__a=None				,			__a=None				,			__a = "eval"    ):






           '''simple docstring'''



           lowerCamelCase								=		self.eval_dataset if eval_dataset is None else eval_dataset
           lowerCamelCase								=		self.get_eval_dataloader(__a    )
           lowerCamelCase								=		self.eval_examples if eval_examples is None else eval_examples

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

           if self.post_process_function is not None and self.compute_metrics is not None:
                 lowerCamelCase								=		self.post_process_function(__a				,			__a				,			output.predictions    )
                 lowerCamelCase								=		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}_"""    ):
                             lowerCamelCase								=		metrics.pop(__a    )

                 self.log(__a    )
           else:
                 lowerCamelCase								=		{}

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

           lowerCamelCase								=		self.callback_handler.on_evaluate(self.args				,			self.state				,			self.control				,			__a    )
           return metrics







     def 						_a      (self				,			__a				,			__a				,			__a=None				,			__a = "test"    ):






           '''simple docstring'''



           lowerCamelCase								=		self.get_test_dataloader(__a    )

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

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

           lowerCamelCase								=		self.post_process_function(__a				,			__a				,			output.predictions				,			"predict"    )
           lowerCamelCase								=		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}_"""    ):
                       lowerCamelCase								=		metrics.pop(__a    )

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







     def 						_a      (self				,			__a="./"    ):






           '''simple docstring'''



           lowerCamelCase								=		self.eval_dataset
           lowerCamelCase								=		self.get_eval_dataloader(__a    )

           lowerCamelCase								=		next(iter(__a    )    )

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

           # convert to tuple
           lowerCamelCase								=		tuple(v.to(__a    ) for k, v in batch.items()    )

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

           lowerCamelCase								=		True

           lowerCamelCase								=		self.model.to(__a    )

           model.eval()
           model.float()

           lowerCamelCase								=		model.module if hasattr(__a				,			"module"    ) else model
           quant_trainer.configure_model(__a				,			self.quant_trainer_args    )

           lowerCamelCase								=		os.path.join(__a				,			"model.onnx"    )
           logger.info(F"""exporting model to {output_model_file}"""    )

           lowerCamelCase								=		{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"    )




import argparse
import os

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_task_guides.py
a_						:     Tuple	     =		'src/transformers'
a_						:     Any	     =		'docs/source/en/tasks'

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__ ):


      """simple docstring"""



      with open(UpperCAmelCase__   ,   "r"   ,   encoding="utf-8"   ,   newline="\n" ) as f:
            lowerCamelCase								=		f.readlines()
      # Find the start prompt.
      lowerCamelCase								=		0
      while not lines[start_index].startswith(UpperCAmelCase__ ):
            start_index += 1
      start_index += 1

      lowerCamelCase								=		start_index
      while not lines[end_index].startswith(UpperCAmelCase__ ):
            end_index += 1
      end_index -= 1

      while len(lines[start_index] ) <= 1:
            start_index += 1
      while len(lines[end_index] ) <= 1:
            end_index -= 1
      end_index += 1
      return "".join(lines[start_index:end_index] ), start_index, end_index, lines


# This is to make sure the transformers module imported is the one in the repo.
a_						:     Union[str, Any]	     =		direct_transformers_import(TRANSFORMERS_PATH)

a_						:     List[Any]	     =		{
    'asr.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_CTC_MAPPING_NAMES,
    'audio_classification.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES,
    'language_modeling.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
    'image_classification.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES,
    'masked_language_modeling.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_MASKED_LM_MAPPING_NAMES,
    'multiple_choice.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES,
    'object_detection.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES,
    'question_answering.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES,
    'semantic_segmentation.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES,
    'sequence_classification.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES,
    'summarization.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
    'token_classification.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES,
    'translation.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
    'video_classification.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES,
    'document_question_answering.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES,
    'monocular_depth_estimation.md': transformers_module.models.auto.modeling_auto.MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES,
}

# This list contains model types used in some task guides that are not in `CONFIG_MAPPING_NAMES` (therefore not in any
# `MODEL_MAPPING_NAMES` or any `MODEL_FOR_XXX_MAPPING_NAMES`).
a_						:     Union[str, Any]	     =		{
    'summarization.md': ('nllb',),
    'translation.md': ('nllb',),
}

def     __lowercase(	UpperCAmelCase__ ):


      """simple docstring"""



      lowerCamelCase								=		TASK_GUIDE_TO_MODELS[task_guide]
      lowerCamelCase								=		SPECIAL_TASK_GUIDE_TO_MODEL_TYPES.get(UpperCAmelCase__   ,   set() )
      lowerCamelCase								=		{
          code: name
          for code, name in transformers_module.MODEL_NAMES_MAPPING.items()
          if (code in model_maping_names or code in special_model_types)
      }
      return ", ".join([F"""[{name}](../model_doc/{code})""" for code, name in model_names.items()] ) + "\n"

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__=False ):


      """simple docstring"""



      lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase								=		_find_text_in_file(
          filename=os.path.join(UpperCAmelCase__   ,   UpperCAmelCase__ )   ,   start_prompt="<!--This tip is automatically generated by `make fix-copies`, do not fill manually!-->"   ,   end_prompt="<!--End of the generated tip-->"   ,   )

      lowerCamelCase								=		get_model_list_for_task(UpperCAmelCase__ )

      if current_list != new_list:
            if overwrite:
                  with open(os.path.join(UpperCAmelCase__   ,   UpperCAmelCase__ )   ,   "w"   ,   encoding="utf-8"   ,   newline="\n" ) as f:
                        f.writelines(lines[:start_index] + [new_list] + lines[end_index:] )
            else:
                  raise ValueError(
                      F"""The list of models that can be used in the {task_guide} guide needs an update. Run `make fix-copies`"""
                      " to fix this." )


if __name__ == "__main__":
     a_						:     str	     =		argparse.ArgumentParser()
     parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
     a_						:     str	     =		parser.parse_args()

     for task_guide in TASK_GUIDE_TO_MODELS.keys():
          check_model_list_for_task(task_guide, args.fix_and_overwrite)