






"""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


_snake_case							:   Dict	  =       5E-5  # larger error tolerance than in our internal tests, to avoid flaky user-facing errors




def     lowerCamelCase_					(	__A	:		Namespace      ):
					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 							__lowercase(							_UpperCamelCase	):






						@staticmethod
						def 		lowerCAmelCase (       _SCREAMING_SNAKE_CASE)      ->	List[Any]:




											"""simple docstring"""





											lowerCamelCase__					=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=_SCREAMING_SNAKE_CASE     ,    required=_SCREAMING_SNAKE_CASE     ,    help='The model name, including owner/organization, as seen on the hub.'     ,    )
											train_parser.add_argument(
											    '--local-dir'     ,    type=_SCREAMING_SNAKE_CASE     ,    default=''     ,    help='Optional local directory of the model repository. Defaults to /tmp/{model_name}'     ,    )
											train_parser.add_argument(
											    '--max-error'     ,    type=_SCREAMING_SNAKE_CASE     ,    default=_SCREAMING_SNAKE_CASE     ,    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=_SCREAMING_SNAKE_CASE     ,    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=_SCREAMING_SNAKE_CASE     ,    default=_SCREAMING_SNAKE_CASE     ,    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=_SCREAMING_SNAKE_CASE)






						@staticmethod
						def 		lowerCAmelCase (       _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Optional[Any]:




											"""simple docstring"""





											lowerCamelCase__					=set(pt_outputs.keys())
											lowerCamelCase__					=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(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE=""):
																# 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(_SCREAMING_SNAKE_CASE     ,    torch.Tensor):
																					lowerCamelCase__					=np.max(np.abs(pt_out.numpy() - tf_out.numpy()))
																					lowerCamelCase__					=tensor_difference
																else:
																					lowerCamelCase__					=attr_name
																					for i, pt_item in enumerate(_SCREAMING_SNAKE_CASE):
																										# If it is a named attribute, we keep the name. Otherwise, just its index.
																										if isinstance(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE):
																															lowerCamelCase__					=root_name + pt_item
																															lowerCamelCase__					=tf_out[pt_item]
																															lowerCamelCase__					=pt_out[pt_item]
																										else:
																															lowerCamelCase__					=root_name + F"[{i}]"
																															lowerCamelCase__					=tf_out[i]
																										lowerCamelCase__					=_find_pt_tf_differences(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)

																return differences

											return _find_pt_tf_differences(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    {})






						def __init__(       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    *_SCREAMING_SNAKE_CASE     ,    )      ->	List[str]:




											"""simple docstring"""





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






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Dict:




											"""simple docstring"""





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

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

											lowerCamelCase__					=set(inspect.signature(pt_model.forward).parameters.keys())
											lowerCamelCase__					={}
											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:
																lowerCamelCase__					=load_dataset('cifar10'     ,    'plain_text'     ,    split='test')[:2]['img']
																processor_inputs.update({'images': sample_images})
											if "input_features" in model_forward_signature:
																lowerCamelCase__					=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:
																					lowerCamelCase__					=feature_extractor_signature['padding'].default
																					if default_strategy is not False and default_strategy is not None:
																										lowerCamelCase__					=default_strategy
																					else:
																										lowerCamelCase__					=True
																else:
																					lowerCamelCase__					=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})
											lowerCamelCase__					=processor(**_SCREAMING_SNAKE_CASE     ,    return_tensors='pt')
											lowerCamelCase__					=processor(**_SCREAMING_SNAKE_CASE     ,    return_tensors='tf')

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

											return pt_input, tf_input






						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											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
											lowerCamelCase__					=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
											lowerCamelCase__					=AutoConfig.from_pretrained(self._local_dir)
											lowerCamelCase__					=config.architectures
											if self._override_model_class is not None:
																if self._override_model_class.startswith('TF'):
																					lowerCamelCase__					=[self._override_model_class[2:]]
																else:
																					lowerCamelCase__					=[self._override_model_class]
																try:
																					lowerCamelCase__					=getattr(import_module('transformers')     ,    architectures[0])
																except AttributeError:
																					raise ValueError(F"Model class {self._override_model_class} not found in transformers.")
																try:
																					lowerCamelCase__					=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
																lowerCamelCase__					=getattr(import_module('transformers')     ,    'AutoModel')
																lowerCamelCase__					=getattr(import_module('transformers')     ,    'TFAutoModel')
																self._logger.warning('No detected architecture, using AutoModel/TFAutoModel')
											else:  # Architecture defined -- use it
																if len(_SCREAMING_SNAKE_CASE) > 1:
																					raise ValueError(F"More than one architecture was found, aborting. (architectures = {architectures})")
																self._logger.warning(F"Detected architecture: {architectures[0]}")
																lowerCamelCase__					=getattr(import_module('transformers')     ,    architectures[0])
																try:
																					lowerCamelCase__					=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
											lowerCamelCase__					=tf_class.from_config(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=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
											lowerCamelCase__					=pt_class.from_pretrained(self._local_dir)
											pt_model.eval()

											lowerCamelCase__					,  lowerCamelCase__					=self.get_inputs(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)

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

											lowerCamelCase__					=tf_class.from_pretrained(self._local_dir     ,    from_pt=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=tf_from_pt_model(**_SCREAMING_SNAKE_CASE     ,    output_hidden_states=_SCREAMING_SNAKE_CASE     ,    training=_SCREAMING_SNAKE_CASE)

											# Confirms that cross loading PT weights into TF worked.
											lowerCamelCase__					=self.find_pt_tf_differences(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											lowerCamelCase__					={k: v for k, v in crossload_differences.items() if 'hidden' not in k}
											lowerCamelCase__					={k: v for k, v in crossload_differences.items() if 'hidden' in k}
											if len(_SCREAMING_SNAKE_CASE) == 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\'')
											lowerCamelCase__					=max(output_differences.values()) if output_differences else 0.0
											lowerCamelCase__					=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
											lowerCamelCase__					=os.path.join(self._local_dir     ,    _SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=os.path.join(self._local_dir     ,    _SCREAMING_SNAKE_CASE)
											if (not os.path.exists(_SCREAMING_SNAKE_CASE) and not os.path.exists(_SCREAMING_SNAKE_CASE)) 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

											lowerCamelCase__					=tf_class.from_pretrained(self._local_dir)
											lowerCamelCase__					=tf_model(**_SCREAMING_SNAKE_CASE     ,    output_hidden_states=_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=self.find_pt_tf_differences(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											lowerCamelCase__					={k: v for k, v in conversion_differences.items() if 'hidden' not in k}
											lowerCamelCase__					={k: v for k, v in conversion_differences.items() if 'hidden' in k}
											if len(_SCREAMING_SNAKE_CASE) == 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\'')
											lowerCamelCase__					=max(output_differences.values()) if output_differences else 0.0
											lowerCamelCase__					=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]))

											lowerCamelCase__					='Update TF weights' if self._new_weights else 'Add TF weights'
											if self._push:
																repo.git_add(auto_lfs_track=_SCREAMING_SNAKE_CASE)
																repo.git_commit(_SCREAMING_SNAKE_CASE)
																repo.git_push(blocking=_SCREAMING_SNAKE_CASE)  # 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...')
																lowerCamelCase__					=(
																    '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(_SCREAMING_SNAKE_CASE):
																					lowerCamelCase__					=[
																					    CommitOperationAdd(path_in_repo=_SCREAMING_SNAKE_CASE     ,    path_or_fileobj=_SCREAMING_SNAKE_CASE)
																					]
																					for shard_path in tf.io.gfile.glob(self._local_dir + '/tf_model-*.h5'):
																										operations += [
																										    CommitOperationAdd(path_in_repo=os.path.basename(_SCREAMING_SNAKE_CASE)     ,    path_or_fileobj=_SCREAMING_SNAKE_CASE)
																										]
																else:
																					lowerCamelCase__					=[CommitOperationAdd(path_in_repo=_SCREAMING_SNAKE_CASE     ,    path_or_fileobj=_SCREAMING_SNAKE_CASE)]

																lowerCamelCase__					=create_commit(
																    repo_id=self._model_name     ,    operations=_SCREAMING_SNAKE_CASE     ,    commit_message=_SCREAMING_SNAKE_CASE     ,    commit_description=_SCREAMING_SNAKE_CASE     ,    repo_type='model'     ,    create_pr=_SCREAMING_SNAKE_CASE     ,    ).pr_url
																self._logger.warning(F"PR open in {hub_pr_url}")







"""simple docstring"""



import inspect

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
						import jax




@require_flax
class 							__lowercase:






						def 		lowerCAmelCase (       self)      ->	str:




											"""simple docstring"""





											lowerCamelCase__					,  lowerCamelCase__					=self.prepare_init_args_and_inputs_for_common()

											lowerCamelCase__					=self.model_class(**_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=model.init(inputs_dict['prng_key']     ,    inputs_dict['sample'])
											jax.lax.stop_gradient(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=model.apply(_SCREAMING_SNAKE_CASE     ,    inputs_dict['sample'])

											if isinstance(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE):
																lowerCamelCase__					=output.sample

											self.assertIsNotNone(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=inputs_dict['sample'].shape
											self.assertEqual(output.shape     ,    _SCREAMING_SNAKE_CASE     ,    'Input and output shapes do not match')






						def 		lowerCAmelCase (       self)      ->	Optional[Any]:




											"""simple docstring"""





											lowerCamelCase__					,  lowerCamelCase__					=self.prepare_init_args_and_inputs_for_common()

											lowerCamelCase__					=16
											lowerCamelCase__					=(16, 32)

											lowerCamelCase__					=self.model_class(**_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=model.init(inputs_dict['prng_key']     ,    inputs_dict['sample'])
											jax.lax.stop_gradient(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=model.apply(_SCREAMING_SNAKE_CASE     ,    inputs_dict['sample'])

											if isinstance(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE):
																lowerCamelCase__					=output.sample

											self.assertIsNotNone(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=inputs_dict['sample'].shape
											self.assertEqual(output.shape     ,    _SCREAMING_SNAKE_CASE     ,    'Input and output shapes do not match')






						def 		lowerCAmelCase (       self)      ->	Optional[Any]:




											"""simple docstring"""





											lowerCamelCase__					='kwargs' in inspect.signature(self.model_class.__init__).parameters
											lowerCamelCase__					=len(self.model_class._deprecated_kwargs) > 0

											if has_kwarg_in_model_class and not has_deprecated_kwarg:
																raise ValueError(
																    F"{self.model_class} has `**kwargs` in its __init__ method but has not defined any deprecated kwargs"
																    ' under the `_deprecated_kwargs` class attribute. Make sure to either remove `**kwargs` if there are'
																    ' no deprecated arguments or add the deprecated argument with `_deprecated_kwargs ='
																    ' [<deprecated_argument>]`')

											if not has_kwarg_in_model_class and has_deprecated_kwarg:
																raise ValueError(
																    F"{self.model_class} doesn't have `**kwargs` in its __init__ method but has defined deprecated kwargs"
																    ' under the `_deprecated_kwargs` class attribute. Make sure to either add the `**kwargs` argument to'
																    F" {self.model_class}.__init__ if there are deprecated arguments or remove the deprecated argument"
																    ' from `_deprecated_kwargs = [<deprecated_argument>]`')