






# flake8: noqa
# Lint as: python3
A        =			[
    'VerificationMode',
    'Version',
    'disable_progress_bar',
    'enable_progress_bar',
    'is_progress_bar_enabled',
    'experimental',
]

from .info_utils import VerificationMode
from .logging import disable_progress_bar, enable_progress_bar, is_progress_bar_enabled
from .version import Version
from .experimental import experimental







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 							__a							(  __A						):





		'''simple docstring'''



		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Any													=   SMALL_MODEL_IDENTIFIER
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   'pt'
						SCREAMING_SNAKE_CASE_   :	Tuple													=   'tf'


		def   __snake_case			(						self ,    UpperCamelCase__       ):
						SCREAMING_SNAKE_CASE_   :	Dict													=   AutoModel.from_pretrained(self.test_model       )
						model_pt.save_pretrained(UpperCamelCase__       )


		def   __snake_case			(						self ,    UpperCamelCase__       ):
						SCREAMING_SNAKE_CASE_   :	int													=   TFAutoModel.from_pretrained(self.test_model ,    from_pt=UpperCamelCase__       )
						model_tf.save_pretrained(UpperCamelCase__       )


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   'mock_framework'

						# Framework provided - return whatever the user provides
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   FeaturesManager.determine_framework(self.test_model ,    UpperCamelCase__       )
						self.assertEqual(UpperCamelCase__ ,    UpperCamelCase__       )

						# Local checkpoint and framework provided - return provided framework
						# PyTorch checkpoint
						with TemporaryDirectory() as local_pt_ckpt:
										self._setup_pt_ckpt(UpperCamelCase__       )
										SCREAMING_SNAKE_CASE_   :	Dict													=   FeaturesManager.determine_framework(UpperCamelCase__ ,    UpperCamelCase__       )
										self.assertEqual(UpperCamelCase__ ,    UpperCamelCase__       )

						# TensorFlow checkpoint
						with TemporaryDirectory() as local_tf_ckpt:
										self._setup_tf_ckpt(UpperCamelCase__       )
										SCREAMING_SNAKE_CASE_   :	List[str]													=   FeaturesManager.determine_framework(UpperCamelCase__ ,    UpperCamelCase__       )
										self.assertEqual(UpperCamelCase__ ,    UpperCamelCase__       )


		def   __snake_case			(						self       ):
						# PyTorch checkpoint
						with TemporaryDirectory() as local_pt_ckpt:
										self._setup_pt_ckpt(UpperCamelCase__       )
										SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   FeaturesManager.determine_framework(UpperCamelCase__       )
										self.assertEqual(UpperCamelCase__ ,    self.framework_pt       )

						# TensorFlow checkpoint
						with TemporaryDirectory() as local_tf_ckpt:
										self._setup_tf_ckpt(UpperCamelCase__       )
										SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   FeaturesManager.determine_framework(UpperCamelCase__       )
										self.assertEqual(UpperCamelCase__ ,    self.framework_tf       )

						# Invalid local checkpoint
						with TemporaryDirectory() as local_invalid_ckpt:
										with self.assertRaises(UpperCamelCase__       ):
														SCREAMING_SNAKE_CASE_   :	Any													=   FeaturesManager.determine_framework(UpperCamelCase__       )






		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   MagicMock(return_value=UpperCamelCase__       )
						with patch('transformers.onnx.features.is_tf_available' ,    UpperCamelCase__       ):
										SCREAMING_SNAKE_CASE_   :	str													=   FeaturesManager.determine_framework(self.test_model       )
										self.assertEqual(UpperCamelCase__ ,    self.framework_pt       )

						# PyTorch not in environment -> use TensorFlow
						SCREAMING_SNAKE_CASE_   :	List[str]													=   MagicMock(return_value=UpperCamelCase__       )
						with patch('transformers.onnx.features.is_torch_available' ,    UpperCamelCase__       ):
										SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   FeaturesManager.determine_framework(self.test_model       )
										self.assertEqual(UpperCamelCase__ ,    self.framework_tf       )

						# Both in environment -> use PyTorch
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   MagicMock(return_value=UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	Tuple													=   MagicMock(return_value=UpperCamelCase__       )
						with patch('transformers.onnx.features.is_tf_available' ,    UpperCamelCase__       ), patch(
						    'transformers.onnx.features.is_torch_available' ,    UpperCamelCase__       ):
										SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   FeaturesManager.determine_framework(self.test_model       )
										self.assertEqual(UpperCamelCase__ ,    self.framework_pt       )

						# Both not in environment -> raise error
						SCREAMING_SNAKE_CASE_   :	Tuple													=   MagicMock(return_value=UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	Any													=   MagicMock(return_value=UpperCamelCase__       )
						with patch('transformers.onnx.features.is_tf_available' ,    UpperCamelCase__       ), patch(
						    'transformers.onnx.features.is_torch_available' ,    UpperCamelCase__       ):
										with self.assertRaises(UpperCamelCase__       ):
														SCREAMING_SNAKE_CASE_   :	Optional[int]													=   FeaturesManager.determine_framework(self.test_model       )