

from __future__ import annotations

import unittest
import warnings

from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import TFAutoModelForSeqaSeqLM, TFMarianModel, TFMarianMTModel



@require_tf
class 						__lowercase       :




			"""simple docstring"""







			_UpperCAmelCase     :							Any    							=  MarianConfig
			_UpperCAmelCase     :							List[str]    							=  {}
			_UpperCAmelCase     :							List[str]    							=  '''gelu'''



			def __init__(				self						:  Union[str, Any]     ,		lowerCAmelCase__						:  Dict     ,		lowerCAmelCase__						:  str=13     ,		lowerCAmelCase__						:  Tuple=7     ,		lowerCAmelCase__						:  List[str]=True     ,		lowerCAmelCase__						:  str=False     ,		lowerCAmelCase__						:  List[Any]=99     ,		lowerCAmelCase__						:  Any=32     ,		lowerCAmelCase__						:  Optional[Any]=2     ,		lowerCAmelCase__						:  List[Any]=4     ,		lowerCAmelCase__						:  List[str]=37     ,		lowerCAmelCase__						:  Any=0.1     ,		lowerCAmelCase__						:  List[str]=0.1     ,		lowerCAmelCase__						:  str=20     ,		lowerCAmelCase__						:  Dict=2     ,		lowerCAmelCase__						:  str=1     ,		lowerCAmelCase__						:  Tuple=0     ,		):
									SCREAMING_SNAKE_CASE_:				List[str]							= parent
									SCREAMING_SNAKE_CASE_:				Dict							= batch_size
									SCREAMING_SNAKE_CASE_:				Union[str, Any]							= seq_length
									SCREAMING_SNAKE_CASE_:				str							= is_training
									SCREAMING_SNAKE_CASE_:				Any							= use_labels
									SCREAMING_SNAKE_CASE_:				Dict							= vocab_size
									SCREAMING_SNAKE_CASE_:				Optional[Any]							= hidden_size
									SCREAMING_SNAKE_CASE_:				Any							= num_hidden_layers
									SCREAMING_SNAKE_CASE_:				int							= num_attention_heads
									SCREAMING_SNAKE_CASE_:				Any							= intermediate_size

									SCREAMING_SNAKE_CASE_:				Optional[int]							= hidden_dropout_prob
									SCREAMING_SNAKE_CASE_:				List[Any]							= attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE_:				Tuple							= max_position_embeddings
									SCREAMING_SNAKE_CASE_:				Tuple							= eos_token_id
									SCREAMING_SNAKE_CASE_:				Optional[Any]							= pad_token_id
									SCREAMING_SNAKE_CASE_:				str							= bos_token_id



			def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]):
									SCREAMING_SNAKE_CASE_:				Optional[int]							= ids_tensor([self.batch_size, self.seq_length - 1]     ,		self.vocab_size)
									SCREAMING_SNAKE_CASE_:				str							= tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size)     ,		1)
									SCREAMING_SNAKE_CASE_:				List[str]							= tf.concat([input_ids, eos_tensor]     ,		axis=1)

									SCREAMING_SNAKE_CASE_:				Optional[int]							= ids_tensor([self.batch_size, self.seq_length]     ,		self.vocab_size)

									SCREAMING_SNAKE_CASE_:				List[Any]							= self.config_cls(
									    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_ids=[2]     ,		bos_token_id=self.bos_token_id     ,		pad_token_id=self.pad_token_id     ,		decoder_start_token_id=self.pad_token_id     ,		**self.config_updates     ,		)
									SCREAMING_SNAKE_CASE_:				Tuple							= prepare_marian_inputs_dict(lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__)
									return config, inputs_dict




			def  _SCREAMING_SNAKE_CASE    (				self						:  Any     ,		lowerCAmelCase__						:  int     ,		lowerCAmelCase__						:  Union[str, Any]):
									SCREAMING_SNAKE_CASE_:				List[str]							= TFMarianModel(config=lowerCAmelCase__).get_decoder()
									SCREAMING_SNAKE_CASE_:				int							= inputs_dict["input_ids"]

									SCREAMING_SNAKE_CASE_:				List[str]							= input_ids[:1, :]
									SCREAMING_SNAKE_CASE_:				List[str]							= inputs_dict["attention_mask"][:1, :]
									SCREAMING_SNAKE_CASE_:				List[Any]							= inputs_dict["head_mask"]
									SCREAMING_SNAKE_CASE_:				List[str]							= 1

									# first forward pass
									SCREAMING_SNAKE_CASE_:				str							= model(lowerCAmelCase__     ,		attention_mask=lowerCAmelCase__     ,		head_mask=lowerCAmelCase__     ,		use_cache=lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_	,					SCREAMING_SNAKE_CASE_:				str							= outputs.to_tuple()

									# create hypothetical next token and extent to next_input_ids
									SCREAMING_SNAKE_CASE_:				Tuple							= ids_tensor((self.batch_size, 3)     ,		config.vocab_size)
									SCREAMING_SNAKE_CASE_:				List[str]							= tf.cast(ids_tensor((self.batch_size, 3)     ,		2)     ,		tf.inta)

									# append to next input_ids and
									SCREAMING_SNAKE_CASE_:				Tuple							= tf.concat([input_ids, next_tokens]     ,		axis=-1)
									SCREAMING_SNAKE_CASE_:				Dict							= tf.concat([attention_mask, next_attn_mask]     ,		axis=-1)

									SCREAMING_SNAKE_CASE_:				Tuple							= model(lowerCAmelCase__     ,		attention_mask=lowerCAmelCase__)[0]
									SCREAMING_SNAKE_CASE_:				Union[str, Any]							= model(lowerCAmelCase__     ,		attention_mask=lowerCAmelCase__     ,		past_key_values=lowerCAmelCase__)[0]

									self.parent.assertEqual(next_tokens.shape[1]     ,		output_from_past.shape[1])

									# select random slice
									SCREAMING_SNAKE_CASE_:				int							= int(ids_tensor((1,)     ,		output_from_past.shape[-1]))
									SCREAMING_SNAKE_CASE_:				Tuple							= output_from_no_past[:, -3:, random_slice_idx]
									SCREAMING_SNAKE_CASE_:				int							= output_from_past[:, :, random_slice_idx]

									# test that outputs are equal for slice
									tf.debugging.assert_near(lowerCAmelCase__     ,		lowerCAmelCase__     ,		rtol=1E-3)







def     A_   (					_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase=None       ,			_UpperCAmelCase=None       ,			_UpperCAmelCase=None       ,			_UpperCAmelCase=None       ,			_UpperCAmelCase=None       ,			):
						if attention_mask is None:
												SCREAMING_SNAKE_CASE_:				Optional[int]							= tf.cast(tf.math.not_equal(_UpperCAmelCase       ,			config.pad_token_id					)       ,			tf.inta					)
						if decoder_attention_mask is None:
												SCREAMING_SNAKE_CASE_:				List[Any]							= tf.concat(
												    [
												        tf.ones(decoder_input_ids[:, :1].shape       ,			dtype=tf.inta					),
												        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]       ,			config.pad_token_id					)       ,			tf.inta					),
												    ]       ,			axis=-1       ,			)
						if head_mask is None:
												SCREAMING_SNAKE_CASE_:				Tuple							= tf.ones((config.encoder_layers, config.encoder_attention_heads)					)
						if decoder_head_mask is None:
												SCREAMING_SNAKE_CASE_:				List[str]							= tf.ones((config.decoder_layers, config.decoder_attention_heads)					)
						if cross_attn_head_mask is None:
												SCREAMING_SNAKE_CASE_:				Optional[Any]							= tf.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": decoder_attention_mask,
						    "head_mask": head_mask,
						    "decoder_head_mask": decoder_head_mask,
						    "cross_attn_head_mask": cross_attn_head_mask,
						}



@require_tf
class 						__lowercase       (				UpperCAmelCase_	, UpperCAmelCase_	, unittest.TestCase			):




			"""simple docstring"""







			_UpperCAmelCase     :							Optional[Any]    							=  (TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
			_UpperCAmelCase     :							Dict    							=  (TFMarianMTModel,) if is_tf_available() else ()
			_UpperCAmelCase     :							List[Any]    							=  (
			    {
			        '''conversational''': TFMarianMTModel,
			        '''feature-extraction''': TFMarianModel,
			        '''summarization''': TFMarianMTModel,
			        '''text2text-generation''': TFMarianMTModel,
			        '''translation''': TFMarianMTModel,
			    }
			    if is_tf_available()
			    else {}
			)
			_UpperCAmelCase     :							Optional[Any]    							=  True
			_UpperCAmelCase     :							Tuple    							=  False
			_UpperCAmelCase     :							Optional[Any]    							=  False



			def  _SCREAMING_SNAKE_CASE    (				self						:  Union[str, Any]):
									SCREAMING_SNAKE_CASE_:				Dict							= TFMarianModelTester(self)
									SCREAMING_SNAKE_CASE_:				str							= ConfigTester(self     ,		config_class=lowerCAmelCase__)



			def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]):
									self.config_tester.run_common_tests()




			def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]):
									SCREAMING_SNAKE_CASE_:				Union[str, Any]							= self.model_tester.prepare_config_and_inputs_for_common()
									self.model_tester.check_decoder_model_past_large_inputs(*lowerCAmelCase__)







@require_tf
class 						__lowercase       (				unittest.TestCase			):




			"""simple docstring"""







			_UpperCAmelCase     :							Any    							=  1000  # show more chars for failing integration tests



			@classmethod
			def  _SCREAMING_SNAKE_CASE    (				cls						:  Union[str, Any]):
									SCREAMING_SNAKE_CASE_:				Optional[Any]							= F"Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"
									return cls



			@cached_property
			def  _SCREAMING_SNAKE_CASE    (				self						:  Any):
									return AutoTokenizer.from_pretrained(self.model_name)



			@property
			def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]):
									return self.tokenizer.eos_token_id



			@cached_property
			def  _SCREAMING_SNAKE_CASE    (				self						:  Dict):
									warnings.simplefilter("error")
									SCREAMING_SNAKE_CASE_:				TFMarianMTModel							= TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name)
									assert isinstance(lowerCAmelCase__     ,		lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:				Union[str, Any]							= model.config
									self.assertListEqual(c.bad_words_ids     ,		[[c.pad_token_id]])
									self.assertEqual(c.max_length     ,		512)
									self.assertEqual(c.decoder_start_token_id     ,		c.pad_token_id)
									return model



			def  _SCREAMING_SNAKE_CASE    (				self						:  Tuple     ,		**lowerCAmelCase__						:  Union[str, Any]):
									SCREAMING_SNAKE_CASE_:				Union[str, Any]							= self.translate_src_text(**lowerCAmelCase__)
									self.assertListEqual(self.expected_text     ,		lowerCAmelCase__)




			def  _SCREAMING_SNAKE_CASE    (				self						:  Any     ,		**lowerCAmelCase__						:  Optional[int]):
									SCREAMING_SNAKE_CASE_:				Any							= self.tokenizer(self.src_text     ,		**lowerCAmelCase__     ,		padding=lowerCAmelCase__     ,		return_tensors="tf")
									SCREAMING_SNAKE_CASE_:				Optional[int]							= self.model.generate(
									    model_inputs.input_ids     ,		attention_mask=model_inputs.attention_mask     ,		num_beams=2     ,		max_length=128)
									SCREAMING_SNAKE_CASE_:				List[str]							= self.tokenizer.batch_decode(generated_ids.numpy()     ,		skip_special_tokens=lowerCAmelCase__)
									return generated_words







@require_sentencepiece
@require_tokenizers
@require_tf
class 						__lowercase       (				UpperCAmelCase_			):




			"""simple docstring"""







			_UpperCAmelCase     :							Union[str, Any]    							=  '''mt'''
			_UpperCAmelCase     :							List[Any]    							=  '''en'''
			_UpperCAmelCase     :							Tuple    							=  ['''Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem.''']
			_UpperCAmelCase     :							Optional[Any]    							=  ['''Touching gently, Jesus healed a man who was affected by the sad disease of leprosy.''']



			@unittest.skip("Skipping until #12647 is resolved.")
			@slow
			def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]):
									self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class 						__lowercase       (				UpperCAmelCase_			):




			"""simple docstring"""







			_UpperCAmelCase     :							Any    							=  '''en'''
			_UpperCAmelCase     :							Optional[Any]    							=  '''zh'''
			_UpperCAmelCase     :							Optional[int]    							=  ['''My name is Wolfgang and I live in Berlin''']
			_UpperCAmelCase     :							str    							=  ['''我叫沃尔夫冈 我住在柏林''']



			@unittest.skip("Skipping until #12647 is resolved.")
			@slow
			def  _SCREAMING_SNAKE_CASE    (				self						:  Union[str, Any]):
									self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class 						__lowercase       (				UpperCAmelCase_			):




			"""simple docstring"""







			_UpperCAmelCase     :							Dict    							=  '''en'''
			_UpperCAmelCase     :							Optional[Any]    							=  '''ROMANCE'''
			_UpperCAmelCase     :							Optional[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     :							List[str]    							=  [
			    '''Ne passez pas autant de temps à regarder la télé.''',
			    '''A sua mensagem foi enviada.''',
			    '''Es dos años más viejo que yo.''',
			]



			@unittest.skip("Skipping until #12647 is resolved.")
			@slow
			def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[Any]):
									self._assert_generated_batch_equal_expected()




			@unittest.skip("Skipping until #12647 is resolved.")
			@slow
			def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[Any]):
									SCREAMING_SNAKE_CASE_:				Optional[int]							= TranslationPipeline(self.model     ,		self.tokenizer     ,		framework="tf")
									SCREAMING_SNAKE_CASE_:				Tuple							= pipeline(self.src_text)
									self.assertEqual(self.expected_text     ,		[x["translation_text"] for x in output])









import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp







def     A_   (					):
						SCREAMING_SNAKE_CASE_:				Tuple							= ArgumentParser(
						    description=(
						        "PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"
						    )					)

						# Optional arguments for the launch helper
						parser.add_argument("--num_cores"       ,			type=_UpperCAmelCase       ,			default=1       ,			help="Number of TPU cores to use (1 or 8)."					)

						# positional
						parser.add_argument(
						    "training_script"       ,			type=_UpperCAmelCase       ,			help=(
						        "The full path to the single TPU training "
						        "program/script to be launched in parallel, "
						        "followed by all the arguments for the "
						        "training script"
						    )       ,			)

						# rest from the training program
						parser.add_argument("training_script_args"       ,			nargs=_UpperCAmelCase					)

						return parser.parse_args()







def     A_   (					):
						SCREAMING_SNAKE_CASE_:				List[str]							= parse_args()

						# Import training_script as a module.
						SCREAMING_SNAKE_CASE_:				Tuple							= Path(args.training_script					)
						sys.path.append(str(script_fpath.parent.resolve()					)					)
						SCREAMING_SNAKE_CASE_:				Union[str, Any]							= script_fpath.stem
						SCREAMING_SNAKE_CASE_:				Optional[Any]							= importlib.import_module(_UpperCAmelCase					)

						# Patch sys.argv
						SCREAMING_SNAKE_CASE_:				int							= [args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores					)]

						xmp.spawn(mod._mp_fn       ,			args=()       ,			nprocs=args.num_cores					)


if __name__ == "__main__":
						main()






