



import json
import os
import re
import unicodedata
from typing import Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
from ...utils import logging


UpperCamelCase_	       =   logging.get_logger(__name__)

UpperCamelCase_	       =   {
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
}

UpperCamelCase_	       =   {
    "vocab_file": {"openai-gpt": "https://huggingface.co/openai-gpt/resolve/main/vocab.json"},
    "merges_file": {"openai-gpt": "https://huggingface.co/openai-gpt/resolve/main/merges.txt"},
}

UpperCamelCase_	       =   {
    "openai-gpt": 512,
}







def 				A			(		__UpperCAmelCase	)       ->       Optional[int]:



						'''simple docstring'''






						UpperCAmelCase_     							=     text.strip()
						if not text:
												return []
						UpperCAmelCase_     							=     text.split()
						return tokens







class 							a_  (		_snake_case				):

						def __init__(    self	:Optional[int]   ,						_lowercase	:Any=True   ,						_lowercase	:Tuple=None   ,						_lowercase	:str=True   ,						_lowercase	:Tuple=None   ,						_lowercase	:Tuple=True   ,						)		->					str:
												if never_split is None:
																		UpperCAmelCase_     							=     []
												UpperCAmelCase_     							=     do_lower_case
												UpperCAmelCase_     							=     set(_lowercase)
												UpperCAmelCase_     							=     tokenize_chinese_chars
												UpperCAmelCase_     							=     strip_accents
												UpperCAmelCase_     							=     do_split_on_punc

						def 	__a       (    self	:str   ,						_lowercase	:int   ,						_lowercase	:Union[str, Any]=None)		->					List[str]:
												UpperCAmelCase_     							=     self.never_split.union(set(_lowercase)) if never_split else self.never_split
												UpperCAmelCase_     							=     self._clean_text(_lowercase)

												# This was added on November 1st, 2018 for the multilingual and Chinese
												# models. This is also applied to the English models now, but it doesn't
												# matter since the English models were not trained on any Chinese data
												# and generally don't have any Chinese data in them (there are Chinese
												# characters in the vocabulary because Wikipedia does have some Chinese
												# words in the English Wikipedia.).
												if self.tokenize_chinese_chars:
																		UpperCAmelCase_     							=     self._tokenize_chinese_chars(_lowercase)
												# prevents treating the same character with different unicode codepoints as different characters
												UpperCAmelCase_     							=     unicodedata.normalize('''NFC'''   ,						_lowercase)
												UpperCAmelCase_     							=     whitespace_tokenize(_lowercase)
												UpperCAmelCase_     							=     []
												for token in orig_tokens:
																		if token not in never_split:
																								if self.do_lower_case:
																														UpperCAmelCase_     							=     token.lower()
																														if self.strip_accents is not False:
																																				UpperCAmelCase_     							=     self._run_strip_accents(_lowercase)
																								elif self.strip_accents:
																														UpperCAmelCase_     							=     self._run_strip_accents(_lowercase)
																		split_tokens.extend(self._run_split_on_punc(_lowercase   ,						_lowercase))

												UpperCAmelCase_     							=     whitespace_tokenize(''' '''.join(_lowercase))
												return output_tokens

						def 	__a       (    self	:Optional[Any]   ,						_lowercase	:Optional[int])		->					Tuple:
												UpperCAmelCase_     							=     unicodedata.normalize('''NFD'''   ,						_lowercase)
												UpperCAmelCase_     							=     []
												for char in text:
																		UpperCAmelCase_     							=     unicodedata.category(_lowercase)
																		if cat == "Mn":
																								continue
																		output.append(_lowercase)
												return "".join(_lowercase)

						def 	__a       (    self	:Dict   ,						_lowercase	:List[Any]   ,						_lowercase	:int=None)		->					str:
												if not self.do_split_on_punc or (never_split is not None and text in never_split):
																		return [text]
												UpperCAmelCase_     							=     list(_lowercase)
												UpperCAmelCase_     							=     0
												UpperCAmelCase_     							=     True
												UpperCAmelCase_     							=     []
												while i < len(_lowercase):
																		UpperCAmelCase_     							=     chars[i]
																		if _is_punctuation(_lowercase):
																								output.append([char])
																								UpperCAmelCase_     							=     True
																		else:
																								if start_new_word:
																														output.append([])
																								UpperCAmelCase_     							=     False
																								output[-1].append(_lowercase)
																		i += 1

												return ["".join(_lowercase) for x in output]

						def 	__a       (    self	:Union[str, Any]   ,						_lowercase	:int)		->					Optional[int]:
												UpperCAmelCase_     							=     []
												for char in text:
																		UpperCAmelCase_     							=     ord(_lowercase)
																		if self._is_chinese_char(_lowercase):
																								output.append(''' ''')
																								output.append(_lowercase)
																								output.append(''' ''')
																		else:
																								output.append(_lowercase)
												return "".join(_lowercase)

						def 	__a       (    self	:Optional[Any]   ,						_lowercase	:Any)		->					Any:
												# This defines a "chinese character" as anything in the CJK Unicode block:
												#   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
												#
												# Note that the CJK Unicode block is NOT all Japanese and Korean characters,
												# despite its name. The modern Korean Hangul alphabet is a different block,
												# as is Japanese Hiragana and Katakana. Those alphabets are used to write
												# space-separated words, so they are not treated specially and handled
												# like the all of the other languages.
												if (
												    (cp >= 0x4e00 and cp <= 0x9fff)
												    or (cp >= 0x3400 and cp <= 0x4dbf)  #
												    or (cp >= 0x2_0000 and cp <= 0x2_a6df)  #
												    or (cp >= 0x2_a700 and cp <= 0x2_b73f)  #
												    or (cp >= 0x2_b740 and cp <= 0x2_b81f)  #
												    or (cp >= 0x2_b820 and cp <= 0x2_ceaf)  #
												    or (cp >= 0xf900 and cp <= 0xfaff)
												    or (cp >= 0x2_f800 and cp <= 0x2_fa1f)  #
												):  #
																		return True

												return False




						def 	__a       (    self	:str   ,						_lowercase	:Tuple)		->					List[str]:
												UpperCAmelCase_     							=     []
												for char in text:
																		UpperCAmelCase_     							=     ord(_lowercase)
																		if cp == 0 or cp == 0xfffd or _is_control(_lowercase):
																								continue
																		if _is_whitespace(_lowercase):
																								output.append(''' ''')
																		else:
																								output.append(_lowercase)
												return "".join(_lowercase)







def 				A			(		__UpperCAmelCase	)       ->       List[Any]:



						'''simple docstring'''






						UpperCAmelCase_     							=     set()
						UpperCAmelCase_     							=     word[0]
						for char in word[1:]:
												pairs.add((prev_char, char)	)
												UpperCAmelCase_     							=     char
						return pairs







def 				A			(		__UpperCAmelCase	)       ->       List[str]:



						'''simple docstring'''






						UpperCAmelCase_     							=     text.replace('''—'''							,					'''-'''	)
						UpperCAmelCase_     							=     text.replace('''–'''							,					'''-'''	)
						UpperCAmelCase_     							=     text.replace('''―'''							,					'''-'''	)
						UpperCAmelCase_     							=     text.replace('''…'''							,					'''...'''	)
						UpperCAmelCase_     							=     text.replace('''´'''							,					'''\''''	)
						UpperCAmelCase_     							=     re.sub(r'''(-+|~+|!+|"+|;+|\?+|\++|,+|\)+|\(+|\\+|\/+|\*+|\[+|\]+|}+|{+|\|+|_+)'''							,					r''' \1 '''							,					__UpperCAmelCase	)
						UpperCAmelCase_     							=     re.sub(r'''\s*\n\s*'''							,					''' \n '''							,					__UpperCAmelCase	)
						UpperCAmelCase_     							=     re.sub(r'''[^\S\n]+'''							,					''' '''							,					__UpperCAmelCase	)
						return text.strip()







class 							a_  (		_snake_case				):
						UpperCamelCase__      :						Union[str, Any]				       =VOCAB_FILES_NAMES
						UpperCamelCase__      :						Tuple				       =PRETRAINED_VOCAB_FILES_MAP
						UpperCamelCase__      :						List[str]				       =PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						UpperCamelCase__      :						Union[str, Any]				       =["input_ids", "attention_mask"]

						def __init__(    self	:Dict   ,						_lowercase	:int   ,						_lowercase	:Union[str, Any]   ,						_lowercase	:List[str]="<unk>"   ,						**_lowercase	:Tuple)		->					List[str]:
												super().__init__(unk_token=_lowercase   ,						**_lowercase)

												try:
																		import ftfy
																		from spacy.lang.en import English

																		UpperCAmelCase_     							=     English()
																		UpperCAmelCase_     							=     _nlp.tokenizer
																		UpperCAmelCase_     							=     ftfy.fix_text
												except ImportError:
																		logger.warning('''ftfy or spacy is not installed using BERT BasicTokenizer instead of SpaCy & ftfy.''')
																		UpperCAmelCase_     							=     BasicTokenizer(do_lower_case=_lowercase)
																		UpperCAmelCase_     							=     None

												with open(_lowercase   ,						encoding='''utf-8''') as vocab_handle:
																		UpperCAmelCase_     							=     json.load(_lowercase)
												UpperCAmelCase_     							=     {v: k for k, v in self.encoder.items()}
												with open(_lowercase   ,						encoding='''utf-8''') as merges_handle:
																		UpperCAmelCase_     							=     merges_handle.read().split('''\n''')[1:-1]
												UpperCAmelCase_     							=     [tuple(merge.split()) for merge in merges]
												UpperCAmelCase_     							=     dict(zip(_lowercase   ,						range(len(_lowercase))))
												UpperCAmelCase_     							=     {}

						@property
						def 	__a       (    self	:int)		->					int:
												return True

						@property
						def 	__a       (    self	:Optional[int])		->					Tuple:
												return len(self.encoder)

						def 	__a       (    self	:Optional[int])		->					Optional[Any]:
												return dict(self.encoder   ,						**self.added_tokens_encoder)

						def 	__a       (    self	:Any   ,						_lowercase	:Any)		->					Union[str, Any]:
												UpperCAmelCase_     							=     tuple(token[:-1]) + (token[-1] + '''</w>''',)
												if token in self.cache:
																		return self.cache[token]
												UpperCAmelCase_     							=     get_pairs(_lowercase)

												if not pairs:
																		return token + "</w>"

												while True:
																		UpperCAmelCase_     							=     min(_lowercase   ,						key=lambda _lowercase: self.bpe_ranks.get(_lowercase   ,						float('''inf''')))
																		if bigram not in self.bpe_ranks:
																								break
																		UpperCAmelCase_							,      UpperCAmelCase_     							=     bigram
																		UpperCAmelCase_     							=     []
																		UpperCAmelCase_     							=     0
																		while i < len(_lowercase):
																								try:
																														UpperCAmelCase_     							=     word.index(_lowercase   ,						_lowercase)
																								except ValueError:
																														new_word.extend(word[i:])
																														break
																								else:
																														new_word.extend(word[i:j])
																														UpperCAmelCase_     							=     j

																								if word[i] == first and i < len(_lowercase) - 1 and word[i + 1] == second:
																														new_word.append(first + second)
																														i += 2
																								else:
																														new_word.append(word[i])
																														i += 1
																		UpperCAmelCase_     							=     tuple(_lowercase)
																		UpperCAmelCase_     							=     new_word
																		if len(_lowercase) == 1:
																								break
																		else:
																								UpperCAmelCase_     							=     get_pairs(_lowercase)
												UpperCAmelCase_     							=     ''' '''.join(_lowercase)
												if word == "\n  </w>":
																		UpperCAmelCase_     							=     '''\n</w>'''
												UpperCAmelCase_     							=     word
												return word

						def 	__a       (    self	:Optional[Any]   ,						_lowercase	:Any)		->					List[str]:
												UpperCAmelCase_     							=     []
												if self.fix_text is None:
																		# Using BERT's BasicTokenizer
																		UpperCAmelCase_     							=     self.nlp.tokenize(_lowercase)
																		for token in text:
																								split_tokens.extend(list(self.bpe(_lowercase).split(''' ''')))
												else:
																		# Using SpaCy & ftfy (original tokenization process of OpenAI GPT)
																		UpperCAmelCase_     							=     self.nlp(text_standardize(self.fix_text(_lowercase)))
																		for token in text:
																								split_tokens.extend(list(self.bpe(token.text.lower()).split(''' ''')))
												return split_tokens

						def 	__a       (    self	:Tuple   ,						_lowercase	:int)		->					List[str]:
												return self.encoder.get(_lowercase   ,						self.encoder.get(self.unk_token))

						def 	__a       (    self	:Dict   ,						_lowercase	:Dict)		->					Union[str, Any]:
												return self.decoder.get(_lowercase   ,						self.unk_token)

						def 	__a       (    self	:Tuple   ,						_lowercase	:str)		->					int:
												UpperCAmelCase_     							=     ''''''.join(_lowercase).replace('''</w>'''   ,						''' ''').strip()
												return out_string




						def 	__a       (    self	:Optional[int]   ,						_lowercase	:str   ,						_lowercase	:Optional[str] = None)		->					Tuple[str]:
												if not os.path.isdir(_lowercase):
																		logger.error(f"Vocabulary path ({save_directory}) should be a directory")
																		return
												UpperCAmelCase_     							=     os.path.join(
												    _lowercase   ,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file'''])
												UpperCAmelCase_     							=     os.path.join(
												    _lowercase   ,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file'''])

												with open(_lowercase   ,						'''w'''   ,						encoding='''utf-8''') as f:
																		f.write(json.dumps(self.encoder   ,						indent=2   ,						sort_keys=_lowercase   ,						ensure_ascii=_lowercase) + '''\n''')

												UpperCAmelCase_     							=     0
												with open(_lowercase   ,						'''w'''   ,						encoding='''utf-8''') as writer:
																		writer.write('''#version: 0.2\n''')
																		for bpe_tokens, token_index in sorted(self.bpe_ranks.items()   ,						key=lambda _lowercase: kv[1]):
																								if index != token_index:
																														logger.warning(
																														    f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
																														    ''' Please check that the tokenizer is not corrupted!''')
																														UpperCAmelCase_     							=     token_index
																								writer.write(''' '''.join(_lowercase) + '''\n''')
																								index += 1

												return vocab_file, merge_file










import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MarianConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
							import os

							# The slow tests are often failing with OOM error on GPU
							# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
							# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
							UpperCamelCase_	       =   "platform"

							import jax
							import jax.numpy as jnp

							from transformers import MarianTokenizer
							from transformers.models.marian.modeling_flax_marian import FlaxMarianModel, FlaxMarianMTModel, shift_tokens_right







def 				A			(		__UpperCAmelCase							,					__UpperCAmelCase							,					__UpperCAmelCase=None							,					__UpperCAmelCase=None							,					__UpperCAmelCase=None							,					__UpperCAmelCase=None							,					__UpperCAmelCase=None							,					__UpperCAmelCase=None							,					)       ->       Optional[int]:



						'''simple docstring'''






						if attention_mask is None:
												UpperCAmelCase_     							=     np.where(input_ids != config.pad_token_id							,					1							,					0	)
						if decoder_attention_mask is None:
												UpperCAmelCase_     							=     np.where(decoder_input_ids != config.pad_token_id							,					1							,					0	)
						if head_mask is None:
												UpperCAmelCase_     							=     np.ones((config.encoder_layers, config.encoder_attention_heads)	)
						if decoder_head_mask is None:
												UpperCAmelCase_     							=     np.ones((config.decoder_layers, config.decoder_attention_heads)	)
						if cross_attn_head_mask is None:
												UpperCAmelCase_     							=     np.ones((config.decoder_layers, config.decoder_attention_heads)	)
						return {
						    "input_ids": input_ids,
						    "decoder_input_ids": decoder_input_ids,
						    "attention_mask": attention_mask,
						    "decoder_attention_mask": attention_mask,
						}







class 							a_  :

						def __init__(    self	:int   ,						_lowercase	:List[str]   ,						_lowercase	:Optional[Any]=13   ,						_lowercase	:Union[str, Any]=7   ,						_lowercase	:Any=True   ,						_lowercase	:List[Any]=False   ,						_lowercase	:Any=99   ,						_lowercase	:List[Any]=16   ,						_lowercase	:List[str]=2   ,						_lowercase	:List[str]=4   ,						_lowercase	:Any=4   ,						_lowercase	:Tuple="gelu"   ,						_lowercase	:Optional[int]=0.1   ,						_lowercase	:Dict=0.1   ,						_lowercase	:Union[str, Any]=32   ,						_lowercase	:Union[str, Any]=2   ,						_lowercase	:str=1   ,						_lowercase	:str=0   ,						_lowercase	:Dict=0.02   ,						)		->					int:
												UpperCAmelCase_     							=     parent
												UpperCAmelCase_     							=     batch_size
												UpperCAmelCase_     							=     seq_length
												UpperCAmelCase_     							=     is_training
												UpperCAmelCase_     							=     use_labels
												UpperCAmelCase_     							=     vocab_size
												UpperCAmelCase_     							=     hidden_size
												UpperCAmelCase_     							=     num_hidden_layers
												UpperCAmelCase_     							=     num_attention_heads
												UpperCAmelCase_     							=     intermediate_size
												UpperCAmelCase_     							=     hidden_act
												UpperCAmelCase_     							=     hidden_dropout_prob
												UpperCAmelCase_     							=     attention_probs_dropout_prob
												UpperCAmelCase_     							=     max_position_embeddings
												UpperCAmelCase_     							=     eos_token_id
												UpperCAmelCase_     							=     pad_token_id
												UpperCAmelCase_     							=     bos_token_id
												UpperCAmelCase_     							=     initializer_range

						def 	__a       (    self	:Optional[Any])		->					int:
												UpperCAmelCase_     							=     np.clip(ids_tensor([self.batch_size, self.seq_length - 1]   ,						self.vocab_size)   ,						3   ,						self.vocab_size)
												UpperCAmelCase_     							=     np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)   ,						dtype=np.intaa))   ,						-1)

												UpperCAmelCase_     							=     shift_tokens_right(_lowercase   ,						1   ,						2)

												UpperCAmelCase_     							=     MarianConfig(
												    vocab_size=self.vocab_size   ,						d_model=self.hidden_size   ,						encoder_layers=self.num_hidden_layers   ,						decoder_layers=self.num_hidden_layers   ,						encoder_attention_heads=self.num_attention_heads   ,						decoder_attention_heads=self.num_attention_heads   ,						encoder_ffn_dim=self.intermediate_size   ,						decoder_ffn_dim=self.intermediate_size   ,						dropout=self.hidden_dropout_prob   ,						attention_dropout=self.attention_probs_dropout_prob   ,						max_position_embeddings=self.max_position_embeddings   ,						eos_token_id=self.eos_token_id   ,						bos_token_id=self.bos_token_id   ,						pad_token_id=self.pad_token_id   ,						initializer_range=self.initializer_range   ,						use_cache=_lowercase   ,						)
												UpperCAmelCase_     							=     prepare_marian_inputs_dict(_lowercase   ,						_lowercase   ,						_lowercase)
												return config, inputs_dict

						def 	__a       (    self	:List[str])		->					Optional[Any]:
												UpperCAmelCase_							,      UpperCAmelCase_     							=     self.prepare_config_and_inputs()
												return config, inputs_dict

						def 	__a       (    self	:Optional[Any]   ,						_lowercase	:Any   ,						_lowercase	:Any   ,						_lowercase	:List[str])		->					List[Any]:
												UpperCAmelCase_     							=     20
												UpperCAmelCase_     							=     model_class_name(_lowercase)

												UpperCAmelCase_     							=     model.encode(inputs_dict['''input_ids'''])

												UpperCAmelCase_							,      UpperCAmelCase_     							=     (
												    inputs_dict['''decoder_input_ids'''],
												    inputs_dict['''decoder_attention_mask'''],
												)

												UpperCAmelCase_     							=     model.init_cache(decoder_input_ids.shape[0]   ,						_lowercase   ,						_lowercase)
												UpperCAmelCase_     							=     jnp.ones((decoder_input_ids.shape[0], max_decoder_length)   ,						dtype='''i4''')

												UpperCAmelCase_     							=     jnp.broadcast_to(
												    jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]   ,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)   ,						)
												UpperCAmelCase_     							=     model.decode(
												    decoder_input_ids[:, :-1]   ,						_lowercase   ,						decoder_attention_mask=_lowercase   ,						past_key_values=_lowercase   ,						decoder_position_ids=_lowercase   ,						)

												UpperCAmelCase_     							=     jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]   ,						dtype='''i4''')
												UpperCAmelCase_     							=     model.decode(
												    decoder_input_ids[:, -1:]   ,						_lowercase   ,						decoder_attention_mask=_lowercase   ,						past_key_values=outputs_cache.past_key_values   ,						decoder_position_ids=_lowercase   ,						)

												UpperCAmelCase_     							=     model.decode(_lowercase   ,						_lowercase)

												UpperCAmelCase_     							=     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
												self.parent.assertTrue(diff < 1E-3   ,						msg=f"Max diff is {diff}")




						def 	__a       (    self	:str   ,						_lowercase	:Union[str, Any]   ,						_lowercase	:List[Any]   ,						_lowercase	:Dict)		->					List[Any]:
												UpperCAmelCase_     							=     20
												UpperCAmelCase_     							=     model_class_name(_lowercase)

												UpperCAmelCase_     							=     model.encode(inputs_dict['''input_ids'''])

												UpperCAmelCase_							,      UpperCAmelCase_     							=     (
												    inputs_dict['''decoder_input_ids'''],
												    inputs_dict['''decoder_attention_mask'''],
												)

												UpperCAmelCase_     							=     jnp.concatenate(
												    [
												        decoder_attention_mask,
												        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])),
												    ]   ,						axis=-1   ,						)

												UpperCAmelCase_     							=     model.init_cache(decoder_input_ids.shape[0]   ,						_lowercase   ,						_lowercase)
												UpperCAmelCase_     							=     jnp.broadcast_to(
												    jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]   ,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)   ,						)

												UpperCAmelCase_     							=     model.decode(
												    decoder_input_ids[:, :-1]   ,						_lowercase   ,						decoder_attention_mask=_lowercase   ,						past_key_values=_lowercase   ,						decoder_position_ids=_lowercase   ,						)
												UpperCAmelCase_     							=     jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]   ,						dtype='''i4''')
												UpperCAmelCase_     							=     model.decode(
												    decoder_input_ids[:, -1:]   ,						_lowercase   ,						past_key_values=outputs_cache.past_key_values   ,						decoder_attention_mask=_lowercase   ,						decoder_position_ids=_lowercase   ,						)

												UpperCAmelCase_     							=     model.decode(_lowercase   ,						_lowercase   ,						decoder_attention_mask=_lowercase)

												UpperCAmelCase_     							=     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
												self.parent.assertTrue(diff < 1E-3   ,						msg=f"Max diff is {diff}")







@require_flax
class 							a_  (		_snake_case					, unittest.TestCase					, _snake_case				):
						UpperCamelCase__      :						Optional[Any]				       =True
						UpperCamelCase__      :						str				       =(FlaxMarianModel, FlaxMarianMTModel) if is_flax_available() else ()
						UpperCamelCase__      :						Any				       =(FlaxMarianMTModel,) if is_flax_available() else ()

						def 	__a       (    self	:List[Any])		->					Tuple:
												UpperCAmelCase_     							=     FlaxMarianModelTester(self)

						def 	__a       (    self	:Tuple)		->					Optional[Any]:
												UpperCAmelCase_							,      UpperCAmelCase_     							=     self.model_tester.prepare_config_and_inputs()
												for model_class in self.all_model_classes:
																		self.model_tester.check_use_cache_forward(_lowercase   ,						_lowercase   ,						_lowercase)

						def 	__a       (    self	:Tuple)		->					str:
												UpperCAmelCase_							,      UpperCAmelCase_     							=     self.model_tester.prepare_config_and_inputs()
												for model_class in self.all_model_classes:
																		self.model_tester.check_use_cache_forward_with_attn_mask(_lowercase   ,						_lowercase   ,						_lowercase)

						def 	__a       (    self	:Dict)		->					str:
												UpperCAmelCase_							,      UpperCAmelCase_     							=     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		with self.subTest(model_class.__name__):
																								UpperCAmelCase_     							=     self._prepare_for_class(_lowercase   ,						_lowercase)
																								UpperCAmelCase_     							=     model_class(_lowercase)

																								@jax.jit
																								def encode_jitted(_lowercase	:Optional[Any]   ,						_lowercase	:Optional[Any]=None   ,						**_lowercase	:Any):
																														return model.encode(input_ids=_lowercase   ,						attention_mask=_lowercase)

																								with self.subTest('''JIT Enabled'''):
																														UpperCAmelCase_     							=     encode_jitted(**_lowercase).to_tuple()

																								with self.subTest('''JIT Disabled'''):
																														with jax.disable_jit():
																																				UpperCAmelCase_     							=     encode_jitted(**_lowercase).to_tuple()

																								self.assertEqual(len(_lowercase)   ,						len(_lowercase))
																								for jitted_output, output in zip(_lowercase   ,						_lowercase):
																														self.assertEqual(jitted_output.shape   ,						output.shape)

						def 	__a       (    self	:Optional[int])		->					str:
												UpperCAmelCase_							,      UpperCAmelCase_     							=     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		with self.subTest(model_class.__name__):
																								UpperCAmelCase_     							=     model_class(_lowercase)
																								UpperCAmelCase_     							=     model.encode(inputs_dict['''input_ids''']   ,						inputs_dict['''attention_mask'''])

																								UpperCAmelCase_     							=     {
																								    '''decoder_input_ids''': inputs_dict['''decoder_input_ids'''],
																								    '''decoder_attention_mask''': inputs_dict['''decoder_attention_mask'''],
																								    '''encoder_outputs''': encoder_outputs,
																								}

																								@jax.jit
																								def decode_jitted(_lowercase	:Dict   ,						_lowercase	:Optional[Any]   ,						_lowercase	:int):
																														return model.decode(
																														    decoder_input_ids=_lowercase   ,						decoder_attention_mask=_lowercase   ,						encoder_outputs=_lowercase   ,						)

																								with self.subTest('''JIT Enabled'''):
																														UpperCAmelCase_     							=     decode_jitted(**_lowercase).to_tuple()

																								with self.subTest('''JIT Disabled'''):
																														with jax.disable_jit():
																																				UpperCAmelCase_     							=     decode_jitted(**_lowercase).to_tuple()

																								self.assertEqual(len(_lowercase)   ,						len(_lowercase))
																								for jitted_output, output in zip(_lowercase   ,						_lowercase):
																														self.assertEqual(jitted_output.shape   ,						output.shape)




						@slow
						def 	__a       (    self	:int)		->					Any:
												for model_class_name in self.all_model_classes:
																		UpperCAmelCase_     							=     model_class_name.from_pretrained('''Helsinki-NLP/opus-mt-en-de''')
																		# FlaxMarianForSequenceClassification expects eos token in input_ids
																		UpperCAmelCase_     							=     np.ones((1, 1)) * model.config.eos_token_id
																		UpperCAmelCase_     							=     model(_lowercase)
																		self.assertIsNotNone(_lowercase)







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		unittest.TestCase				):
						UpperCamelCase__      :						str				       =None
						UpperCamelCase__      :						List[str]				       =None

						@classmethod
						def 	__a       (    cls	:Optional[int])		->					None:
												UpperCAmelCase_     							=     f"Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"
												return cls

						@cached_property
						def 	__a       (    self	:str)		->					str:
												return MarianTokenizer.from_pretrained(self.model_name)

						@property
						def 	__a       (    self	:Optional[int])		->					int:
												return self.tokenizer.eos_token_id

						@cached_property
						def 	__a       (    self	:Dict)		->					Optional[int]:
												UpperCAmelCase_     							=     FlaxMarianMTModel.from_pretrained(self.model_name)
												self.assertEqual(model.config.decoder_start_token_id   ,						model.config.pad_token_id)
												return model

						def 	__a       (    self	:List[Any]   ,						**_lowercase	:Any)		->					Union[str, Any]:
												UpperCAmelCase_     							=     self.translate_src_text(**_lowercase)
												self.assertListEqual(self.expected_text   ,						_lowercase)




						def 	__a       (    self	:List[str]   ,						**_lowercase	:List[str])		->					Any:
												UpperCAmelCase_     							=     self.tokenizer(self.src_text   ,						padding=_lowercase   ,						return_tensors='''np'''   ,						**_lowercase)
												UpperCAmelCase_     							=     self.model.generate(
												    model_inputs.input_ids   ,						attention_mask=model_inputs.attention_mask   ,						num_beams=2   ,						max_length=128   ,						).sequences
												UpperCAmelCase_     							=     self.tokenizer.batch_decode(_lowercase   ,						skip_special_tokens=_lowercase)
												return generated_words







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		_snake_case				):
						UpperCamelCase__      :						str				       ="en"
						UpperCamelCase__      :						List[str]				       ="fr"
						UpperCamelCase__      :						Dict				       =[
						    "I am a small frog.",
						    "Now I can forget the 100 words of german that I know.",
						]
						UpperCamelCase__      :						Optional[Any]				       =[
						    "Je suis une petite grenouille.",
						    "Maintenant, je peux oublier les 100 mots d'allemand que je connais.",
						]

						@slow
						def 	__a       (    self	:List[Any])		->					Dict:
												self._assert_generated_batch_equal_expected()







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		_snake_case				):
						UpperCamelCase__      :						Tuple				       ="fr"
						UpperCamelCase__      :						Tuple				       ="en"
						UpperCamelCase__      :						Union[str, Any]				       =[
						    "Donnez moi le micro.",
						    "Tom et Mary étaient assis à une table.",  # Accents
						]
						UpperCamelCase__      :						Tuple				       =[
						    "Give me the microphone.",
						    "Tom and Mary were sitting at a table.",
						]

						@slow
						def 	__a       (    self	:Union[str, Any])		->					int:
												self._assert_generated_batch_equal_expected()







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		_snake_case				):
						UpperCamelCase__      :						List[str]				       ="mt"
						UpperCamelCase__      :						List[Any]				       ="en"
						UpperCamelCase__      :						List[Any]				       =["Billi messu b'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem."]
						UpperCamelCase__      :						int				       =["Touching gently, Jesus healed a man who was affected by the sad disease of leprosy."]

						@slow
						def 	__a       (    self	:Union[str, Any])		->					Union[str, Any]:
												self._assert_generated_batch_equal_expected()







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		_snake_case				):
						UpperCamelCase__      :						str				       ="en"
						UpperCamelCase__      :						List[str]				       ="de"
						UpperCamelCase__      :						List[str]				       =[
						    "I am a small frog.",
						    "Now I can forget the 100 words of german that I know.",
						    "Tom asked his teacher for advice.",
						    "That's how I would do it.",
						    "Tom really admired Mary's courage.",
						    "Turn around and close your eyes.",
						]
						UpperCamelCase__      :						Any				       =[
						    "Ich bin ein kleiner Frosch.",
						    "Jetzt kann ich die 100 Wörter des Deutschen vergessen, die ich kenne.",
						    "Tom bat seinen Lehrer um Rat.",
						    "So würde ich das machen.",
						    "Tom bewunderte Marias Mut wirklich.",
						    "Drehen Sie sich um und schließen Sie die Augen.",
						]

						@slow
						def 	__a       (    self	:List[Any])		->					Dict:
												self._assert_generated_batch_equal_expected()







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		_snake_case				):
						UpperCamelCase__      :						List[Any]				       ="en"
						UpperCamelCase__      :						Union[str, Any]				       ="zh"
						UpperCamelCase__      :						Optional[Any]				       =["My name is Wolfgang and I live in Berlin"]
						UpperCamelCase__      :						List[Any]				       =["我叫沃尔夫冈 我住在柏林"]

						@slow
						def 	__a       (    self	:Tuple)		->					Optional[Any]:
												self._assert_generated_batch_equal_expected()







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		_snake_case				):
						UpperCamelCase__      :						Dict				       ="ru"
						UpperCamelCase__      :						Optional[int]				       ="fr"
						UpperCamelCase__      :						Optional[Any]				       =["Он показал мне рукопись своей новой пьесы."]
						UpperCamelCase__      :						int				       =["Il m'a montré le manuscrit de sa nouvelle pièce."]

						@slow
						def 	__a       (    self	:Optional[Any])		->					Union[str, Any]:
												self._assert_generated_batch_equal_expected()







@require_flax
@require_sentencepiece
@require_tokenizers
class 							a_  (		_snake_case				):
						UpperCamelCase__      :						Dict				       ="en"
						UpperCamelCase__      :						Optional[Any]				       ="ROMANCE"
						UpperCamelCase__      :						int				       =[
						    ">>fr<< Don't spend so much time watching TV.",
						    ">>pt<< Your message has been sent.",
						    ">>es<< He's two years older than me.",
						]
						UpperCamelCase__      :						Tuple				       =[
						    "Ne passez pas autant de temps à regarder la télé.",
						    "A sua mensagem foi enviada.",
						    "Es dos años más viejo que yo.",
						]

						@slow
						def 	__a       (    self	:List[str])		->					List[str]:
												self._assert_generated_batch_equal_expected()





