



import unittest

from transformers import ErnieMTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


UpperCAmelCase							=       get_tests_dir("""fixtures/spiece.model""")

@require_sentencepiece
@require_tokenizers
class 			__snake_case   ( SCREAMING_SNAKE_CASE      ,unittest.TestCase):







							'''simple docstring'''






							UpperCamelCase__					:   List[Any]       =       ErnieMTokenizer
							UpperCamelCase__					:   List[Any]       =       False
							UpperCamelCase__					:   List[Any]       =       True
							UpperCamelCase__					:   int       =       False
							UpperCamelCase__					:   Union[str, Any]       =       False


							def _a   (				self		):
											super().setUp()

											# We have a SentencePiece fixture for testing
											a__					   =		ErnieMTokenizer(a_	,	unk_token="""<unk>"""	,	pad_token="""<pad>"""		)
											tokenizer.save_pretrained(self.tmpdirname		)


							def _a   (				self	,	a_		):
											a__					   =		"""this is a test"""
											a__					   =		"""this is a test"""
											return input_text, output_text


							def _a   (				self		):
											a__					   =		"""<pad>"""
											a__					   =		0

											self.assertEqual(self.get_tokenizer()._convert_token_to_id(a_		)	,	a_		)
											self.assertEqual(self.get_tokenizer()._convert_id_to_token(a_		)	,	a_		)


							def _a   (				self		):
											a__					   =		list(self.get_tokenizer().get_vocab().keys()		)

											self.assertEqual(vocab_keys[0]	,	"""<pad>"""		)
											self.assertEqual(vocab_keys[1]	,	"""<unk>"""		)
											self.assertEqual(vocab_keys[-1]	,	"""▁eloquent"""		)
											self.assertEqual(len(a_		)	,	30_000		)


							def _a   (				self		):
											self.assertEqual(self.get_tokenizer().vocab_size	,	30_000		)


							def _a   (				self		):
											if not self.test_rust_tokenizer:
															return

											a__					   =		self.get_tokenizer()
											a__					   =		self.get_rust_tokenizer()

											a__					   =		"""I was born in 92000, and this is falsé."""

											a__					   =		tokenizer.tokenize(a_		)
											a__					   =		rust_tokenizer.tokenize(a_		)
											self.assertListEqual(a_	,	a_		)

											a__					   =		tokenizer.encode(a_	,	add_special_tokens=a_		)
											a__					   =		rust_tokenizer.encode(a_	,	add_special_tokens=a_		)
											self.assertListEqual(a_	,	a_		)

											a__					   =		self.get_rust_tokenizer()
											a__					   =		tokenizer.encode(a_		)
											a__					   =		rust_tokenizer.encode(a_		)
											self.assertListEqual(a_	,	a_		)


							def _a   (				self		):
											a__					   =		ErnieMTokenizer(a_	,	do_lower_case=a_	,	unk_token="""<unk>"""	,	pad_token="""<pad>"""		)

											a__					   =		tokenizer.tokenize("""This is a test"""		)
											self.assertListEqual(a_	,	["""▁this""", """▁is""", """▁a""", """▁test"""]		)

											self.assertListEqual(tokenizer.convert_tokens_to_ids(a_		)	,	[48, 25, 21, 1_289]		)

											a__					   =		tokenizer.tokenize("""I was born in 92000, and this is falsé."""		)
											# ErnieMTokenizer(paddlenlp implementation) outputs '9' instead of '_9' so to mimic that '_9' is changed to '9'
											self.assertListEqual(
											    a_	,	["""▁i""", """▁was""", """▁born""", """▁in""", """9""", """2000""", """,""", """▁and""", """▁this""", """▁is""", """▁fal""", """s""", """é""", """."""]		)
											a__					   =		tokenizer.convert_tokens_to_ids(a_		)
											self.assertListEqual(a_	,	[31, 23, 386, 19, 518, 3_050, 15, 17, 48, 25, 8_256, 18, 1, 9]		)

											a__					   =		tokenizer.convert_ids_to_tokens(a_		)
											self.assertListEqual(
											    a_	,	["""▁i""", """▁was""", """▁born""", """▁in""", """9""", """2000""", """,""", """▁and""", """▁this""", """▁is""", """▁fal""", """s""", """<unk>""", """."""]	,	)


							def _a   (				self		):
											a__					   =		ErnieMTokenizer(a_	,	unk_token="""<unk>"""	,	pad_token="""<pad>"""		)

											a__					   =		tokenizer.encode("""sequence builders"""		)
											a__					   =		tokenizer.encode("""multi-sequence build"""		)

											a__					   =		tokenizer.build_inputs_with_special_tokens(a_		)
											a__					   =		tokenizer.build_inputs_with_special_tokens(a_	,	a_		)

											assert encoded_sentence == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]
											assert encoded_pair == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + [
											    tokenizer.sep_token_id
											] + text_a + [tokenizer.sep_token_id]







							@slow
							def _a   (				self		):
											# fmt: off
											a__					   =		{"""input_ids""": [[0, 11_062, 82_772, 7, 15, 82_772, 538, 51_529, 237, 17_198, 1_290, 206, 9, 215_175, 1_314, 136, 17_198, 1_290, 206, 9, 56_359, 42, 122_009, 9, 16_466, 16, 87_344, 4_537, 9, 4_717, 78_381, 6, 159_958, 7, 15, 24_480, 618, 4, 527, 22_693, 9, 304, 4, 2_777, 24_480, 9_874, 4, 43_523, 594, 4, 803, 18_392, 33_189, 18, 4, 43_523, 24_447, 5, 5, 5, 16, 100, 24_955, 83_658, 9_626, 144_057, 15, 839, 22_335, 16, 136, 24_955, 83_658, 83_479, 15, 39_102, 724, 16, 678, 645, 6_460, 1_328, 4_589, 42, 122_009, 115_774, 23, 3_559, 1_328, 46_876, 7, 136, 53_894, 1_940, 42_227, 41_159, 17_721, 823, 425, 4, 27_512, 98_722, 206, 136, 5_531, 4_970, 919, 17_336, 5, 2], [0, 20_080, 618, 83, 82_775, 47, 479, 9, 1_517, 73, 53_894, 333, 80_581, 110_117, 18_811, 5_256, 1_295, 51, 152_526, 297, 7_986, 390, 124_416, 538, 35_431, 214, 98, 15_044, 25_737, 136, 7_108, 43_701, 23, 756, 135_355, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 581, 63_773, 119_455, 6, 147_797, 88_203, 7, 645, 70, 21, 3_285, 10_269, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}
											# fmt: on

											self.tokenizer_integration_test_util(
											    expected_encoding=a_	,	model_name="""susnato/ernie-m-base_pytorch"""	,	sequences=[
											        """Transformers (formerly known as pytorch-transformers and pytorch-pretrained-bert) provides """
											        """general-purpose architectures (BERT, GPT-2, RoBERTa, XLM, DistilBert, XLNet...) for Natural """
											        """Language Understanding (NLU) and Natural Language Generation (NLG) with over32+ pretrained """
											        """models in100+ languages and deep interoperability between Jax, PyTorch and TensorFlow.""",
											        """BERT is designed to pre-train deep bidirectional representations from unlabeled text by jointly """
											        """conditioning on both left and right context in all layers.""",
											        """The quick brown fox jumps over the lazy dog.""",
											    ]	,	)






import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp

def       A_   (						__a  :					str			,		__a  :					int			,		__a  :					List[Any]			,		__a  :					Dict			,		__a  :					int = 4_096		):

				"""simple docstring"""
				a__   ,      a__   ,      a__					   =		key.shape[-3:]
				a__					   =		value.shape[-1]
				a__					   =		min(__a			,		__a		)
				a__					   =		query / jnp.sqrt(__a		)

				@functools.partial(jax.checkpoint			,		prevent_cse=__a		)
				def summarize_chunk(__a  :					Optional[Any]			,		__a  :					Tuple			,		__a  :					Dict		):
								a__					   =		jnp.einsum("""...qhd,...khd->...qhk"""			,		__a			,		__a			,		precision=__a		)

								a__					   =		jnp.max(__a			,		axis=-1			,		keepdims=__a		)
								a__					   =		jax.lax.stop_gradient(__a		)
								a__					   =		jnp.exp(attn_weights - max_score		)

								a__					   =		jnp.einsum("""...vhf,...qhv->...qhf"""			,		__a			,		__a			,		precision=__a		)
								a__					   =		jnp.einsum("""...qhk->...qh"""			,		__a		)

								return (exp_values, exp_weights.sum(axis=-1		), max_score)

				def chunk_scanner(__a  :					Tuple		):
								# julienne key array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]			,		slice_sizes=list(key.shape[:-3]		) + [key_chunk_size, num_heads, k_features]			,		)

								# julienne value array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]			,		slice_sizes=list(value.shape[:-3]		) + [key_chunk_size, num_heads, v_features]			,		)

								return summarize_chunk(__a			,		__a			,		__a		)

				a__   ,      a__   ,      a__					   =		jax.lax.map(f=__a			,		xs=jnp.arange(0			,		__a			,		__a		)		)

				a__					   =		jnp.max(__a			,		axis=0			,		keepdims=__a		)
				a__					   =		jnp.exp(chunk_max - global_max		)

				chunk_values *= jnp.expand_dims(__a			,		axis=-1		)
				chunk_weights *= max_diffs

				a__					   =		chunk_values.sum(axis=0		)
				a__					   =		jnp.expand_dims(__a			,		-1		).sum(axis=0		)

				return all_values / all_weights

def       A_   (						__a  :					Tuple			,		__a  :					List[Any]			,		__a  :					int			,		__a  :					List[str]=jax.lax.Precision.HIGHEST			,		__a  :					int = 1_024			,		__a  :					int = 4_096		):

				"""simple docstring"""
				a__   ,      a__   ,      a__					   =		query.shape[-3:]

				def chunk_scanner(__a  :					List[str]			,		__a  :					List[str]		):
								# julienne query array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]			,		slice_sizes=list(query.shape[:-3]		) + [min(__a			,		__a		), num_heads, q_features]			,		)

								return (
								    chunk_idx + query_chunk_size,  # unused ignore it
								    _query_chunk_attention(
								        query=__a			,		key=__a			,		value=__a			,		precision=__a			,		key_chunk_size=__a		),
								)

				a__   ,      a__					   =		jax.lax.scan(
				    f=__a			,		init=0			,		xs=__a			,		length=math.ceil(num_q / query_chunk_size		)		)

				return jnp.concatenate(__a			,		axis=-3		)  # fuse the chunked result back

class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int       =       8
							UpperCamelCase__					:   int       =       6_4
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											a__					   =		self.dim_head * self.heads
											a__					   =		self.dim_head**-0.5

											# Weights were exported with old names {to_q, to_k, to_v, to_out}
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_q"""		)
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_k"""		)
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_v"""		)

											a__					   =		nn.Dense(self.query_dim	,	dtype=self.dtype	,	name="""to_out_0"""		)
											a__					   =		nn.Dropout(rate=self.dropout		)


							def _a   (				self	,	a_		):
											a__   ,      a__   ,      a__					   =		tensor.shape
											a__					   =		self.heads
											a__					   =		tensor.reshape(a_	,	a_	,	a_	,	dim // head_size		)
											a__					   =		jnp.transpose(a_	,	(0, 2, 1, 3)		)
											a__					   =		tensor.reshape(batch_size * head_size	,	a_	,	dim // head_size		)
											return tensor


							def _a   (				self	,	a_		):
											a__   ,      a__   ,      a__					   =		tensor.shape
											a__					   =		self.heads
											a__					   =		tensor.reshape(batch_size // head_size	,	a_	,	a_	,	a_		)
											a__					   =		jnp.transpose(a_	,	(0, 2, 1, 3)		)
											a__					   =		tensor.reshape(batch_size // head_size	,	a_	,	dim * head_size		)
											return tensor







							def __call__(				self	,	a_	,	a_=None	,	a_=True		):
											a__					   =		hidden_states if context is None else context

											a__					   =		self.query(a_		)
											a__					   =		self.key(a_		)
											a__					   =		self.value(a_		)

											a__					   =		self.reshape_heads_to_batch_dim(a_		)
											a__					   =		self.reshape_heads_to_batch_dim(a_		)
											a__					   =		self.reshape_heads_to_batch_dim(a_		)

											if self.use_memory_efficient_attention:
															a__					   =		query_states.transpose(1	,	0	,	2		)
															a__					   =		key_states.transpose(1	,	0	,	2		)
															a__					   =		value_states.transpose(1	,	0	,	2		)

															# this if statement create a chunk size for each layer of the unet
															# the chunk size is equal to the query_length dimension of the deepest layer of the unet

															a__					   =		query_states.shape[-3]
															if flatten_latent_dim % 64 == 0:
																			a__					   =		int(flatten_latent_dim / 64		)
															elif flatten_latent_dim % 16 == 0:
																			a__					   =		int(flatten_latent_dim / 16		)
															elif flatten_latent_dim % 4 == 0:
																			a__					   =		int(flatten_latent_dim / 4		)
															else:
																			a__					   =		int(a_		)

															a__					   =		jax_memory_efficient_attention(
															    a_	,	a_	,	a_	,	query_chunk_size=a_	,	key_chunk_size=4_096 * 4		)

															a__					   =		hidden_states.transpose(1	,	0	,	2		)
											else:
															# compute attentions
															a__					   =		jnp.einsum("""b i d, b j d->b i j"""	,	a_	,	a_		)
															a__					   =		attention_scores * self.scale
															a__					   =		nn.softmax(a_	,	axis=2		)

															# attend to values
															a__					   =		jnp.einsum("""b i j, b j d -> b i d"""	,	a_	,	a_		)

											a__					   =		self.reshape_batch_dim_to_heads(a_		)
											a__					   =		self.proj_attn(a_		)
											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa
							UpperCamelCase__					:   bool       =       False


							def _a   (				self		):
											# self attention (or cross_attention if only_cross_attention is True)
											a__					   =		FlaxAttention(
											    self.dim	,	self.n_heads	,	self.d_head	,	self.dropout	,	self.use_memory_efficient_attention	,	dtype=self.dtype		)
											# cross attention
											a__					   =		FlaxAttention(
											    self.dim	,	self.n_heads	,	self.d_head	,	self.dropout	,	self.use_memory_efficient_attention	,	dtype=self.dtype		)
											a__					   =		FlaxFeedForward(dim=self.dim	,	dropout=self.dropout	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_	,	a_=True		):
											# self attention
											a__					   =		hidden_states
											if self.only_cross_attention:
															a__					   =		self.attna(self.norma(a_		)	,	a_	,	deterministic=a_		)
											else:
															a__					   =		self.attna(self.norma(a_		)	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											# cross attention
											a__					   =		hidden_states
											a__					   =		self.attna(self.norma(a_		)	,	a_	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											# feed forward
											a__					   =		hidden_states
											a__					   =		self.ff(self.norma(a_		)	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int       =       1
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa
							UpperCamelCase__					:   bool       =       False


							def _a   (				self		):
											a__					   =		nn.GroupNorm(num_groups=32	,	epsilon=1E-5		)

											a__					   =		self.n_heads * self.d_head
											if self.use_linear_projection:
															a__					   =		nn.Dense(a_	,	dtype=self.dtype		)
											else:
															a__					   =		nn.Conv(
															    a_	,	kernel_size=(1, 1)	,	strides=(1, 1)	,	padding="""VALID"""	,	dtype=self.dtype	,	)

											a__					   =		[
											    FlaxBasicTransformerBlock(
											        a_	,	self.n_heads	,	self.d_head	,	dropout=self.dropout	,	only_cross_attention=self.only_cross_attention	,	dtype=self.dtype	,	use_memory_efficient_attention=self.use_memory_efficient_attention	,	)
											    for _ in range(self.depth		)
											]

											if self.use_linear_projection:
															a__					   =		nn.Dense(a_	,	dtype=self.dtype		)
											else:
															a__					   =		nn.Conv(
															    a_	,	kernel_size=(1, 1)	,	strides=(1, 1)	,	padding="""VALID"""	,	dtype=self.dtype	,	)

											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_	,	a_=True		):
											a__   ,      a__   ,      a__   ,      a__					   =		hidden_states.shape
											a__					   =		hidden_states
											a__					   =		self.norm(a_		)
											if self.use_linear_projection:
															a__					   =		hidden_states.reshape(a_	,	height * width	,	a_		)
															a__					   =		self.proj_in(a_		)
											else:
															a__					   =		self.proj_in(a_		)
															a__					   =		hidden_states.reshape(a_	,	height * width	,	a_		)

											for transformer_block in self.transformer_blocks:
															a__					   =		transformer_block(a_	,	a_	,	deterministic=a_		)

											if self.use_linear_projection:
															a__					   =		self.proj_out(a_		)
															a__					   =		hidden_states.reshape(a_	,	a_	,	a_	,	a_		)
											else:
															a__					   =		hidden_states.reshape(a_	,	a_	,	a_	,	a_		)
															a__					   =		self.proj_out(a_		)

											a__					   =		hidden_states + residual
											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											# The second linear layer needs to be called
											# net_2 for now to match the index of the Sequential layer
											a__					   =		FlaxGEGLU(self.dim	,	self.dropout	,	self.dtype		)
											a__					   =		nn.Dense(self.dim	,	dtype=self.dtype		)







							def __call__(				self	,	a_	,	a_=True		):
											a__					   =		self.net_a(a_	,	deterministic=a_		)
											a__					   =		self.net_a(a_		)
											return hidden_states




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											a__					   =		self.dim * 4
											a__					   =		nn.Dense(inner_dim * 2	,	dtype=self.dtype		)
											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_=True		):
											a__					   =		self.proj(a_		)
											a__   ,      a__					   =		jnp.split(a_	,	2	,	axis=2		)
											return self.dropout_layer(hidden_linear * nn.gelu(a_		)	,	deterministic=a_		)

