

"""simple docstring"""
import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import List

import timm
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from torch import Tensor

from transformers import AutoImageProcessor, ResNetConfig, ResNetForImageClassification
from transformers.utils import logging


logging.set_verbosity_info()
_A							=  logging.get_logger()

@dataclass
class       lowerCamelCase	:


  '''simple docstring'''







  SCREAMING_SNAKE_CASE					      =  42
  SCREAMING_SNAKE_CASE					      =  field(default_factory=lowerCAmelCase__   )
  SCREAMING_SNAKE_CASE					      =  field(default_factory=lowerCAmelCase__   )



  def  _a      (self							,       _lowerCamelCase							,       _lowerCamelCase							,       _lowerCamelCase	):


        """simple docstring"""

        UpperCAmelCase__			:							str							=    len(list(m.modules()	)	) == 1 or isinstance(_lowerCamelCase							,       nn.Convad	) or isinstance(_lowerCamelCase							,       nn.BatchNormad	)
        if has_not_submodules:
              self.traced.append(_lowerCamelCase	)



  def __call__(self							,       _lowerCamelCase	):


        """simple docstring"""

        for m in self.module.modules():
              self.handles.append(m.register_forward_hook(self._forward_hook	)	)
        self.module(_lowerCamelCase	)
        [x.remove() for x in self.handles]
        return self




  @property
  def  _a      (self	):


        """simple docstring"""

        return list(filter(lambda _lowerCamelCase	: len(list(x.state_dict().keys()	)	) > 0							,       self.traced	)	)




@dataclass
class       lowerCamelCase	:


  '''simple docstring'''







  SCREAMING_SNAKE_CASE					      =  42
  SCREAMING_SNAKE_CASE					      =  42
  SCREAMING_SNAKE_CASE					      =  0
  SCREAMING_SNAKE_CASE					      =  field(default_factory=lowerCAmelCase__   )
  SCREAMING_SNAKE_CASE					      =  field(default_factory=lowerCAmelCase__   )



  def __call__(self							,       _lowerCamelCase	):


        """simple docstring"""

        UpperCAmelCase__			:							int							=    Tracker(self.dest	)(_lowerCamelCase	).parametrized
        UpperCAmelCase__			:							int							=    Tracker(self.src	)(_lowerCamelCase	).parametrized

        UpperCAmelCase__			:							Optional[int]							=    list(filter(lambda _lowerCamelCase	: type(_lowerCamelCase	) not in self.src_skip							,       _lowerCamelCase	)	)
        UpperCAmelCase__			:							Optional[Any]							=    list(filter(lambda _lowerCamelCase	: type(_lowerCamelCase	) not in self.dest_skip							,       _lowerCamelCase	)	)

        if len(_lowerCamelCase	) != len(_lowerCamelCase	):
              raise Exception(
                  F"""Numbers of operations are different. Source module has {len(_lowerCamelCase	)} operations while"""
                  F""" destination module has {len(_lowerCamelCase	)}."""	)

        for dest_m, src_m in zip(_lowerCamelCase							,       _lowerCamelCase	):
              dest_m.load_state_dict(src_m.state_dict()	)
              if self.verbose == 1:
                    print(F"""Transfered from={src_m} to={dest_m}"""	)




def 					a__						(						lowerCAmelCase			,	lowerCAmelCase			,	lowerCAmelCase			,	lowerCAmelCase = True    )	-> Any:
      print(F"""Converting {name}..."""    )
      with torch.no_grad():
            UpperCAmelCase__			:							Optional[int]							=    timm.create_model(lowerCAmelCase			,	pretrained=lowerCAmelCase    ).eval()
            UpperCAmelCase__			:							str							=    ResNetForImageClassification(lowerCAmelCase    ).eval()
            UpperCAmelCase__			:							List[str]							=    ModuleTransfer(src=lowerCAmelCase			,	dest=lowerCAmelCase    )
            UpperCAmelCase__			:							Any							=    torch.randn((1, 3, 2_24, 2_24)    )
            module_transfer(lowerCAmelCase    )

      assert torch.allclose(from_model(lowerCAmelCase    )			,	our_model(lowerCAmelCase    ).logits    ), "The model logits don't match the original one."

      UpperCAmelCase__			:							List[str]							=    F"""resnet{'-'.join(name.split('resnet'    )    )}"""
      print(lowerCAmelCase    )

      if push_to_hub:
            our_model.push_to_hub(
                repo_path_or_name=save_directory / checkpoint_name			,	commit_message="""Add model"""			,	use_temp_dir=lowerCAmelCase			,	)

            # we can use the convnext one
            UpperCAmelCase__			:							Dict							=    AutoImageProcessor.from_pretrained("""facebook/convnext-base-224-22k-1k"""    )
            image_processor.push_to_hub(
                repo_path_or_name=save_directory / checkpoint_name			,	commit_message="""Add image processor"""			,	use_temp_dir=lowerCAmelCase			,	)

            print(F"""Pushed {checkpoint_name}"""    )




def 					a__						(						lowerCAmelCase			,	lowerCAmelCase = None			,	lowerCAmelCase = True    )	-> Dict:
      UpperCAmelCase__			:							Any							=    """imagenet-1k-id2label.json"""
      UpperCAmelCase__			:							Any							=    10_00
      UpperCAmelCase__			:							List[str]							=    (1, num_labels)

      UpperCAmelCase__			:							Union[str, Any]							=    """huggingface/label-files"""
      UpperCAmelCase__			:							int							=    num_labels
      UpperCAmelCase__			:							List[Any]							=    json.load(open(hf_hub_download(lowerCAmelCase			,	lowerCAmelCase			,	repo_type="""dataset"""    )			,	"""r"""    )    )
      UpperCAmelCase__			:							Optional[Any]							=    {int(lowerCAmelCase    ): v for k, v in idalabel.items()}

      UpperCAmelCase__			:							str							=    idalabel
      UpperCAmelCase__			:							Tuple							=    {v: k for k, v in idalabel.items()}

      UpperCAmelCase__			:							Dict							=    partial(lowerCAmelCase			,	num_labels=lowerCAmelCase			,	idalabel=lowerCAmelCase			,	labelaid=lowerCAmelCase    )

      UpperCAmelCase__			:							Dict							=    {
          """resnet18""": ImageNetPreTrainedConfig(
              depths=[2, 2, 2, 2]			,	hidden_sizes=[64, 1_28, 2_56, 5_12]			,	layer_type="""basic"""    ),
          """resnet26""": ImageNetPreTrainedConfig(
              depths=[2, 2, 2, 2]			,	hidden_sizes=[2_56, 5_12, 10_24, 20_48]			,	layer_type="""bottleneck"""    ),
          """resnet34""": ImageNetPreTrainedConfig(
              depths=[3, 4, 6, 3]			,	hidden_sizes=[64, 1_28, 2_56, 5_12]			,	layer_type="""basic"""    ),
          """resnet50""": ImageNetPreTrainedConfig(
              depths=[3, 4, 6, 3]			,	hidden_sizes=[2_56, 5_12, 10_24, 20_48]			,	layer_type="""bottleneck"""    ),
          """resnet101""": ImageNetPreTrainedConfig(
              depths=[3, 4, 23, 3]			,	hidden_sizes=[2_56, 5_12, 10_24, 20_48]			,	layer_type="""bottleneck"""    ),
          """resnet152""": ImageNetPreTrainedConfig(
              depths=[3, 8, 36, 3]			,	hidden_sizes=[2_56, 5_12, 10_24, 20_48]			,	layer_type="""bottleneck"""    ),
      }

      if model_name:
            convert_weight_and_push(lowerCAmelCase			,	names_to_config[model_name]			,	lowerCAmelCase			,	lowerCAmelCase    )
      else:
            for model_name, config in names_to_config.items():
                  convert_weight_and_push(lowerCAmelCase			,	lowerCAmelCase			,	lowerCAmelCase			,	lowerCAmelCase    )
      return config, expected_shape


if __name__ == "__main__":
    _A							=  argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--model_name""",
        default=None,
        type=str,
        help=(
            """The name of the model you wish to convert, it must be one of the supported resnet* architecture,"""
            """ currently: resnet18,26,34,50,101,152. If `None`, all of them will the converted."""
        ),
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""",
        default=None,
        type=Path,
        required=True,
        help="""Path to the output PyTorch model directory.""",
    )
    parser.add_argument(
        """--push_to_hub""",
        default=True,
        type=bool,
        required=False,
        help="""If True, push model and image processor to the hub.""",
    )

    _A							=  parser.parse_args()
    _A							=  args.pytorch_dump_folder_path
    pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
    convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)




"""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							=  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	(       lowerCAmelCase__   ):


  '''simple docstring'''







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


        """simple docstring"""

        super().__init__(*_lowerCamelCase							,       **_lowerCamelCase	)
        UpperCAmelCase__			:							Any							=    eval_examples
        UpperCAmelCase__			:							Union[str, Any]							=    post_process_function
        UpperCAmelCase__			:							int							=    quant_trainer_args
        UpperCAmelCase__			:							Dict							=    128  # default number of calibration samples



  def  _a      (self							,       _lowerCamelCase=None	):


        """simple docstring"""

        if calib_dataset is None and self.calib_dataset is None:
              raise ValueError("""Trainer: calibration requires an calib_dataset."""	)
        UpperCAmelCase__			:							List[str]							=    calib_dataset if calib_dataset is not None else self.calib_dataset

        UpperCAmelCase__			:							Optional[int]							=    self._remove_unused_columns(_lowerCamelCase							,       description="""Calibration"""	)

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



  def  _a      (self							,       _lowerCamelCase=None	):


        """simple docstring"""

        UpperCAmelCase__			:							int							=    self.train_dataset if calib_dataset is None else calib_dataset
        UpperCAmelCase__			:							List[str]							=    self.get_calib_dataloader(_lowerCamelCase	)

        UpperCAmelCase__			:							List[str]							=    self.model
        quant_trainer.configure_model(_lowerCamelCase							,       self.quant_trainer_args							,       calib=_lowerCamelCase	)
        model.eval()
        quant_trainer.enable_calibration(_lowerCamelCase	)

        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(_lowerCamelCase	):
              # Prediction step
              UpperCAmelCase__					,		UpperCAmelCase__					,		UpperCAmelCase__			:							List[str]							=    self.prediction_step(_lowerCamelCase							,       _lowerCamelCase							,       prediction_loss_only=_lowerCamelCase	)
              if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                    break

        quant_trainer.finish_calibration(_lowerCamelCase							,       self.quant_trainer_args	)
        UpperCAmelCase__			:							Tuple							=    model



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


        """simple docstring"""

        UpperCAmelCase__			:							Tuple							=    self.eval_dataset if eval_dataset is None else eval_dataset
        UpperCAmelCase__			:							List[Any]							=    self.get_eval_dataloader(_lowerCamelCase	)
        UpperCAmelCase__			:							Tuple							=    self.eval_examples if eval_examples is None else eval_examples

        # Temporarily disable metric computation, we will do it in the loop here.
        UpperCAmelCase__			:							Optional[Any]							=    self.compute_metrics
        UpperCAmelCase__			:							int							=    None
        UpperCAmelCase__			:							int							=    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
              UpperCAmelCase__			:							str							=    eval_loop(
                  _lowerCamelCase							,       description="""Evaluation"""							,       prediction_loss_only=True if compute_metrics is None else None							,       ignore_keys=_lowerCamelCase							,       )
        finally:
              UpperCAmelCase__			:							List[str]							=    compute_metrics

        if self.post_process_function is not None and self.compute_metrics is not None:
              UpperCAmelCase__			:							int							=    self.post_process_function(_lowerCamelCase							,       _lowerCamelCase							,       output.predictions	)
              UpperCAmelCase__			:							List[Any]							=    self.compute_metrics(_lowerCamelCase	)

              # Prefix all keys with metric_key_prefix + '_'
              for key in list(metrics.keys()	):
                    if not key.startswith(F"""{metric_key_prefix}_"""	):
                          UpperCAmelCase__			:							Optional[int]							=    metrics.pop(_lowerCamelCase	)

              self.log(_lowerCamelCase	)
        else:
              UpperCAmelCase__			:							List[Any]							=    {}

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

        UpperCAmelCase__			:							Optional[int]							=    self.callback_handler.on_evaluate(self.args							,       self.state							,       self.control							,       _lowerCamelCase	)
        return metrics



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


        """simple docstring"""

        UpperCAmelCase__			:							Union[str, Any]							=    self.get_test_dataloader(_lowerCamelCase	)

        # Temporarily disable metric computation, we will do it in the loop here.
        UpperCAmelCase__			:							Union[str, Any]							=    self.compute_metrics
        UpperCAmelCase__			:							Optional[Any]							=    None
        UpperCAmelCase__			:							Any							=    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
              UpperCAmelCase__			:							Optional[Any]							=    eval_loop(
                  _lowerCamelCase							,       description="""Prediction"""							,       prediction_loss_only=True if compute_metrics is None else None							,       ignore_keys=_lowerCamelCase							,       )
        finally:
              UpperCAmelCase__			:							Dict							=    compute_metrics

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

        UpperCAmelCase__			:							Dict							=    self.post_process_function(_lowerCamelCase							,       _lowerCamelCase							,       output.predictions							,       """predict"""	)
        UpperCAmelCase__			:							Any							=    self.compute_metrics(_lowerCamelCase	)

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()	):
              if not key.startswith(F"""{metric_key_prefix}_"""	):
                    UpperCAmelCase__			:							Any							=    metrics.pop(_lowerCamelCase	)

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




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


        """simple docstring"""

        UpperCAmelCase__			:							Tuple							=    self.eval_dataset
        UpperCAmelCase__			:							int							=    self.get_eval_dataloader(_lowerCamelCase	)

        UpperCAmelCase__			:							Union[str, Any]							=    next(iter(_lowerCamelCase	)	)

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

        # convert to tuple
        UpperCAmelCase__			:							List[Any]							=    tuple(v.to(_lowerCamelCase	) for k, v in batch.items()	)

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

        UpperCAmelCase__			:							Optional[Any]							=    True

        UpperCAmelCase__			:							Tuple							=    self.model.to(_lowerCamelCase	)

        model.eval()
        model.float()

        UpperCAmelCase__			:							str							=    model.module if hasattr(_lowerCamelCase							,       """module"""	) else model
        quant_trainer.configure_model(_lowerCamelCase							,       self.quant_trainer_args	)

        UpperCAmelCase__			:							List[str]							=    os.path.join(_lowerCamelCase							,       """model.onnx"""	)
        logger.info(F"""exporting model to {output_model_file}"""	)

        UpperCAmelCase__			:							Optional[Any]							=    {0: """batch_size""", 1: """seq_len"""}

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

