






"""simple docstring"""



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 				__SCREAMING_SNAKE_CASE    :
							snake_case_  =				MarianConfig
							snake_case_  =				{}
							snake_case_  =				'gelu'


							def __init__(      self   :    List[Any]      ,     snake_case   :    List[str]      ,     snake_case   :    Any=13      ,     snake_case   :    Union[str, Any]=7      ,     snake_case   :    Tuple=True      ,     snake_case   :    Optional[Any]=False      ,     snake_case   :    int=99      ,     snake_case   :    Any=32      ,     snake_case   :    Any=2      ,     snake_case   :    List[Any]=4      ,     snake_case   :    Any=37      ,     snake_case   :    List[str]=0.1      ,     snake_case   :    List[str]=0.1      ,     snake_case   :    Optional[Any]=20      ,     snake_case   :    Optional[Any]=2      ,     snake_case   :    List[Any]=1      ,     snake_case   :    Dict=0      ,     ):
													'''simple docstring'''





													A__    :	int			=	parent
													A__    :	Tuple			=	batch_size
													A__    :	List[str]			=	seq_length
													A__    :	int			=	is_training
													A__    :	Optional[Any]			=	use_labels
													A__    :	Any			=	vocab_size
													A__    :	List[str]			=	hidden_size
													A__    :	Tuple			=	num_hidden_layers
													A__    :	str			=	num_attention_heads
													A__    :	Optional[Any]			=	intermediate_size

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


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





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

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

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


							def     _UpperCamelCase							(      self   :    int      ,     snake_case   :    List[Any]      ,     snake_case   :    List[Any]				):
													'''simple docstring'''





													A__    :	str			=	TFMarianModel(config=snake_case				).get_decoder()
													A__    :	Optional[Any]			=	inputs_dict["""input_ids"""]

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

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

													A__	,	A__    :	List[Any]			=	outputs.to_tuple()

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

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

													A__    :	str			=	model(snake_case      ,     attention_mask=snake_case				)[0]
													A__    :	Union[str, Any]			=	model(snake_case      ,     attention_mask=snake_case      ,     past_key_values=snake_case				)[0]

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

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

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




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


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__    :	List[str]			=	TFMarianModelTester(self				)
													A__    :	Optional[Any]			=	ConfigTester(self      ,     config_class=snake_case				)


							def     _UpperCamelCase							(      self   :    Dict				):
													'''simple docstring'''





													self.config_tester.run_common_tests()


							def     _UpperCamelCase							(      self   :    Optional[int]				):
													'''simple docstring'''





													A__    :	Any			=	self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.check_decoder_model_past_large_inputs(*snake_case				)






@require_tf
class 				__SCREAMING_SNAKE_CASE    (					unittest.TestCase		):
							snake_case_  =				1000  # show more chars for failing integration tests


							@classmethod
							def     _UpperCamelCase							(      cls   :    Optional[int]				):
													'''simple docstring'''





													A__    :	Tuple			=	F'Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}'
													return cls


							@cached_property
							def     _UpperCamelCase							(      self   :    Dict				):
													'''simple docstring'''





													return AutoTokenizer.from_pretrained(self.model_name				)


							@property
							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													return self.tokenizer.eos_token_id


							@cached_property
							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													warnings.simplefilter("""error"""				)
													A__    :	TFMarianMTModel			=	TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name				)
													assert isinstance(snake_case      ,     snake_case				)
													A__    :	List[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     _UpperCamelCase							(      self   :    Optional[int]      ,     **snake_case   :    Union[str, Any]				):
													'''simple docstring'''





													A__    :	Tuple			=	self.translate_src_text(**snake_case				)
													self.assertListEqual(self.expected_text      ,     snake_case				)


							def     _UpperCamelCase							(      self   :    str      ,     **snake_case   :    List[Any]				):
													'''simple docstring'''





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






@require_sentencepiece
@require_tokenizers
@require_tf
class 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase		):
							snake_case_  =				'mt'
							snake_case_  =				'en'
							snake_case_  =				['Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem.']
							snake_case_  =				['Touching gently, Jesus healed a man who was affected by the sad disease of leprosy.']


							@unittest.skip("""Skipping until #12647 is resolved."""				)
							@slow
							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													self._assert_generated_batch_equal_expected()






@require_sentencepiece
@require_tokenizers
@require_tf
class 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase		):
							snake_case_  =				'en'
							snake_case_  =				'zh'
							snake_case_  =				['My name is Wolfgang and I live in Berlin']
							snake_case_  =				['我叫沃尔夫冈 我住在柏林']


							@unittest.skip("""Skipping until #12647 is resolved."""				)
							@slow
							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													self._assert_generated_batch_equal_expected()






@require_sentencepiece
@require_tokenizers
@require_tf
class 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase		):
							snake_case_  =				'en'
							snake_case_  =				'ROMANCE'
							snake_case_  =				[
							    '>>fr<< Don\'t spend so much time watching TV.',
							    '>>pt<< Your message has been sent.',
							    '>>es<< He\'s two years older than me.',
							]
							snake_case_  =				[
							    '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     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													self._assert_generated_batch_equal_expected()


							@unittest.skip("""Skipping until #12647 is resolved."""				)
							@slow
							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__    :	List[Any]			=	TranslationPipeline(self.model      ,     self.tokenizer      ,     framework="""tf"""				)
													A__    :	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"""   )   )






