

import sys
from pathlib import Path


A__						=							Path(__file__).resolve().parents[3] / "src"
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(42)

A__						=							{"base": "patrickvonplaten/wav2vec2_tiny_random", "robust": "patrickvonplaten/wav2vec2_tiny_random_robust"}

A__						=							"zero2"
A__						=							"zero3"
A__						=							[ZEROa, ZEROa]




def 					_lowercase     (				a_				:					str					,a_				:					Optional[Any]					,a_				:					Optional[int]				)				->	int:




     '''simple docstring'''







     __magic_name__							=    parameterized.to_safe_name('_'.join(str(__UpperCAmelCase				) for x in param.args				)				)
     return F'{func.__name__}_{param_based_name}'


# Cartesian-product of zero stages with models to test
A__						=							list(itertools.product(stages, models.keys()))


@slow
@require_deepspeed
@require_torch_gpu
class 				__UpperCamelCase				(					_snake_case			):







 @parameterized.expand(__UpperCamelCase  , name_func=__UpperCamelCase      )
 def 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: Tuple      ):

      '''simple docstring'''





      self.run_and_check(
          stage=__UpperCamelCase  , model=__UpperCamelCase  , distributed=__UpperCamelCase  , fpaa=__UpperCamelCase  , )







 @require_torch_multi_gpu
 @parameterized.expand(__UpperCamelCase  , name_func=__UpperCamelCase      )
 def 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Dict      ):

      '''simple docstring'''





      self.run_and_check(
          stage=__UpperCamelCase  , model=__UpperCamelCase  , distributed=__UpperCamelCase  , fpaa=__UpperCamelCase  , )







 @parameterized.expand(__UpperCamelCase  , name_func=__UpperCamelCase      )
 def 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: List[Any]      ):

      '''simple docstring'''





      self.run_and_check(
          stage=__UpperCamelCase  , model=__UpperCamelCase  , distributed=__UpperCamelCase  , fpaa=__UpperCamelCase  , )







 @require_torch_multi_gpu
 @parameterized.expand(__UpperCamelCase  , name_func=__UpperCamelCase      )
 def 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: Optional[Any]      ):

      '''simple docstring'''





      self.run_and_check(
          stage=__UpperCamelCase  , model=__UpperCamelCase  , distributed=__UpperCamelCase  , fpaa=__UpperCamelCase  , )







 def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]  , __UpperCamelCase: Dict      ):

      '''simple docstring'''





      pass







 def 	_SCREAMING_SNAKE_CASE							(   self: Tuple  , __UpperCamelCase: str  , __UpperCamelCase: str  , __UpperCamelCase: int = 10  , __UpperCamelCase: bool = True  , __UpperCamelCase: bool = True  , __UpperCamelCase: bool = True  , ):

      '''simple docstring'''





      __magic_name__							=    models[model]

      __magic_name__							=    self.run_trainer(
          stage=__UpperCamelCase  , model_name=__UpperCamelCase  , eval_steps=__UpperCamelCase  , num_train_epochs=1  , distributed=__UpperCamelCase  , fpaa=__UpperCamelCase  , )

      self.do_checks(__UpperCamelCase      )

      return output_dir







 def 	_SCREAMING_SNAKE_CASE							(   self: List[str]  , __UpperCamelCase: str  , __UpperCamelCase: str  , __UpperCamelCase: int = 10  , __UpperCamelCase: int = 1  , __UpperCamelCase: bool = True  , __UpperCamelCase: bool = True  , ):

      '''simple docstring'''





      __magic_name__							=    self.get_auto_remove_tmp_dir('./xxx'  , after=__UpperCamelCase      )
      __magic_name__							=    F'\n            --model_name_or_path {model_name}\n            --dataset_name hf-internal-testing/librispeech_asr_dummy\n            --dataset_config_name clean\n            --train_split_name validation\n            --validation_split_name validation\n            --output_dir {output_dir}\n            --num_train_epochs {str(__UpperCamelCase      )}\n            --per_device_train_batch_size 2\n            --per_device_eval_batch_size 2\n            --evaluation_strategy steps\n            --learning_rate 5e-4\n            --warmup_steps 8\n            --orthography timit\n            --preprocessing_num_workers 1\n            --group_by_length\n            --freeze_feature_extractor\n            --report_to none\n            --save_steps 0\n            --eval_steps {eval_steps}\n            --report_to none\n        '.split()

      if fpaa:
           args.extend(['--fp16']      )

      # currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
      # hence the separate config files
      __magic_name__							=    F'--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json'.split()
      __magic_name__							=    [F'{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py']
      __magic_name__							=    self.get_launcher(__UpperCamelCase      )

      __magic_name__							=    launcher + script + args + ds_args
      # keep for quick debug
      # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
      execute_subprocess_async(__UpperCamelCase  , env=self.get_env()      )

      return output_dir







 def 	_SCREAMING_SNAKE_CASE							(   self: Tuple  , __UpperCamelCase: List[Any]=False      ):

      '''simple docstring'''





      __magic_name__							=    min(2  , get_gpu_count()      ) if distributed else 1
      return F'deepspeed --num_nodes 1 --num_gpus {num_gpus}'.split()







import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
						import jax

						from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
						from transformers.models.bert.modeling_flax_bert import FlaxBertModel
						from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel

@require_flax
class 				__UpperCamelCase				(					unittest.TestCase			):







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: Any      ):

						'''simple docstring'''





						for model_name in ["bert-base-cased", "bert-large-uncased"]:
											with self.subTest(__UpperCamelCase      ):
																__magic_name__							=    AutoConfig.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )

																__magic_name__							=    FlaxAutoModel.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						for model_name in ["roberta-base", "roberta-large"]:
											with self.subTest(__UpperCamelCase      ):
																__magic_name__							=    AutoConfig.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )

																__magic_name__							=    FlaxAutoModel.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						for model_name in ["bert-base-cased", "bert-large-uncased"]:
											__magic_name__							=    AutoTokenizer.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    FlaxBertModel.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    tokenizer('Do you support jax jitted function?'  , return_tensors=TensorType.JAX      )

											@jax.jit
											def eval(**__UpperCamelCase: Optional[Any]      ):
																return model(**__UpperCamelCase      )

											eval(**__UpperCamelCase      ).block_until_ready()







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						for model_name in ["roberta-base", "roberta-large"]:
											__magic_name__							=    AutoTokenizer.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    FlaxRobertaModel.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    tokenizer('Do you support jax jitted function?'  , return_tensors=TensorType.JAX      )

											@jax.jit
											def eval(**__UpperCamelCase: Any      ):
																return model(**__UpperCamelCase      )

											eval(**__UpperCamelCase      ).block_until_ready()







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						with self.assertRaisesRegex(
						    __UpperCamelCase  , 'bert-base is not a local folder and is not a valid model identifier'      ):
											__magic_name__							=    FlaxAutoModel.from_pretrained('bert-base'      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						with self.assertRaisesRegex(
						    __UpperCamelCase  , r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'      ):
											__magic_name__							=    FlaxAutoModel.from_pretrained(__UpperCamelCase  , revision='aaaaaa'      )







	def 	_SCREAMING_SNAKE_CASE							(   self: str      ):

						'''simple docstring'''





						with self.assertRaisesRegex(
						    __UpperCamelCase  , 'hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'  , ):
											__magic_name__							=    FlaxAutoModel.from_pretrained('hf-internal-testing/config-no-model'      )







	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						with self.assertRaisesRegex(__UpperCamelCase  , 'Use `from_pt=True` to load this model'      ):
											__magic_name__							=    FlaxAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'      )




