



"""simple docstring"""





from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch



@require_torch
@require_tf
class    _lowerCAmelCase    (				__A		):





		"""simple docstring"""






		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : Optional[int]						=			SMALL_MODEL_IDENTIFIER
								snake_case_   : int						=			'''pt'''
								snake_case_   : Tuple						=			'''tf'''





		def    UpperCAmelCase__							(					self							,    _lowercase   ) ->						int:







								'''simple docstring'''

								snake_case_   : List[Any]						=			AutoModel.from_pretrained(self.test_model   )
								model_pt.save_pretrained(_lowercase   )





		def    UpperCAmelCase__							(					self							,    _lowercase   ) ->						int:







								'''simple docstring'''

								snake_case_   : Tuple						=			TFAutoModel.from_pretrained(self.test_model							,    from_pt=_lowercase   )
								model_tf.save_pretrained(_lowercase   )





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : Optional[Any]						=			'''mock_framework'''

								# Framework provided - return whatever the user provides
								snake_case_   : Optional[int]						=			FeaturesManager.determine_framework(self.test_model							,    _lowercase   )
								self.assertEqual(_lowercase							,    _lowercase   )

								# Local checkpoint and framework provided - return provided framework
								# PyTorch checkpoint
								with TemporaryDirectory() as local_pt_ckpt:
														self._setup_pt_ckpt(_lowercase   )
														snake_case_   : Union[str, Any]						=			FeaturesManager.determine_framework(_lowercase							,    _lowercase   )
														self.assertEqual(_lowercase							,    _lowercase   )

								# TensorFlow checkpoint
								with TemporaryDirectory() as local_tf_ckpt:
														self._setup_tf_ckpt(_lowercase   )
														snake_case_   : Optional[Any]						=			FeaturesManager.determine_framework(_lowercase							,    _lowercase   )
														self.assertEqual(_lowercase							,    _lowercase   )





		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								with TemporaryDirectory() as local_pt_ckpt:
														self._setup_pt_ckpt(_lowercase   )
														snake_case_   : List[str]						=			FeaturesManager.determine_framework(_lowercase   )
														self.assertEqual(_lowercase							,    self.framework_pt   )

								# TensorFlow checkpoint
								with TemporaryDirectory() as local_tf_ckpt:
														self._setup_tf_ckpt(_lowercase   )
														snake_case_   : Optional[Any]						=			FeaturesManager.determine_framework(_lowercase   )
														self.assertEqual(_lowercase							,    self.framework_tf   )

								# Invalid local checkpoint
								with TemporaryDirectory() as local_invalid_ckpt:
														with self.assertRaises(_lowercase   ):
																				snake_case_   : Union[str, Any]						=			FeaturesManager.determine_framework(_lowercase   )





		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								snake_case_   : List[str]						=			MagicMock(return_value=_lowercase   )
								with patch("""transformers.onnx.features.is_tf_available"""							,    _lowercase   ):
														snake_case_   : List[Any]						=			FeaturesManager.determine_framework(self.test_model   )
														self.assertEqual(_lowercase							,    self.framework_pt   )

								# PyTorch not in environment -> use TensorFlow
								snake_case_   : Optional[int]						=			MagicMock(return_value=_lowercase   )
								with patch("""transformers.onnx.features.is_torch_available"""							,    _lowercase   ):
														snake_case_   : Tuple						=			FeaturesManager.determine_framework(self.test_model   )
														self.assertEqual(_lowercase							,    self.framework_tf   )

								# Both in environment -> use PyTorch
								snake_case_   : List[str]						=			MagicMock(return_value=_lowercase   )
								snake_case_   : List[str]						=			MagicMock(return_value=_lowercase   )
								with patch("""transformers.onnx.features.is_tf_available"""							,    _lowercase   ), patch(
								    """transformers.onnx.features.is_torch_available"""							,    _lowercase   ):
														snake_case_   : Tuple						=			FeaturesManager.determine_framework(self.test_model   )
														self.assertEqual(_lowercase							,    self.framework_pt   )

								# Both not in environment -> raise error
								snake_case_   : str						=			MagicMock(return_value=_lowercase   )
								snake_case_   : str						=			MagicMock(return_value=_lowercase   )
								with patch("""transformers.onnx.features.is_tf_available"""							,    _lowercase   ), patch(
								    """transformers.onnx.features.is_torch_available"""							,    _lowercase   ):
														with self.assertRaises(_lowercase   ):
																				snake_case_   : List[str]						=			FeaturesManager.determine_framework(self.test_model   )






"""simple docstring"""





import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
					import optax
					from flax.training.common_utils import onehot

					from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
					from transformers.models.ta.modeling_flax_ta import shift_tokens_right



@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		@slow
		def    UpperCAmelCase__							(					self   ) ->						Optional[int]:







								'''simple docstring'''

								snake_case_   : int						=			FlaxMTaForConditionalGeneration.from_pretrained("""google/mt5-small"""   )
								snake_case_   : List[Any]						=			AutoTokenizer.from_pretrained("""google/mt5-small"""   )

								snake_case_   : Dict						=			tokenizer("""Hello there"""							,    return_tensors="""np"""   ).input_ids
								snake_case_   : List[str]						=			tokenizer("""Hi I am"""							,    return_tensors="""np"""   ).input_ids

								snake_case_   : Optional[Any]						=			shift_tokens_right(_lowercase							,    model.config.pad_token_id							,    model.config.decoder_start_token_id   )

								snake_case_   : Tuple						=			model(_lowercase							,    decoder_input_ids=_lowercase   ).logits
								snake_case_   : Tuple						=			optax.softmax_cross_entropy(_lowercase							,    onehot(_lowercase							,    logits.shape[-1]   )   ).mean()

								snake_case_   : List[str]						=			-(labels.shape[-1] * loss.item())

								snake_case_   : Optional[int]						=			-84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE   ) < 1E-4   )

