






'''simple docstring'''







import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


UpperCamelCase     :   Optional[Any]	      =		logging.get_logger(__name__)




class 		UpperCamelCase__     :




		'''simple docstring'''





		_UpperCamelCase						=       None



@experimental
def     A__      (					__lowerCAmelCase		:      str ,     __lowerCAmelCase		:      Any ,     __lowerCAmelCase		:      Union[str, Any] ,     __lowerCAmelCase		:      Optional[int] ,     __lowerCAmelCase		:      Optional[Any] ,     __lowerCAmelCase		:      Optional[Any] ,     __lowerCAmelCase		:      Tuple    ):
						if ParallelBackendConfig.backend_name is None:
												return _map_with_multiprocessing_pool(
												    __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase    )

						return _map_with_joblib(__lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase    )
def     A__      (					__lowerCAmelCase		:      int ,     __lowerCAmelCase		:      List[str] ,     __lowerCAmelCase		:      Tuple ,     __lowerCAmelCase		:      List[str] ,     __lowerCAmelCase		:      int ,     __lowerCAmelCase		:      str ,     __lowerCAmelCase		:      List[str]    ):
						lowerCamelCase__								=	num_proc if num_proc <= len(__lowerCAmelCase    ) else len(__lowerCAmelCase    )
						lowerCamelCase__								=	[]  # We organize the splits ourselve (contiguous splits)
						for index in range(__lowerCAmelCase    ):
												lowerCamelCase__								=	len(__lowerCAmelCase    ) // num_proc
												lowerCamelCase__								=	len(__lowerCAmelCase    ) % num_proc
												lowerCamelCase__								=	div * index + min(__lowerCAmelCase ,     __lowerCAmelCase    )
												lowerCamelCase__								=	start + div + (1 if index < mod else 0)
												split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)    )

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

						logger.info(
						    F'''Spawning {num_proc} processes for {len(__lowerCAmelCase    )} objects in slices of {[len(i[1]    ) for i in split_kwds]}'''    )
						lowerCamelCase__							,      lowerCamelCase__								=	None, None
						if not disable_tqdm:
												lowerCamelCase__							,      lowerCamelCase__								=	(RLock(),), tqdm.set_lock
						with Pool(__lowerCAmelCase ,     initargs=__lowerCAmelCase ,     initializer=__lowerCAmelCase    ) as pool:
												lowerCamelCase__								=	pool.map(__lowerCAmelCase ,     __lowerCAmelCase    )
						logger.info(F'''Finished {num_proc} processes'''    )
						lowerCamelCase__								=	[obj for proc_res in mapped for obj in proc_res]
						logger.info(F'''Unpacked {len(__lowerCAmelCase    )} objects'''    )

						return mapped
def     A__      (					__lowerCAmelCase		:      str ,     __lowerCAmelCase		:      str ,     __lowerCAmelCase		:      Tuple ,     __lowerCAmelCase		:      Tuple ,     __lowerCAmelCase		:      str ,     __lowerCAmelCase		:      Dict ,     __lowerCAmelCase		:      Optional[int]    ):
						# progress bar is not yet supported for _map_with_joblib, because tqdm couldn't accurately be applied to joblib,
						# and it requires monkey-patching joblib internal classes which is subject to change
						import joblib

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





@experimental
@contextlib.contextmanager
def     A__      (					__lowerCAmelCase		:      str    ):
						lowerCamelCase__								=	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:
												lowerCamelCase__								=	None







import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin
class 						UpperCAmelCase  (							UpperCAmelCase__			,  UpperCAmelCase__			,  UpperCAmelCase__     ):







						'''simple docstring'''





						@register_to_config
						def __init__(		self				:	Tuple			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	float			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	str			,      __lowercase				:	bool = False			,      ):






							"""simple docstring"""






							super().__init__()

							snake_case_						       =	nn.Embedding(__lowercase			,      __lowercase		)

							snake_case_						       =	nn.Embedding(__lowercase			,      __lowercase		)
							snake_case_						       =	False

							snake_case_						       =	nn.Dropout(p=__lowercase		)

							snake_case_						       =	TaConfig(
							    vocab_size=__lowercase			,      d_model=__lowercase			,      num_heads=__lowercase			,      d_kv=__lowercase			,      d_ff=__lowercase			,      dropout_rate=__lowercase			,      feed_forward_proj=__lowercase			,      is_decoder=__lowercase			,      is_encoder_decoder=__lowercase			,      )

							snake_case_						       =	nn.ModuleList()
							for lyr_num in range(__lowercase		):
								snake_case_						       =	TaBlock(__lowercase		)
								self.encoders.append(__lowercase		)

							snake_case_						       =	TaLayerNorm(__lowercase		)
							snake_case_						       =	nn.Dropout(p=__lowercase		)






						def 			snake_case__		(		self				:	str			,      __lowercase				:	Dict			,      __lowercase				:	str		):






							"""simple docstring"""






							snake_case_						       =	self.token_embedder(__lowercase		)

							snake_case_						       =	encoder_input_tokens.shape[1]
							snake_case_						       =	torch.arange(__lowercase			,      device=encoder_input_tokens.device		)
							x += self.position_encoding(__lowercase		)

							snake_case_						       =	self.dropout_pre(__lowercase		)

							# inverted the attention mask
							snake_case_						       =	encoder_input_tokens.size()
							snake_case_						       =	self.get_extended_attention_mask(__lowercase			,      __lowercase		)

							for lyr in self.encoders:
								snake_case_						       =	lyr(__lowercase			,      __lowercase		)[0]
							snake_case_						       =	self.layer_norm(__lowercase		)

							return self.dropout_post(__lowercase		), encoder_inputs_mask
