






import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


UpperCAmelCase_ :						Dict      =      logging.get_logger(__name__)

class     _SCREAMING_SNAKE_CASE       :
	snake_case__							:     Dict     	=  None

@experimental
def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Optional[int]	,			__magic_name__					:  Optional[Any]	,			__magic_name__					:  Optional[int]	,			__magic_name__					:  Union[str, Any]	,			__magic_name__					:  List[str]	,			__magic_name__					:  Any	,			__magic_name__					:  Dict					)		->			Optional[int]:






	"""simple docstring"""



	if ParallelBackendConfig.backend_name is None:
		return _map_with_multiprocessing_pool(
		    __magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__					)

	return _map_with_joblib(__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__					)

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Any	,			__magic_name__					:  Tuple	,			__magic_name__					:  List[Any]	,			__magic_name__					:  str	,			__magic_name__					:  List[Any]	,			__magic_name__					:  str	,			__magic_name__					:  Dict					)		->			Optional[int]:






	"""simple docstring"""



	UpperCamelCase     :Optional[Any]          =	num_proc if num_proc <= len(__magic_name__					) else len(__magic_name__					)
	UpperCamelCase     :int          =	[]  # We organize the splits ourselve (contiguous splits)
	for index in range(__magic_name__					):
		UpperCamelCase     :str          =	len(__magic_name__					) // num_proc
		UpperCamelCase     :List[Any]          =	len(__magic_name__					) % num_proc
		UpperCamelCase     :Union[str, Any]          =	div * index + min(__magic_name__	,			__magic_name__					)
		UpperCamelCase     :Optional[Any]          =	start + div + (1 if index < mod else 0)
		split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)					)

	if len(__magic_name__					) != sum(len(i[1]					) for i in split_kwds					):
		raise ValueError(
		    f"""Error dividing inputs iterable among processes. """
		    f"""Total number of objects {len(__magic_name__					)}, """
		    f"""length: {sum(len(i[1]					) for i in split_kwds					)}"""					)

	logger.info(
	    f"""Spawning {num_proc} processes for {len(__magic_name__					)} objects in slices of {[len(i[1]					) for i in split_kwds]}"""					)
	UpperCamelCase     ,      UpperCamelCase     :List[str]          =	None, None
	if not disable_tqdm:
		UpperCamelCase     ,      UpperCamelCase     :Dict          =	(RLock(),), tqdm.set_lock
	with Pool(__magic_name__	,			initargs=__magic_name__	,			initializer=__magic_name__					) as pool:
		UpperCamelCase     :Dict          =	pool.map(__magic_name__	,			__magic_name__					)
	logger.info(f"""Finished {num_proc} processes"""					)
	UpperCamelCase     :Optional[Any]          =	[obj for proc_res in mapped for obj in proc_res]
	logger.info(f"""Unpacked {len(__magic_name__					)} objects"""					)

	return mapped

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Dict	,			__magic_name__					:  int	,			__magic_name__					:  int	,			__magic_name__					:  Union[str, Any]	,			__magic_name__					:  Any	,			__magic_name__					:  int	,			__magic_name__					:  Union[str, Any]					)		->			Union[str, Any]:






	"""simple docstring"""



	import joblib

	with joblib.parallel_backend(ParallelBackendConfig.backend_name	,			n_jobs=__magic_name__					):
		return joblib.Parallel()(
		    joblib.delayed(__magic_name__					)((function, obj, types, None, True, None)					) for obj in iterable					)

@experimental
@contextlib.contextmanager
def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  str					)		->			List[Any]:






	"""simple docstring"""



	UpperCamelCase     :Union[str, Any]          =	backend_name

	if backend_name == "spark":
		from joblibspark import register_spark

		register_spark()

		# TODO: call create_cache_and_write_probe if "download" in steps
		# TODO: raise NotImplementedError when Dataset.map etc is called

	try:
		yield
	finally:
		UpperCamelCase     :Union[str, Any]          =	None


'''simple docstring'''




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  A				(							unittest.TestCase						):





     @slow
     def 				__lowerCAmelCase							(							self	) ->				List[Any]:






         """simple docstring"""






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

                 A       :   List[str]					=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                 self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                 self.assertIsInstance(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






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

                 A       :   Any					=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                 self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                 self.assertIsInstance(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         for model_name in ["bert-base-cased", "bert-large-uncased"]:
             A       :   Optional[int]					=						AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   List[str]					=						FlaxBertModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   Optional[Any]					=						tokenizer('''Do you support jax jitted function?'''						,			return_tensors=TensorType.JAX	)

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

             eval(**SCREAMING_SNAKE_CASE	).block_until_ready()





     @slow
     def 				__lowerCAmelCase							(							self	) ->				List[str]:






         """simple docstring"""






         for model_name in ["roberta-base", "roberta-large"]:
             A       :   List[str]					=						AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   Union[str, Any]					=						FlaxRobertaModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   int					=						tokenizer('''Do you support jax jitted function?'''						,			return_tensors=TensorType.JAX	)

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

             eval(**SCREAMING_SNAKE_CASE	).block_until_ready()





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






         """simple docstring"""






         with self.assertRaisesRegex(
             SCREAMING_SNAKE_CASE						,			'''bert-base is not a local folder and is not a valid model identifier'''	):
             A       :   List[Any]					=						FlaxAutoModel.from_pretrained('''bert-base'''	)





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         with self.assertRaisesRegex(
             SCREAMING_SNAKE_CASE						,			R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''	):
             A       :   Optional[int]					=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE						,			revision='''aaaaaa'''	)





     def 				__lowerCAmelCase							(							self	) ->				str:






         """simple docstring"""






         with self.assertRaisesRegex(
             SCREAMING_SNAKE_CASE						,			'''hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'''						,			):
             A       :   List[str]					=						FlaxAutoModel.from_pretrained('''hf-internal-testing/config-no-model'''	)





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






         """simple docstring"""






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




