


"""simple docstring"""







import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput

class 	UpperCamelCase    (					lowerCAmelCase__		):



				def __init__(self       :   Any			,						_A       :   Dict			,						_A       :   Any			,						_A       :   List[Any])					->			Union[str, Any]:
					__snake_case    :							Optional[int]   	=				dataset
					__snake_case    :							Optional[int]   	=				process
					__snake_case    :							Any   	=				params



				def __len__(self       :   Union[str, Any])					->			str:
					return len(self.dataset)






				def __getitem__(self       :   Optional[Any]			,						_A       :   Optional[int])					->			Dict:
					__snake_case    :							List[str]   	=				self.dataset[i]
					__snake_case    :							List[str]   	=				self.process(_SCREAMING_SNAKE_CASE			,						**self.params)
					return processed







class 	UpperCamelCase    (					lowerCAmelCase__		):



				def __init__(self       :   Optional[int]			,						_A       :   Dict			,						_A       :   Dict			,						_A       :   str			,						_A       :   Dict=None)					->			List[str]:
					__snake_case    :							Tuple   	=				loader
					__snake_case    :							Optional[int]   	=				infer
					__snake_case    :							Tuple   	=				params
					if loader_batch_size == 1:
						# Let's spare some time by deactivating altogether
						__snake_case    :							List[str]   	=				None
					__snake_case    :							str   	=				loader_batch_size

					# Internal bookkeeping
					__snake_case    :							Optional[Any]   	=				None
					__snake_case    :							str   	=				None



				def __len__(self       :   Dict)					->			List[Any]:
					return len(self.loader)



				def __iter__(self       :   Tuple)					->			Tuple:
					__snake_case    :							Dict   	=				iter(self.loader)
					return self



				def       _lowercase  (self       :   Any)					->			int:
					if isinstance(self._loader_batch_data			,						torch.Tensor):
						# Batch data is simple tensor, just fetch the slice
						__snake_case    :							Union[str, Any]   	=				self._loader_batch_data[self._loader_batch_index]
					else:
						# Batch data is assumed to be BaseModelOutput (or dict)
						__snake_case    :							Dict   	=				{}
						for k, element in self._loader_batch_data.items():
							if isinstance(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE):
								# Convert ModelOutput to tuple first
								__snake_case    :							List[str]   	=				element.to_tuple()
								if isinstance(element[0]			,						torch.Tensor):
									__snake_case    :							Any   	=				tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
								elif isinstance(element[0]			,						np.ndarray):
									__snake_case    :							Union[str, Any]   	=				tuple(np.expand_dims(el[self._loader_batch_index]			,						0) for el in element)
								continue
							if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE):
								# Those are stored as lists of tensors so need specific unbatching.
								if isinstance(element[0]			,						torch.Tensor):
									__snake_case    :							int   	=				tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
								elif isinstance(element[0]			,						np.ndarray):
									__snake_case    :							Dict   	=				tuple(np.expand_dims(el[self._loader_batch_index]			,						0) for el in element)
								continue
							if element is None:
								# This can happen for optional data that get passed around
								__snake_case    :							Union[str, Any]   	=				None
							elif isinstance(element[self._loader_batch_index]			,						torch.Tensor):
								# Take correct batch data, but make it looked like batch_size=1
								# For compatibility with other methods within transformers

								__snake_case    :							Optional[Any]   	=				element[self._loader_batch_index].unsqueeze(0)
							elif isinstance(element[self._loader_batch_index]			,						np.ndarray):
								# Take correct batch data, but make it looked like batch_size=1
								# For compatibility with other methods within transformers
								__snake_case    :							str   	=				np.expand_dims(element[self._loader_batch_index]			,						0)
							else:
								# This is typically a list, so no need to `unsqueeze`.
								__snake_case    :							List[Any]   	=				element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
						__snake_case    :							Tuple   	=				self._loader_batch_data.__class__(_SCREAMING_SNAKE_CASE)
					self._loader_batch_index += 1
					return result






				def       _lowercase  (self       :   str)					->			Union[str, Any]:
					if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
						# We are currently unrolling a batch so we just need to return
						# the current item within a batch
						return self.loader_batch_item()

					# We're out of items within a batch
					__snake_case    :							Dict   	=				next(self.iterator)
					__snake_case    :							int   	=				self.infer(_SCREAMING_SNAKE_CASE			,						**self.params)
					# We now have a batch of "inferred things".
					if self.loader_batch_size is not None:
						# Try to infer the size of the batch
						if isinstance(_SCREAMING_SNAKE_CASE			,						torch.Tensor):
							__snake_case    :							str   	=				processed
						else:
							__snake_case    :							int   	=				list(processed.keys())[0]
							__snake_case    :							Any   	=				processed[key]
						if isinstance(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE):
							__snake_case    :							str   	=				len(_SCREAMING_SNAKE_CASE)
						else:
							__snake_case    :							List[str]   	=				first_tensor.shape[0]
						if 0 < observed_batch_size < self.loader_batch_size:
							# could be last batch so we can't unroll as many
							# elements.
							__snake_case    :							Optional[Any]   	=				observed_batch_size
						# Setting internal index to unwrap the batch
						__snake_case    :							Optional[Any]   	=				processed
						__snake_case    :							List[Any]   	=				0
						return self.loader_batch_item()
					else:
						# We're not unrolling batches
						return processed







class 	UpperCamelCase    (					lowerCAmelCase__		):



				def __init__(self       :   int			,						_A       :   Any			,						_A       :   Optional[Any]			,						_A       :   List[Any]			,						_A       :   str=None)					->			List[Any]:
					super().__init__(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE)



				def __iter__(self       :   List[str])					->			int:
					__snake_case    :							Dict   	=				iter(self.loader)
					__snake_case    :							str   	=				None
					return self






				def       _lowercase  (self       :   str)					->			Optional[Any]:
					if self.subiterator is None:
						__snake_case    :							Any   	=				self.infer(next(self.iterator)			,						**self.params)
					try:
						# Try to return next item
						__snake_case    :							Union[str, Any]   	=				next(self.subiterator)
					except StopIteration:
						# When a preprocess iterator ends, we can start lookig at the next item
						# ChunkIterator will keep feeding until ALL elements of iterator
						# all have created their subiterator and have been iterating against.
						#
						# Another way to look at it, is we're basically flattening lists of lists
						# into a single list, but with generators
						__snake_case    :							Optional[Any]   	=				self.infer(next(self.iterator)			,						**self.params)
						__snake_case    :							List[Any]   	=				next(self.subiterator)
					return processed







class 	UpperCamelCase    (					lowerCAmelCase__		):



				def __iter__(self       :   int)					->			Any:
					__snake_case    :							int   	=				iter(self.loader)
					return self






				def       _lowercase  (self       :   List[str])					->			List[str]:
					# Extremely similar to PipelineIterator in its unpacking mechanism
					# BUT, we have an extra required item which is the presence of `is_last`
					# That is because everything is flattened by `PipelineChunkIterator` we
					# need to keep track of how to regroup here in the original `process`
					# boundaries so that `process` and `postprocess` see the same data.

					# This iterator accumulates items (possibly while unbatching) until it
					# its a `is_last` and then just passes it on to the caller.
					__snake_case    :							Optional[Any]   	=				False
					__snake_case    :							Dict   	=				[]
					if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
						while self._loader_batch_index < self.loader_batch_size:
							__snake_case    :							Union[str, Any]   	=				self.loader_batch_item()
							__snake_case    :							List[str]   	=				item.pop('is_last')
							accumulator.append(_SCREAMING_SNAKE_CASE)
							if is_last:
								return accumulator

					while not is_last:
						__snake_case    :							int   	=				self.infer(next(self.iterator)			,						**self.params)
						if self.loader_batch_size is not None:
							if isinstance(_SCREAMING_SNAKE_CASE			,						torch.Tensor):
								__snake_case    :							List[str]   	=				processed
							else:
								__snake_case    :							str   	=				list(processed.keys())[0]
								__snake_case    :							Optional[int]   	=				processed[key]
							if isinstance(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE):
								__snake_case    :							str   	=				len(_SCREAMING_SNAKE_CASE)
							else:
								__snake_case    :							Optional[int]   	=				first_tensor.shape[0]
							if 0 < observed_batch_size < self.loader_batch_size:
								# could be last batch so we can't unroll as many
								# elements.
								__snake_case    :							int   	=				observed_batch_size
							__snake_case    :							Dict   	=				processed
							__snake_case    :							List[str]   	=				0
							while self._loader_batch_index < self.loader_batch_size:
								__snake_case    :							Optional[Any]   	=				self.loader_batch_item()
								__snake_case    :							Any   	=				item.pop('is_last')
								accumulator.append(_SCREAMING_SNAKE_CASE)
								if is_last:
									return accumulator
						else:
							__snake_case    :							Optional[int]   	=				processed
							__snake_case    :							Tuple   	=				item.pop('is_last')
							accumulator.append(_SCREAMING_SNAKE_CASE)
					return accumulator







class 	UpperCamelCase    (					lowerCAmelCase__		):



				def __init__(self       :   Union[str, Any]			,						_A       :   Optional[int]			,						_A       :   Dict)					->			List[str]:
					__snake_case    :							int   	=				dataset
					__snake_case    :							int   	=				key



				def __len__(self       :   str)					->			Optional[int]:
					return len(self.dataset)






				def __getitem__(self       :   Dict			,						_A       :   Any)					->			Any:
					return self.dataset[i][self.key]







class 	UpperCamelCase    (					lowerCAmelCase__		):



				def __init__(self       :   Optional[int]			,						_A       :   Dict			,						_A       :   int			,						_A       :   Optional[Any])					->			Any:
					__snake_case    :							Any   	=				dataset
					__snake_case    :							List[Any]   	=				keya
					__snake_case    :							Union[str, Any]   	=				keya



				def __len__(self       :   Tuple)					->			List[Any]:
					return len(self.dataset)






				def __getitem__(self       :   Any			,						_A       :   Optional[int])					->			Union[str, Any]:
					return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}






"""simple docstring"""







import math

import tensorflow as tf
from packaging import version
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]						)							->    Any:

	'''simple docstring'''





	__snake_case    :							List[str]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							List[str]   	=				0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0						)    ,							x.dtype						)						))

	return x * cdf
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    Optional[Any]:

	'''simple docstring'''





	__snake_case    :							Optional[int]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Any   	=				tf.cast(math.pi    ,							x.dtype						)
	__snake_case    :							List[Any]   	=				tf.cast(0.044_715    ,							x.dtype						)
	__snake_case    :							Union[str, Any]   	=				0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi						) * (x + coeff * tf.pow(UpperCAmelCase_    ,							3						))						))

	return x * cdf
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    List[Any]:

	'''simple docstring'''





	__snake_case    :							Optional[int]   	=				tf.convert_to_tensor(UpperCAmelCase_						)

	return x * tf.tanh(tf.math.softplus(UpperCAmelCase_						)						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    List[str]:

	'''simple docstring'''





	__snake_case    :							Union[str, Any]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Dict   	=				tf.cast(0.044_715    ,							x.dtype						)
	__snake_case    :							int   	=				tf.cast(0.7_978_845_608    ,							x.dtype						)

	return 0.5 * x * (1.0 + tf.tanh(x * coeffa * (1.0 + coeffa * x * x)						))
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       int						)							->    Dict:

	'''simple docstring'''





	__snake_case    :							Any   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Dict   	=				tf.cast(1.702    ,							x.dtype						)
	return x * tf.math.sigmoid(coeff * x						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]						)							->    List[Any]:

	'''simple docstring'''





	return tf.clip_by_value(_gelu(UpperCAmelCase_						)    ,							-10    ,							10						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]    ,							UpperCAmelCase_					:       Any=-1						)							->    int:

	'''simple docstring'''





	__snake_case	,       __snake_case    :							str   	=				tf.split(UpperCAmelCase_    ,							2    ,							axis=UpperCAmelCase_						)
	return a * tf.math.sigmoid(UpperCAmelCase_						)


if version.parse(tf.version.VERSION) >= version.parse("2.4"):
	def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    int:

		'''simple docstring'''





		return tf.keras.activations.gelu(UpperCAmelCase_    ,							approximate=UpperCAmelCase_						)

	_a  :  Tuple=				tf.keras.activations.gelu

	_a  :  Tuple=				approximate_gelu_wrap
else:
	_a  :  Any=				_gelu
	_a  :  str=				_gelu_new


_a  :  List[str]=				{
    "gelu": gelu,
    "gelu_10": gelu_aa,
    "gelu_fast": gelu_fast,
    "gelu_new": gelu_new,
    "glu": glu,
    "mish": mish,
    "quick_gelu": quick_gelu,
    "relu": tf.keras.activations.relu,
    "sigmoid": tf.keras.activations.sigmoid,
    "silu": tf.keras.activations.swish,
    "swish": tf.keras.activations.swish,
    "tanh": tf.keras.activations.tanh,
}

def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Optional[Any]						)							->    List[Any]:

	'''simple docstring'''





	if activation_string in ACTaFN:
		return ACTaFN[activation_string]
	else:
		raise KeyError(F"function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()						)}"						)





