'''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     A__							(       lowerCamelCase_   ):





   """simple docstring"""





   def 			_lowerCAmelCase		( self     :		Any			)       ->    int:







      """simple docstring"""




      _UpperCAmelCase							:	str							=   SMALL_MODEL_IDENTIFIER
      _UpperCAmelCase							:	Optional[int]							=   "pt"
      _UpperCAmelCase							:	Dict							=   "tf"




   def 			_lowerCAmelCase		( self     :		int						,     lowerCAmelCase__     :		int			)       ->    int:







      """simple docstring"""




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




   def 			_lowerCAmelCase		( self     :		Any						,     lowerCAmelCase__     :		Union[str, Any]			)       ->    Dict:







      """simple docstring"""




      _UpperCAmelCase							:	Dict							=   TFAutoModel.from_pretrained(self.test_model						,     from_pt=_lowercase			)
      model_tf.save_pretrained(_lowercase			)




   def 			_lowerCAmelCase		( self     :		Any			)       ->    Optional[int]:







      """simple docstring"""




      _UpperCAmelCase							:	Tuple							=   "mock_framework"

      # Framework provided - return whatever the user provides
      _UpperCAmelCase							:	Tuple							=   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			)
         _UpperCAmelCase							:	int							=   FeaturesManager.determine_framework(_lowercase						,     _lowercase			)
         self.assertEqual(_lowercase						,     _lowercase			)

      # TensorFlow checkpoint
      with TemporaryDirectory() as local_tf_ckpt:
         self._setup_tf_ckpt(_lowercase			)
         _UpperCAmelCase							:	Tuple							=   FeaturesManager.determine_framework(_lowercase						,     _lowercase			)
         self.assertEqual(_lowercase						,     _lowercase			)




   def 			_lowerCAmelCase		( self     :		Tuple			)       ->    int:







      """simple docstring"""




      with TemporaryDirectory() as local_pt_ckpt:
         self._setup_pt_ckpt(_lowercase			)
         _UpperCAmelCase							:	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			)
         _UpperCAmelCase							:	int							=   FeaturesManager.determine_framework(_lowercase			)
         self.assertEqual(_lowercase						,     self.framework_tf			)

      # Invalid local checkpoint
      with TemporaryDirectory() as local_invalid_ckpt:
         with self.assertRaises(_lowercase			):
            _UpperCAmelCase							:	int							=   FeaturesManager.determine_framework(_lowercase			)




   def 			_lowerCAmelCase		( self     :		Dict			)       ->    List[Any]:







      """simple docstring"""




      _UpperCAmelCase							:	int							=   MagicMock(return_value=_lowercase			)
      with patch("transformers.onnx.features.is_tf_available"						,     _lowercase			):
         _UpperCAmelCase							:	Dict							=   FeaturesManager.determine_framework(self.test_model			)
         self.assertEqual(_lowercase						,     self.framework_pt			)

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

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

      # Both not in environment -> raise error
      _UpperCAmelCase							:	List[str]							=   MagicMock(return_value=_lowercase			)
      _UpperCAmelCase							:	Union[str, Any]							=   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			):
            _UpperCAmelCase							:	List[str]							=   FeaturesManager.determine_framework(self.test_model			)
import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp



def 	__SCREAMING_SNAKE_CASE  (					)							->					Any:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__										=      ArgumentParser(
      description=(
          """PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"""
      )	)

  # Optional arguments for the launch helper
  parser.add_argument("""--num_cores""" ,     type=__UpperCamelCase ,     default=1 ,     help="""Number of TPU cores to use (1 or 8)."""	)

  # positional
  parser.add_argument(
      """training_script""" ,     type=__UpperCamelCase ,     help=(
          """The full path to the single TPU training """
          """program/script to be launched in parallel, """
          """followed by all the arguments for the """
          """training script"""
      ) ,     )

  # rest from the training program
  parser.add_argument("""training_script_args""" ,     nargs=__UpperCamelCase	)

  return parser.parse_args()



def 	__SCREAMING_SNAKE_CASE  (					)							->					Dict:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__										=      parse_args()

  # Import training_script as a module.
  SCREAMING_SNAKE_CASE__										=      Path(args.training_script	)
  sys.path.append(str(script_fpath.parent.resolve()	)	)
  SCREAMING_SNAKE_CASE__										=      script_fpath.stem
  SCREAMING_SNAKE_CASE__										=      importlib.import_module(__UpperCamelCase	)

  # Patch sys.argv
  SCREAMING_SNAKE_CASE__										=      [args.training_script] + args.training_script_args + ["""--tpu_num_cores""", str(args.num_cores	)]

  xmp.spawn(mod._mp_fn ,     args=() ,     nprocs=args.num_cores	)


if __name__ == "__main__":
   main()
