
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    _UpperCamelCase:
						__A:						Tuple				   =      MarianConfig
						__A:						Optional[Any]				   =      {}
						__A:						Union[str, Any]				   =      """gelu"""

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

													_UpperCAmelCase  :			Dict	       =		hidden_dropout_prob
													_UpperCAmelCase  :			Any	       =		attention_probs_dropout_prob
													_UpperCAmelCase  :			Dict	       =		max_position_embeddings
													_UpperCAmelCase  :			int	       =		eos_token_id
													_UpperCAmelCase  :			List[str]	       =		pad_token_id
													_UpperCAmelCase  :			Optional[Any]	       =		bos_token_id

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

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

													_UpperCAmelCase  :			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	,						)
													_UpperCAmelCase  :			str	       =		prepare_marian_inputs_dict(_lowerCamelCase	,						_lowerCamelCase	,						_lowerCamelCase			)
													return config, inputs_dict






						def 						a__ (					self		:			int	,						_lowerCamelCase		:			List[Any]	,						_lowerCamelCase		:			List[Any]			):
													_UpperCAmelCase  :			str	       =		TFMarianModel(config=_lowerCamelCase			).get_decoder()
													_UpperCAmelCase  :			Optional[Any]	       =		inputs_dict["""input_ids"""]

													_UpperCAmelCase  :			Any	       =		input_ids[:1, :]
													_UpperCAmelCase  :			Union[str, Any]	       =		inputs_dict["""attention_mask"""][:1, :]
													_UpperCAmelCase  :			Dict	       =		inputs_dict["""head_mask"""]
													_UpperCAmelCase  :			Optional[Any]	       =		1

													# first forward pass
													_UpperCAmelCase  :			Union[str, Any]	       =		model(_lowerCamelCase	,						attention_mask=_lowerCamelCase	,						head_mask=_lowerCamelCase	,						use_cache=_lowerCamelCase			)

													_UpperCAmelCase  :			List[Any]	       =		outputs.to_tuple()

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

													# append to next input_ids and
													_UpperCAmelCase  :			Optional[int]	       =		tf.concat([input_ids, next_tokens]	,						axis=-1			)
													_UpperCAmelCase  :			Optional[Any]	       =		tf.concat([attention_mask, next_attn_mask]	,						axis=-1			)

													_UpperCAmelCase  :			str	       =		model(_lowerCamelCase	,						attention_mask=_lowerCamelCase			)[0]
													_UpperCAmelCase  :			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
													_UpperCAmelCase  :			List[Any]	       =		int(ids_tensor((1,)	,						output_from_past.shape[-1]			)			)
													_UpperCAmelCase  :			int	       =		output_from_no_past[:, -3:, random_slice_idx]
													_UpperCAmelCase  :			List[Any]	       =		output_from_past[:, :, random_slice_idx]

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







def 		lowerCAmelCase_     (     _SCREAMING_SNAKE_CASE      ,	_SCREAMING_SNAKE_CASE      ,	_SCREAMING_SNAKE_CASE      ,	_SCREAMING_SNAKE_CASE=None      ,	_SCREAMING_SNAKE_CASE=None      ,	_SCREAMING_SNAKE_CASE=None      ,	_SCREAMING_SNAKE_CASE=None      ,	_SCREAMING_SNAKE_CASE=None      ,	)   ->      Optional[Any]:

							"""simple docstring"""






							if attention_mask is None:
														_UpperCAmelCase  :			Optional[int]	       =		tf.cast(tf.math.not_equal(UpperCAmelCase__      ,	config.pad_token_id  )      ,	tf.inta  )
							if decoder_attention_mask is None:
														_UpperCAmelCase  :			List[str]	       =		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:
														_UpperCAmelCase  :			List[Any]	       =		tf.ones((config.encoder_layers, config.encoder_attention_heads)  )
							if decoder_head_mask is None:
														_UpperCAmelCase  :			Tuple	       =		tf.ones((config.decoder_layers, config.decoder_attention_heads)  )
							if cross_attn_head_mask is None:
														_UpperCAmelCase  :			str	       =		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    _UpperCamelCase(   SCREAMING_SNAKE_CASE		,						SCREAMING_SNAKE_CASE		,						unittest.TestCase				):
						__A:						Dict				   =      (TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
						__A:						Optional[int]				   =      (TFMarianMTModel,) if is_tf_available() else ()
						__A:						List[Any]				   =      (
						    {
						        """conversational""": TFMarianMTModel,
						        """feature-extraction""": TFMarianModel,
						        """summarization""": TFMarianMTModel,
						        """text2text-generation""": TFMarianMTModel,
						        """translation""": TFMarianMTModel,
						    }
						    if is_tf_available()
						    else {}
						)
						__A:						str				   =      True
						__A:						List[Any]				   =      False
						__A:						List[Any]				   =      False

						def 						a__ (					self		:			List[str]			):
													_UpperCAmelCase  :			List[str]	       =		TFMarianModelTester(self			)
													_UpperCAmelCase  :			Optional[Any]	       =		ConfigTester(self	,						config_class=_lowerCamelCase			)

						def 						a__ (					self		:			Dict			):
													self.config_tester.run_common_tests()






						def 						a__ (					self		:			Optional[int]			):
													_UpperCAmelCase  :			Any	       =		self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.check_decoder_model_past_large_inputs(*_lowerCamelCase			)







@require_tf
class    _UpperCamelCase(   unittest.TestCase				):
						__A:						str				   =      10_00  # show more chars for failing integration tests

						@classmethod
						def 						a__ (					cls		:			Optional[int]			):
													_UpperCAmelCase  :			Tuple	       =		f"""Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"""
													return cls

						@cached_property
						def 						a__ (					self		:			Dict			):
													return AutoTokenizer.from_pretrained(self.model_name			)

						@property
						def 						a__ (					self		:			str			):
													return self.tokenizer.eos_token_id

						@cached_property
						def 						a__ (					self		:			Union[str, Any]			):
													warnings.simplefilter("error"			)
													_UpperCAmelCase  :			TFMarianMTModel	       =		TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name			)
													assert isinstance(_lowerCamelCase	,						_lowerCamelCase			)
													_UpperCAmelCase  :			List[Any]	       =		model.config
													self.assertListEqual(c.bad_words_ids	,						[[c.pad_token_id]]			)
													self.assertEqual(c.max_length	,						5_12			)
													self.assertEqual(c.decoder_start_token_id	,						c.pad_token_id			)
													return model

						def 						a__ (					self		:			Optional[int]	,						**_lowerCamelCase		:			Union[str, Any]			):
													_UpperCAmelCase  :			Tuple	       =		self.translate_src_text(**_lowerCamelCase			)
													self.assertListEqual(self.expected_text	,						_lowerCamelCase			)






						def 						a__ (					self		:			str	,						**_lowerCamelCase		:			List[Any]			):
													_UpperCAmelCase  :			int	       =		self.tokenizer(self.src_text	,						**_lowerCamelCase	,						padding=_lowerCamelCase	,						return_tensors="tf"			)
													_UpperCAmelCase  :			Union[str, Any]	       =		self.model.generate(
													    model_inputs.input_ids	,						attention_mask=model_inputs.attention_mask	,						num_beams=2	,						max_length=1_28			)
													_UpperCAmelCase  :			Optional[int]	       =		self.tokenizer.batch_decode(generated_ids.numpy()	,						skip_special_tokens=_lowerCamelCase			)
													return generated_words







@require_sentencepiece
@require_tokenizers
@require_tf
class    _UpperCamelCase(   SCREAMING_SNAKE_CASE				):
						__A:						List[Any]				   =      """mt"""
						__A:						List[Any]				   =      """en"""
						__A:						Union[str, Any]				   =      ["""Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem."""]
						__A:						Union[str, 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 						a__ (					self		:			Any			):
													self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class    _UpperCamelCase(   SCREAMING_SNAKE_CASE				):
						__A:						Any				   =      """en"""
						__A:						Union[str, Any]				   =      """zh"""
						__A:						str				   =      ["""My name is Wolfgang and I live in Berlin"""]
						__A:						int				   =      ["""我叫沃尔夫冈 我住在柏林"""]

						@unittest.skip("Skipping until #12647 is resolved."			)
						@slow
						def 						a__ (					self		:			int			):
													self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class    _UpperCamelCase(   SCREAMING_SNAKE_CASE				):
						__A:						List[Any]				   =      """en"""
						__A:						Union[str, Any]				   =      """ROMANCE"""
						__A:						Any				   =      [
						    """>>fr<< Don\'t spend so much time watching TV.""",
						    """>>pt<< Your message has been sent.""",
						    """>>es<< He\'s two years older than me.""",
						]
						__A:						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 						a__ (					self		:			int			):
													self._assert_generated_batch_equal_expected()






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












"""simple docstring"""



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


A_     =				get_logger(__name__)




def   _lowerCAmelCase		(				UpperCAmelCase__							:       str,    UpperCAmelCase__							:       List[dict],    UpperCAmelCase__							:       list,    UpperCAmelCase__							:       list,    UpperCAmelCase__							:       int,    UpperCAmelCase__							:       GradScaler = None,    )		->Any:
						# Model states
						for i, state in enumerate(UpperCAmelCase__   ):
												A__    :	Optional[int]			=	f'{MODEL_NAME}.bin' if i == 0 else f'{MODEL_NAME}_{i}.bin'
												A__    :	int			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Model weights saved in {output_model_file}'   )
						# Optimizer states
						for i, opt in enumerate(UpperCAmelCase__   ):
												A__    :	str			=	opt.state_dict()
												A__    :	Any			=	f'{OPTIMIZER_NAME}.bin' if i == 0 else f'{OPTIMIZER_NAME}_{i}.bin'
												A__    :	Optional[int]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Optimizer state saved in {output_optimizer_file}'   )
						# Scheduler states
						for i, scheduler in enumerate(UpperCAmelCase__   ):
												A__    :	Union[str, Any]			=	scheduler.state_dict()
												A__    :	Union[str, Any]			=	f'{SCHEDULER_NAME}.bin' if i == 0 else f'{SCHEDULER_NAME}_{i}.bin'
												A__    :	Tuple			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Scheduler state saved in {output_scheduler_file}'   )
						# GradScaler state
						if scaler is not None:
												A__    :	List[str]			=	scaler.state_dict()
												A__    :	Optional[Any]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												torch.save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Gradient scaler state saved in {output_scaler_file}'   )
						# Random number generator states
						A__    :	str			=	{}
						A__    :	str			=	f'{RNG_STATE_NAME}_{process_index}.pkl'
						A__    :	Optional[Any]			=	random.getstate()
						A__    :	str			=	np.random.get_state()
						A__    :	List[str]			=	torch.get_rng_state()
						if is_xpu_available():
												A__    :	Optional[Any]			=	torch.xpu.get_rng_state_all()
						else:
												A__    :	str			=	torch.cuda.get_rng_state_all()
						if is_tpu_available():
												A__    :	Tuple			=	xm.get_rng_state()
						A__    :	str			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
						torch.save(UpperCAmelCase__,    UpperCAmelCase__   )
						logger.info(f'Random states saved in {output_states_file}'   )
						return output_dir




def   _lowerCAmelCase		(				UpperCAmelCase__							:       Any,    UpperCAmelCase__							:       Union[str, Any],    UpperCAmelCase__							:       Union[str, Any],    UpperCAmelCase__							:       Dict,    UpperCAmelCase__							:       Tuple,    UpperCAmelCase__							:       Optional[Any]=None,    UpperCAmelCase__							:       str=None,    **UpperCAmelCase__							:       Optional[int],    )		->List[str]:
						if map_location not in [None, "cpu", "on_device"]:
												raise TypeError(
												    """Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"""   )
						if map_location is None:
												A__    :	Union[str, Any]			=	"""cpu"""
						elif map_location == "on_device":
												A__    :	Any			=	PartialState().device
						# Model states
						for i, model in enumerate(UpperCAmelCase__   ):
												A__    :	Any			=	f'{MODEL_NAME}.bin' if i == 0 else f'{MODEL_NAME}_{i}.bin'
												A__    :	List[str]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												models[i].load_state_dict(torch.load(UpperCAmelCase__,    map_location=UpperCAmelCase__   ),    **UpperCAmelCase__   )
						logger.info("""All model weights loaded successfully"""   )

						# Optimizer states
						for i, opt in enumerate(UpperCAmelCase__   ):
												A__    :	Union[str, Any]			=	f'{OPTIMIZER_NAME}.bin' if i == 0 else f'{OPTIMIZER_NAME}_{i}.bin'
												A__    :	Dict			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												A__    :	List[Any]			=	torch.load(UpperCAmelCase__,    map_location=UpperCAmelCase__   )
												optimizers[i].load_state_dict(UpperCAmelCase__   )
						logger.info("""All optimizer states loaded successfully"""   )

						# Scheduler states
						for i, scheduler in enumerate(UpperCAmelCase__   ):
												A__    :	List[Any]			=	f'{SCHEDULER_NAME}.bin' if i == 0 else f'{SCHEDULER_NAME}_{i}.bin'
												A__    :	Optional[Any]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												scheduler.load_state_dict(torch.load(UpperCAmelCase__   )   )
						logger.info("""All scheduler states loaded successfully"""   )

						# GradScaler state
						if scaler is not None:
												A__    :	List[str]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												scaler.load_state_dict(torch.load(UpperCAmelCase__   )   )
												logger.info("""GradScaler state loaded successfully"""   )

						# Random states
						try:
												A__    :	Any			=	torch.load(os.path.join(UpperCAmelCase__,    f'{RNG_STATE_NAME}_{process_index}.pkl'   )   )
												random.setstate(states["""random_state"""]   )
												np.random.set_state(states["""numpy_random_seed"""]   )
												torch.set_rng_state(states["""torch_manual_seed"""]   )
												if is_xpu_available():
																		torch.xpu.set_rng_state_all(states["""torch_xpu_manual_seed"""]   )
												else:
																		torch.cuda.set_rng_state_all(states["""torch_cuda_manual_seed"""]   )
												if is_tpu_available():
																		xm.set_rng_state(states["""xm_seed"""]   )
												logger.info("""All random states loaded successfully"""   )
						except Exception:
												logger.info("""Could not load random states"""   )




def   _lowerCAmelCase		(				UpperCAmelCase__							:       Tuple,    UpperCAmelCase__							:       Dict,    UpperCAmelCase__							:       int = 0   )		->Dict:
						A__    :	Dict			=	Path(UpperCAmelCase__   ) / f'custom_checkpoint_{index}.pkl'
						logger.info(f'Saving the state of {get_pretty_name(UpperCAmelCase__   )} to {save_location}'   )
						torch.save(obj.state_dict(),    UpperCAmelCase__   )




def   _lowerCAmelCase		(				UpperCAmelCase__							:       Tuple,    UpperCAmelCase__							:       str,    UpperCAmelCase__							:       int = 0   )		->str:
						A__    :	Optional[int]			=	f'{path}/custom_checkpoint_{index}.pkl'
						logger.info(f'Loading the state of {get_pretty_name(UpperCAmelCase__   )} from {load_location}'   )
						obj.load_state_dict(torch.load(UpperCAmelCase__,    map_location="""cpu"""   )   )






