






"""simple docstring"""


import inspect
import os
from argparse import ArgumentParser, Namespace
from importlib import import_module

import huggingface_hub
import numpy as np
from packaging import version

from .. import (
    FEATURE_EXTRACTOR_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutoProcessor,
    AutoTokenizer,
    is_datasets_available,
    is_tf_available,
    is_torch_available,
)
from ..utils import TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, logging
from . import BaseTransformersCLICommand


if is_tf_available():
       import tensorflow as tf

       tf.config.experimental.enable_tensor_float_aa_execution(False)

if is_torch_available():
       import torch

if is_datasets_available():
       from datasets import load_dataset


lowercase_    					=					5e-5  # larger error tolerance than in our internal tests, to avoid flaky user-facing errors



def  lowerCAmelCase    (   __UpperCamelCase			):




  """simple docstring"""







  return PTtoTFCommand(
      args.model_name		,	args.local_dir		,	args.max_error		,	args.new_weights		,	args.no_pr		,	args.push		,	args.extra_commit_description		,	args.override_model_class		,	)





class  snake_case		(							_lowerCAmelCase						):


   '''simple docstring'''

   @staticmethod
   def     _SCREAMING_SNAKE_CASE    (     _lowerCamelCase			:					ArgumentParser  ):

     '''simple docstring'''
     __A						      =			parser.add_parser(
         '''pt-to-tf''',	help=(
             '''CLI tool to run convert a transformers model from a PyTorch checkpoint to a TensorFlow checkpoint.'''
             ''' Can also be used to validate existing weights without opening PRs, with --no-pr.'''
         ),	)
     train_parser.add_argument(
         '''--model-name''',	type=_lowerCamelCase,	required=_lowerCamelCase,	help='''The model name, including owner/organization, as seen on the hub.''',	)
     train_parser.add_argument(
         '''--local-dir''',	type=_lowerCamelCase,	default='''''',	help='''Optional local directory of the model repository. Defaults to /tmp/{model_name}''',	)
     train_parser.add_argument(
         '''--max-error''',	type=_lowerCamelCase,	default=_lowerCamelCase,	help=(
             f'Maximum error tolerance. Defaults to {MAX_ERROR}. This flag should be avoided, use at your own risk.'
         ),	)
     train_parser.add_argument(
         '''--new-weights''',	action='''store_true''',	help='''Optional flag to create new TensorFlow weights, even if they already exist.''',	)
     train_parser.add_argument(
         '''--no-pr''',	action='''store_true''',	help='''Optional flag to NOT open a PR with converted weights.'''  )
     train_parser.add_argument(
         '''--push''',	action='''store_true''',	help='''Optional flag to push the weights directly to `main` (requires permissions)''',	)
     train_parser.add_argument(
         '''--extra-commit-description''',	type=_lowerCamelCase,	default='''''',	help='''Optional additional commit description to use when opening a PR (e.g. to tag the owner).''',	)
     train_parser.add_argument(
         '''--override-model-class''',	type=_lowerCamelCase,	default=_lowerCamelCase,	help='''If you think you know better than the auto-detector, you can specify the model class here. '''
         '''Can be either an AutoModel class or a specific model class like BertForSequenceClassification.''',	)
     train_parser.set_defaults(func=_lowerCamelCase  )


   @staticmethod
   def     _SCREAMING_SNAKE_CASE    (     _lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					Optional[int]  ):

     '''simple docstring'''
     __A						      =			set(pt_outputs.keys()  )
     __A						      =			set(tf_outputs.keys()  )
     if pt_out_attrs != tf_out_attrs:
       raise ValueError(
           f'The model outputs have different attributes, aborting. (Pytorch: {pt_out_attrs}, TensorFlow:'
           f' {tf_out_attrs})'  )

     # 2. For each output attribute, computes the difference
     def _find_pt_tf_differences(_lowerCamelCase			:					Dict,	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Union[str, Any],	_lowerCamelCase			:					Union[str, Any]=""  ):
       # If the current attribute is a tensor, it is a leaf and we make the comparison. Otherwise, we will dig in
       # recursivelly, keeping the name of the attribute.
       if isinstance(_lowerCamelCase,	torch.Tensor  ):
         __A						      =			np.max(np.abs(pt_out.numpy() - tf_out.numpy()  )  )
         __A						      =			tensor_difference
       else:
         __A						      =			attr_name
         for i, pt_item in enumerate(_lowerCamelCase  ):
           # If it is a named attribute, we keep the name. Otherwise, just its index.
           if isinstance(_lowerCamelCase,	_lowerCamelCase  ):
             __A						      =			root_name + pt_item
             __A						      =			tf_out[pt_item]
             __A						      =			pt_out[pt_item]
           else:
             __A						      =			root_name + f'[{i}]'
             __A						      =			tf_out[i]
           __A						      =			_find_pt_tf_differences(_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase  )

       return differences

     return _find_pt_tf_differences(_lowerCamelCase,	_lowerCamelCase,	{}  )


   def __init__(     self			:					List[Any],	_lowerCamelCase			:					str,	_lowerCamelCase			:					str,	_lowerCamelCase			:					float,	_lowerCamelCase			:					bool,	_lowerCamelCase			:					bool,	_lowerCamelCase			:					bool,	_lowerCamelCase			:					str,	_lowerCamelCase			:					str,	*_lowerCamelCase			:					Dict,	):

     '''simple docstring'''
     __A						      =			logging.get_logger('''transformers-cli/pt_to_tf'''  )
     __A						      =			model_name
     __A						      =			local_dir if local_dir else os.path.join('''/tmp''',	_lowerCamelCase  )
     __A						      =			max_error
     __A						      =			new_weights
     __A						      =			no_pr
     __A						      =			push
     __A						      =			extra_commit_description
     __A						      =			override_model_class


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[Any],	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Dict  ):

     '''simple docstring'''

     def _get_audio_input():
       __A						      =			load_dataset('''hf-internal-testing/librispeech_asr_dummy''',	'''clean''',	split='''validation'''  )
       __A						      =			ds.sort('''id'''  ).select(range(2  )  )[:2]['''audio''']
       __A						      =			[x['''array'''] for x in speech_samples]
       return raw_samples

     __A						      =			type(pt_model.config  )
     if model_config_class in PROCESSOR_MAPPING:
       __A						      =			AutoProcessor.from_pretrained(self._local_dir  )
       if model_config_class in TOKENIZER_MAPPING and processor.tokenizer.pad_token is None:
         __A						      =			processor.tokenizer.eos_token
     elif model_config_class in IMAGE_PROCESSOR_MAPPING:
       __A						      =			AutoImageProcessor.from_pretrained(self._local_dir  )
     elif model_config_class in FEATURE_EXTRACTOR_MAPPING:
       __A						      =			AutoFeatureExtractor.from_pretrained(self._local_dir  )
     elif model_config_class in TOKENIZER_MAPPING:
       __A						      =			AutoTokenizer.from_pretrained(self._local_dir  )
       if processor.pad_token is None:
         __A						      =			processor.eos_token
     else:
       raise ValueError(f'Unknown data processing type (model config type: {model_config_class})'  )

     __A						      =			set(inspect.signature(pt_model.forward  ).parameters.keys()  )
     __A						      =			{}
     if "input_ids" in model_forward_signature:
       processor_inputs.update(
           {
               '''text''': ['''Hi there!''', '''I am a batch with more than one row and different input lengths.'''],
               '''padding''': True,
               '''truncation''': True,
           }  )
     if "pixel_values" in model_forward_signature:
       __A						      =			load_dataset('''cifar10''',	'''plain_text''',	split='''test'''  )[:2]['''img''']
       processor_inputs.update({'''images''': sample_images}  )
     if "input_features" in model_forward_signature:
       __A						      =			inspect.signature(processor.feature_extractor  ).parameters
       # Pad to the largest input length by default but take feature extractor default
       # padding value if it exists e.g. "max_length" and is not False or None
       if "padding" in feature_extractor_signature:
         __A						      =			feature_extractor_signature['''padding'''].default
         if default_strategy is not False and default_strategy is not None:
           __A						      =			default_strategy
         else:
           __A						      =			True
       else:
         __A						      =			True
       processor_inputs.update({'''audio''': _get_audio_input(), '''padding''': padding_strategy}  )
     if "input_values" in model_forward_signature:  # Wav2Vec2 audio input
       processor_inputs.update({'''audio''': _get_audio_input(), '''padding''': True}  )
     __A						      =			processor(**_lowerCamelCase,	return_tensors='''pt'''  )
     __A						      =			processor(**_lowerCamelCase,	return_tensors='''tf'''  )

     # Extra input requirements, in addition to the input modality
     if (
         config.is_encoder_decoder
         or (hasattr(_lowerCamelCase,	'''encoder'''  ) and hasattr(_lowerCamelCase,	'''decoder'''  ))
         or "decoder_input_ids" in tf_dummy_inputs
     ):
       __A						      =			np.asarray([[1], [1]],	dtype=_lowerCamelCase  ) * (pt_model.config.decoder_start_token_id or 0)
       pt_input.update({'''decoder_input_ids''': torch.tensor(_lowerCamelCase  )}  )
       tf_input.update({'''decoder_input_ids''': tf.convert_to_tensor(_lowerCamelCase  )}  )

     return pt_input, tf_input






   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     # hub version 0.9.0 introduced the possibility of programmatically opening PRs with normal write tokens.
     if version.parse(huggingface_hub.__version__  ) < version.parse('''0.9.0'''  ):
       raise ImportError(
           '''The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub'''
           ''' installation.'''  )
     else:
       from huggingface_hub import Repository, create_commit
       from huggingface_hub._commit_api import CommitOperationAdd

     # Fetch remote data
     __A						      =			Repository(local_dir=self._local_dir,	clone_from=self._model_name  )

     # Load config and get the appropriate architecture -- the latter is needed to convert the head's weights
     __A						      =			AutoConfig.from_pretrained(self._local_dir  )
     __A						      =			config.architectures
     if self._override_model_class is not None:
       if self._override_model_class.startswith('''TF'''  ):
         __A						      =			[self._override_model_class[2:]]
       else:
         __A						      =			[self._override_model_class]
       try:
         __A						      =			getattr(import_module('''transformers'''  ),	architectures[0]  )
       except AttributeError:
         raise ValueError(f'Model class {self._override_model_class} not found in transformers.'  )
       try:
         __A						      =			getattr(import_module('''transformers'''  ),	'''TF''' + architectures[0]  )
       except AttributeError:
         raise ValueError(f'TF model class TF{self._override_model_class} not found in transformers.'  )
     elif architectures is None:  # No architecture defined -- use auto classes
       __A						      =			getattr(import_module('''transformers'''  ),	'''AutoModel'''  )
       __A						      =			getattr(import_module('''transformers'''  ),	'''TFAutoModel'''  )
       self._logger.warning('''No detected architecture, using AutoModel/TFAutoModel'''  )
     else:  # Architecture defined -- use it
       if len(_lowerCamelCase  ) > 1:
         raise ValueError(f'More than one architecture was found, aborting. (architectures = {architectures})'  )
       self._logger.warning(f'Detected architecture: {architectures[0]}'  )
       __A						      =			getattr(import_module('''transformers'''  ),	architectures[0]  )
       try:
         __A						      =			getattr(import_module('''transformers'''  ),	'''TF''' + architectures[0]  )
       except AttributeError:
         raise AttributeError(f'The TensorFlow equivalent of {architectures[0]} doesn\'t exist in transformers.'  )

        # Check the TF dummy inputs to see what keys we need in the forward pass
     __A						      =			tf_class.from_config(_lowerCamelCase  )
     __A						      =			tf_from_pt_model.dummy_inputs

     del tf_from_pt_model  # Try to keep only one model in memory at a time

     # Load the model and get some basic inputs
     __A						      =			pt_class.from_pretrained(self._local_dir  )
     pt_model.eval()

     __A      ,   __A						      =			self.get_inputs(_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase  )

     with torch.no_grad():
       __A						      =			pt_model(**_lowerCamelCase,	output_hidden_states=_lowerCamelCase  )
     del pt_model  # will no longer be used, and may have a large memory footprint

     __A						      =			tf_class.from_pretrained(self._local_dir,	from_pt=_lowerCamelCase  )
     __A						      =			tf_from_pt_model(**_lowerCamelCase,	output_hidden_states=_lowerCamelCase,	training=_lowerCamelCase  )

     # Confirms that cross loading PT weights into TF worked.
     __A						      =			self.find_pt_tf_differences(_lowerCamelCase,	_lowerCamelCase  )
     __A						      =			{k: v for k, v in crossload_differences.items() if '''hidden''' not in k}
     __A						      =			{k: v for k, v in crossload_differences.items() if '''hidden''' in k}
     if len(_lowerCamelCase  ) == 0 and architectures is not None:
       raise ValueError(
           f'Something went wrong -- the config file has architectures ({architectures}), but no model head'
           ''' output was found. All outputs start with \'hidden\''''  )
     __A						      =			max(output_differences.values()  ) if output_differences else 0.0
     __A						      =			max(hidden_differences.values()  )
     if max_crossload_output_diff > self._max_error or max_crossload_hidden_diff > self._max_error:
       raise ValueError(
           '''The cross-loaded TensorFlow model has different outputs, something went wrong!\n'''
           + f'\nList of maximum output differences above the threshold ({self._max_error}):\n'
           + '''\n'''.join([f'{k}: {v:.3e}' for k, v in output_differences.items() if v > self._max_error]  )
           + f'\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n'
           + '''\n'''.join([f'{k}: {v:.3e}' for k, v in hidden_differences.items() if v > self._max_error]  )  )

     # Save the weights in a TF format (if needed) and confirms that the results are still good
     __A						      =			os.path.join(self._local_dir,	_lowerCamelCase  )
     __A						      =			os.path.join(self._local_dir,	_lowerCamelCase  )
     if (not os.path.exists(_lowerCamelCase  ) and not os.path.exists(_lowerCamelCase  )) or self._new_weights:
       tf_from_pt_model.save_pretrained(self._local_dir  )
     del tf_from_pt_model  # will no longer be used, and may have a large memory footprint

     __A						      =			tf_class.from_pretrained(self._local_dir  )
     __A						      =			tf_model(**_lowerCamelCase,	output_hidden_states=_lowerCamelCase  )

     __A						      =			self.find_pt_tf_differences(_lowerCamelCase,	_lowerCamelCase  )
     __A						      =			{k: v for k, v in conversion_differences.items() if '''hidden''' not in k}
     __A						      =			{k: v for k, v in conversion_differences.items() if '''hidden''' in k}
     if len(_lowerCamelCase  ) == 0 and architectures is not None:
       raise ValueError(
           f'Something went wrong -- the config file has architectures ({architectures}), but no model head'
           ''' output was found. All outputs start with \'hidden\''''  )
     __A						      =			max(output_differences.values()  ) if output_differences else 0.0
     __A						      =			max(hidden_differences.values()  )
     if max_conversion_output_diff > self._max_error or max_conversion_hidden_diff > self._max_error:
       raise ValueError(
           '''The converted TensorFlow model has different outputs, something went wrong!\n'''
           + f'\nList of maximum output differences above the threshold ({self._max_error}):\n'
           + '''\n'''.join([f'{k}: {v:.3e}' for k, v in output_differences.items() if v > self._max_error]  )
           + f'\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n'
           + '''\n'''.join([f'{k}: {v:.3e}' for k, v in hidden_differences.items() if v > self._max_error]  )  )

     __A						      =			'''Update TF weights''' if self._new_weights else '''Add TF weights'''
     if self._push:
       repo.git_add(auto_lfs_track=_lowerCamelCase  )
       repo.git_commit(_lowerCamelCase  )
       repo.git_push(blocking=_lowerCamelCase  )  # this prints a progress bar with the upload
       self._logger.warning(f'TF weights pushed into {self._model_name}'  )
     elif not self._no_pr:
       self._logger.warning('''Uploading the weights into a new PR...'''  )
       __A						      =			(
           '''Model converted by the [`transformers`\' `pt_to_tf`'''
           ''' CLI](https://github.com/huggingface/transformers/blob/main/src/transformers/commands/pt_to_tf.py). '''
           '''All converted model outputs and hidden layers were validated against its PyTorch counterpart.\n\n'''
           f'Maximum crossload output difference={max_crossload_output_diff:.3e}; '
           f'Maximum crossload hidden layer difference={max_crossload_hidden_diff:.3e};\n'
           f'Maximum conversion output difference={max_conversion_output_diff:.3e}; '
           f'Maximum conversion hidden layer difference={max_conversion_hidden_diff:.3e};\n'
       )
       if self._max_error > MAX_ERROR:
         commit_descrition += (
             f'\n\nCAUTION: The maximum admissible error was manually increased to {self._max_error}!'
         )
       if self._extra_commit_description:
         commit_descrition += "\n\n" + self._extra_commit_description

       # sharded model -> adds all related files (index and .h5 shards)
       if os.path.exists(_lowerCamelCase  ):
         __A						      =			[
             CommitOperationAdd(path_in_repo=_lowerCamelCase,	path_or_fileobj=_lowerCamelCase  )
         ]
         for shard_path in tf.io.gfile.glob(self._local_dir + '''/tf_model-*.h5'''  ):
           operations += [
               CommitOperationAdd(path_in_repo=os.path.basename(_lowerCamelCase  ),	path_or_fileobj=_lowerCamelCase  )
           ]
       else:
         __A						      =			[CommitOperationAdd(path_in_repo=_lowerCamelCase,	path_or_fileobj=_lowerCamelCase  )]

       __A						      =			create_commit(
           repo_id=self._model_name,	operations=_lowerCamelCase,	commit_message=_lowerCamelCase,	commit_description=_lowerCamelCase,	repo_type='''model''',	create_pr=_lowerCamelCase,	).pr_url
       self._logger.warning(f'PR open in {hub_pr_url}'  )












"""simple docstring"""


# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


lowercase_    					=					logging.get_logger(__name__)


if is_torch_available():
       import torch
       from torch import nn
       from torch.utils.data import Dataset

       from transformers import Trainer


       class  snake_case		(							_lowerCAmelCase						):


          '''simple docstring'''

          def __init__(     self			:					List[Any],	_lowerCamelCase			:					int = 1_01  ):

            '''simple docstring'''
            __A						      =			length


          def __len__(     self			:					Optional[int]  ):

            '''simple docstring'''
            return self.length






          def __getitem__(     self			:					List[str],	_lowerCamelCase			:					Union[str, Any]  ):

            '''simple docstring'''
            return i





       class  snake_case		:


          '''simple docstring'''

          def __call__(     self			:					Tuple,	_lowerCamelCase			:					Dict  ):

            '''simple docstring'''
            return {"input_ids": torch.tensor(_lowerCamelCase  ), "labels": torch.tensor(_lowerCamelCase  )}





       class  snake_case		(							nn.Module						):


          '''simple docstring'''

          def __init__(     self			:					List[str]  ):

            '''simple docstring'''
            super().__init__()
            # Add some (unused) params otherwise DDP will complain.
            __A						      =			nn.Linear(1_20,	80  )






          def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[Any],	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Optional[Any]=None  ):

            '''simple docstring'''
            if labels is not None:
              return torch.tensor(0.0,	device=input_ids.device  ), input_ids
            else:
              return input_ids



def  lowerCAmelCase    (   ):




  """simple docstring"""







  __A						      =			HfArgumentParser((TrainingArguments,)			)
  sys.argv += ["--output_dir", "./examples"]
  __A						      =			parser.parse_args_into_dataclasses()[0]

  logger.warning(
      f'Process rank: {training_args.local_rank}, device: {training_args.device}, '
      f'tpu_num_cores: {training_args.tpu_num_cores}'		,	)

  # Essentially, what we want to verify in the distributed case is
  # that we get all samples back, in the right order.
  # (this is crucial for prediction for instance)
  for dataset_length in [1_0_0_1, 2_5_6, 1_5]:
    __A						      =			DummyDataset(__UpperCamelCase			)

    def compute_metrics(__UpperCamelCase			) -> Dict:
      __A						      =			list(range(len(__UpperCamelCase			)			)			)
      __A						      =			p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
      return {"success": success}

    __A						      =			Trainer(
        model=DummyModel()		,	args=__UpperCamelCase		,	data_collator=DummyDataCollator()		,	eval_dataset=__UpperCamelCase		,	compute_metrics=__UpperCamelCase		,	)
    __A						      =			trainer.evaluate()
    logger.info(__UpperCamelCase			)
    if metrics["eval_success"] is not True:
      logger.error(__UpperCamelCase			)
      exit(1			)

    __A						      =			trainer.predict(__UpperCamelCase			)
    logger.info(p.metrics			)
    if p.metrics["test_success"] is not True:
      logger.error(p.metrics			)
      exit(1			)

    __A						      =			2

    __A						      =			trainer.evaluate()
    logger.info(__UpperCamelCase			)
    if metrics["eval_success"] is not True:
      logger.error(__UpperCamelCase			)
      exit(1			)

    __A						      =			trainer.predict(__UpperCamelCase			)
    logger.info(p.metrics			)
    if p.metrics["test_success"] is not True:
      logger.error(p.metrics			)
      exit(1			)

    __A						      =			None

  logger.info('''🔥 All distributed tests successful'''			)



def  lowerCAmelCase    (   __UpperCamelCase			):




  """simple docstring"""







  main()


if __name__ == "__main__":
       main()




