






import collections
import os
import re
from pathlib import Path


_SCREAMING_SNAKE_CASE   =  "src/transformers"


# Matches is_xxx_available()
_SCREAMING_SNAKE_CASE   =  re.compile(R"is\_([a-z_]*)_available()")
# Catches a one-line _import_struct = {xxx}
_SCREAMING_SNAKE_CASE   =  re.compile(R"^_import_structure\s+=\s+\{([^\}]+)\}")
# Catches a line with a key-values pattern: "bla": ["foo", "bar"]
_SCREAMING_SNAKE_CASE   =  re.compile(R"\s+\"\S*\":\s+\[([^\]]*)\]")
# Catches a line if not is_foo_available
_SCREAMING_SNAKE_CASE   =  re.compile(R"^\s*if\s+not\s+is\_[a-z_]*\_available\(\)")
# Catches a line _import_struct["bla"].append("foo")
_SCREAMING_SNAKE_CASE   =  re.compile(R"^\s*_import_structure\[\"\S*\"\]\.append\(\"(\S*)\"\)")
# Catches a line _import_struct["bla"].extend(["foo", "bar"]) or _import_struct["bla"] = ["foo", "bar"]
_SCREAMING_SNAKE_CASE   =  re.compile(R"^\s*_import_structure\[\S*\](?:\.extend\(|\s*=\s+)\[([^\]]*)\]")
# Catches a line with an object between quotes and a comma:     "MyModel",
_SCREAMING_SNAKE_CASE   =  re.compile(R"^\s+\"([^\"]+)\",")
# Catches a line with objects between brackets only:    ["foo", "bar"],
_SCREAMING_SNAKE_CASE   =  re.compile(R"^\s+\[([^\]]+)\]")
# Catches a line with from foo import bar, bla, boo
_SCREAMING_SNAKE_CASE   =  re.compile(R"\s+from\s+\S*\s+import\s+([^\(\s].*)\n")
# Catches a line with try:
_SCREAMING_SNAKE_CASE   =  re.compile(R"^\s*try:")
# Catches a line with else:
_SCREAMING_SNAKE_CASE   =  re.compile(R"^\s*else:")





def        _snake_case   (_snake_case				:       Dict)					->	int:
   if _re_test_backend.search(_snake_case) is None:
      return None
   _lowercase              =[b[0] for b in _re_backend.findall(_snake_case)]
   backends.sort()
   return "_and_".join(_snake_case)





def        _snake_case   (_snake_case				:       int)					->	Any:
   with open(_snake_case      ,    'r'      ,    encoding='utf-8'      ,    newline='\n') as f:
      _lowercase              =f.readlines()

   _lowercase              =0
   while line_index < len(_snake_case) and not lines[line_index].startswith('_import_structure = {'):
      line_index += 1

   # If this is a traditional init, just return.
   if line_index >= len(_snake_case):
      return None

   # First grab the objects without a specific backend in _import_structure
   _lowercase              =[]
   while not lines[line_index].startswith('if TYPE_CHECKING') and find_backend(lines[line_index]) is None:
      _lowercase              =lines[line_index]
      # If we have everything on a single line, let's deal with it.
      if _re_one_line_import_struct.search(_snake_case):
         _lowercase              =_re_one_line_import_struct.search(_snake_case).groups()[0]
         _lowercase              =re.findall(R'\[([^\]]+)\]'      ,    _snake_case)
         for imp in imports:
            objects.extend([obj[1:-1] for obj in imp.split(', ')])
         line_index += 1
         continue
      _lowercase              =_re_import_struct_key_value.search(_snake_case)
      if single_line_import_search is not None:
         _lowercase              =[obj[1:-1] for obj in single_line_import_search.groups()[0].split(', ') if len(_snake_case) > 0]
         objects.extend(_snake_case)
      elif line.startswith(' ' * 8 + '"'):
         objects.append(line[9:-3])
      line_index += 1

   _lowercase              ={'none': objects}
   # Let's continue with backend-specific objects in _import_structure
   while not lines[line_index].startswith('if TYPE_CHECKING'):
      # If the line is an if not is_backend_available, we grab all objects associated.
      _lowercase              =find_backend(lines[line_index])
      # Check if the backend declaration is inside a try block:
      if _re_try.search(lines[line_index - 1]) is None:
         _lowercase              =None

      if backend is not None:
         line_index += 1

         # Scroll until we hit the else block of try-except-else
         while _re_else.search(lines[line_index]) is None:
            line_index += 1

         line_index += 1

         _lowercase              =[]
         # Until we unindent, add backend objects to the list
         while len(lines[line_index]) <= 1 or lines[line_index].startswith(' ' * 4):
            _lowercase              =lines[line_index]
            if _re_import_struct_add_one.search(_snake_case) is not None:
               objects.append(_re_import_struct_add_one.search(_snake_case).groups()[0])
            elif _re_import_struct_add_many.search(_snake_case) is not None:
               _lowercase              =_re_import_struct_add_many.search(_snake_case).groups()[0].split(', ')
               _lowercase              =[obj[1:-1] for obj in imports if len(_snake_case) > 0]
               objects.extend(_snake_case)
            elif _re_between_brackets.search(_snake_case) is not None:
               _lowercase              =_re_between_brackets.search(_snake_case).groups()[0].split(', ')
               _lowercase              =[obj[1:-1] for obj in imports if len(_snake_case) > 0]
               objects.extend(_snake_case)
            elif _re_quote_object.search(_snake_case) is not None:
               objects.append(_re_quote_object.search(_snake_case).groups()[0])
            elif line.startswith(' ' * 8 + '"'):
               objects.append(line[9:-3])
            elif line.startswith(' ' * 12 + '"'):
               objects.append(line[13:-3])
            line_index += 1

         _lowercase              =objects
      else:
         line_index += 1

    # At this stage we are in the TYPE_CHECKING part, first grab the objects without a specific backend
   _lowercase              =[]
   while (
       line_index < len(_snake_case)
       and find_backend(lines[line_index]) is None
       and not lines[line_index].startswith('else')
   ):
      _lowercase              =lines[line_index]
      _lowercase              =_re_import.search(_snake_case)
      if single_line_import_search is not None:
         objects.extend(single_line_import_search.groups()[0].split(', '))
      elif line.startswith(' ' * 8):
         objects.append(line[8:-2])
      line_index += 1

   _lowercase              ={'none': objects}
   # Let's continue with backend-specific objects
   while line_index < len(_snake_case):
      # If the line is an if is_backend_available, we grab all objects associated.
      _lowercase              =find_backend(lines[line_index])
      # Check if the backend declaration is inside a try block:
      if _re_try.search(lines[line_index - 1]) is None:
         _lowercase              =None

      if backend is not None:
         line_index += 1

         # Scroll until we hit the else block of try-except-else
         while _re_else.search(lines[line_index]) is None:
            line_index += 1

         line_index += 1

         _lowercase              =[]
         # Until we unindent, add backend objects to the list
         while len(lines[line_index]) <= 1 or lines[line_index].startswith(' ' * 8):
            _lowercase              =lines[line_index]
            _lowercase              =_re_import.search(_snake_case)
            if single_line_import_search is not None:
               objects.extend(single_line_import_search.groups()[0].split(', '))
            elif line.startswith(' ' * 12):
               objects.append(line[12:-2])
            line_index += 1

         _lowercase              =objects
      else:
         line_index += 1

   return import_dict_objects, type_hint_objects





def        _snake_case   (_snake_case				:       Dict      ,    _snake_case				:       List[str])					->	List[Any]:

   def find_duplicates(_snake_case				:       Any):
      return [k for k, v in collections.Counter(_snake_case).items() if v > 1]

   if list(import_dict_objects.keys()) != list(type_hint_objects.keys()):
      return ["Both sides of the init do not have the same backends!"]

   _lowercase              =[]
   for key in import_dict_objects.keys():
      _lowercase              =find_duplicates(import_dict_objects[key])
      if duplicate_imports:
         errors.append(f'''Duplicate _import_structure definitions for: {duplicate_imports}''')
      _lowercase              =find_duplicates(type_hint_objects[key])
      if duplicate_type_hints:
         errors.append(f'''Duplicate TYPE_CHECKING objects for: {duplicate_type_hints}''')

      if sorted(set(import_dict_objects[key])) != sorted(set(type_hint_objects[key])):
         _lowercase              ='base imports' if key == 'none' else f'''{key} backend'''
         errors.append(f'''Differences for {name}:''')
         for a in type_hint_objects[key]:
            if a not in import_dict_objects[key]:
               errors.append(f'''  {a} in TYPE_HINT but not in _import_structure.''')
         for a in import_dict_objects[key]:
            if a not in type_hint_objects[key]:
               errors.append(f'''  {a} in _import_structure but not in TYPE_HINT.''')
   return errors





def        _snake_case   ()					->	Dict:
   _lowercase              =[]
   for root, _, files in os.walk(_snake_case):
      if "__init__.py" in files:
         _lowercase              =os.path.join(_snake_case      ,    '__init__.py')
         _lowercase              =parse_init(_snake_case)
         if objects is not None:
            _lowercase              =analyze_results(*_snake_case)
            if len(_snake_case) > 0:
               _lowercase              =f'''Problem in {fname}, both halves do not define the same objects.\n{errors[0]}'''
               failures.append('\n'.join(_snake_case))
   if len(_snake_case) > 0:
      raise ValueError('\n\n'.join(_snake_case))





def        _snake_case   ()					->	int:
   _lowercase              =[]
   for path, directories, files in os.walk(_snake_case):
      for folder in directories:
         # Ignore private modules
         if folder.startswith('_'):
            directories.remove(_snake_case)
            continue
         # Ignore leftovers from branches (empty folders apart from pycache)
         if len(list((Path(_snake_case) / folder).glob('*.py'))) == 0:
            continue
         _lowercase              =str((Path(_snake_case) / folder).relative_to(_snake_case))
         _lowercase              =short_path.replace(os.path.sep      ,    '.')
         submodules.append(_snake_case)
      for fname in files:
         if fname == "__init__.py":
            continue
         _lowercase              =str((Path(_snake_case) / fname).relative_to(_snake_case))
         _lowercase              =short_path.replace('.py'      ,    '').replace(os.path.sep      ,    '.')
         if len(submodule.split('.')) == 1:
            submodules.append(_snake_case)
   return submodules


_SCREAMING_SNAKE_CASE   =  [
    "convert_pytorch_checkpoint_to_tf2",
    "modeling_flax_pytorch_utils",
    "models.esm.openfold_utils",
]





def        _snake_case   ()					->	List[Any]:
   # This is to make sure the transformers module imported is the one in the repo.
   from transformers.utils import direct_transformers_import

   _lowercase              =direct_transformers_import(_snake_case)

   _lowercase              =set(transformers._import_structure.keys())
   # This contains all the base keys of the _import_structure object defined in the init, but if the user is missing
   # some optional dependencies, they may not have all of them. Thus we read the init to read all additions and
   # (potentiall re-) add them.
   with open(os.path.join(_snake_case      ,    '__init__.py')      ,    'r') as f:
      _lowercase              =f.read()
   import_structure_keys.update(set(re.findall(R'import_structure\[\"([^\"]*)\"\]'      ,    _snake_case)))

   _lowercase              =[
       module
       for module in get_transformers_submodules()
       if module not in IGNORE_SUBMODULES and module not in import_structure_keys
   ]

   if len(_snake_case) > 0:
      _lowercase              ='\n'.join(f'''- {module}''' for module in module_not_registered)
      raise ValueError(
          'The following submodules are not properly registed in the main init of Transformers:\n'
          f'''{list_of_modules}\n'''
          'Make sure they appear somewhere in the keys of `_import_structure` with an empty list as value.')


if __name__ == "__main__":
    check_all_inits()
    check_submodules()










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


_SCREAMING_SNAKE_CASE   =  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 					SCREAMING_SNAKE_CASE_  (       _a ):
       """simple docstring"""


       def __init__(				self	:Dict,      *snake_case	:int,      snake_case	:Any=None,      snake_case	:str=None,      snake_case	:int=None,      **snake_case	:List[str]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Dict,      snake_case	:List[str]=None):



          """simple docstring"""

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

          _lowercase              =self._remove_unused_columns(snake_case,      description='Calibration')

          return DataLoader(
              snake_case,      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=snake_case,      )





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Union[str, Any]=None):



          """simple docstring"""

          _lowercase              =self.train_dataset if calib_dataset is None else calib_dataset
          _lowercase              =self.get_calib_dataloader(snake_case)

          _lowercase              =self.model
          quant_trainer.configure_model(snake_case,      self.quant_trainer_args,      calib=snake_case)
          model.eval()
          quant_trainer.enable_calibration(snake_case)

          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(snake_case):
             # Prediction step
             _lowercase       ,       _lowercase       ,       _lowercase              =self.prediction_step(snake_case,      snake_case,      prediction_loss_only=snake_case)
             if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                break

          quant_trainer.finish_calibration(snake_case,      self.quant_trainer_args)
          _lowercase              =model





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:Union[str, Any]=None,      snake_case	:Optional[int]=None,      snake_case	:Union[str, Any]=None,      snake_case	:str = "eval"):



          """simple docstring"""

          _lowercase              =self.eval_dataset if eval_dataset is None else eval_dataset
          _lowercase              =self.get_eval_dataloader(snake_case)
          _lowercase              =self.eval_examples if eval_examples is None else eval_examples

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

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

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

             self.log(snake_case)
          else:
             _lowercase              ={}

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

          _lowercase              =self.callback_handler.on_evaluate(self.args,      self.state,      self.control,      snake_case)
          return metrics





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:str,      snake_case	:Optional[Any],      snake_case	:Any=None,      snake_case	:str = "test"):



          """simple docstring"""

          _lowercase              =self.get_test_dataloader(snake_case)

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

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

          _lowercase              =self.post_process_function(snake_case,      snake_case,      output.predictions,      'predict')
          _lowercase              =self.compute_metrics(snake_case)

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

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





       def 				UpperCamelCase__						(				self	:List[str],      snake_case	:Dict="./"):



          """simple docstring"""

          _lowercase              =self.eval_dataset
          _lowercase              =self.get_eval_dataloader(snake_case)

          _lowercase              =next(iter(snake_case))

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

          # convert to tuple
          _lowercase              =tuple(v.to(snake_case) for k, v in batch.items())

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

          _lowercase              =True

          _lowercase              =self.model.to(snake_case)

          model.eval()
          model.float()

          _lowercase              =model.module if hasattr(snake_case,      'module') else model
          quant_trainer.configure_model(snake_case,      self.quant_trainer_args)

          _lowercase              =os.path.join(snake_case,      'model.onnx')
          logger.info(f'''exporting model to {output_model_file}''')

          _lowercase              ={0: 'batch_size', 1: 'seq_len'}

          torch.onnx.export(
              snake_case,      snake_case,      snake_case,      export_params=snake_case,      opset_version=13,      do_constant_folding=snake_case,      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=snake_case,      )
          logger.info('onnx export finished')


