



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 	__lowerCAmelCase       :
 lowerCAmelCase__       =					MarianConfig
 lowerCAmelCase__       =					{}
 lowerCAmelCase__       =					"""gelu"""




 def __init__(  self		,			__UpperCAmelCase		,			__UpperCAmelCase=13		,			__UpperCAmelCase=7		,			__UpperCAmelCase=True		,			__UpperCAmelCase=False		,			__UpperCAmelCase=99		,			__UpperCAmelCase=32		,			__UpperCAmelCase=2		,			__UpperCAmelCase=4		,			__UpperCAmelCase=37		,			__UpperCAmelCase=0.1		,			__UpperCAmelCase=0.1		,			__UpperCAmelCase=20		,			__UpperCAmelCase=2		,			__UpperCAmelCase=1		,			__UpperCAmelCase=0		,			):



    '''simple docstring'''
    __lowerCamelCase  							=      parent
    __lowerCamelCase  							=      batch_size
    __lowerCamelCase  							=      seq_length
    __lowerCamelCase  							=      is_training
    __lowerCamelCase  							=      use_labels
    __lowerCamelCase  							=      vocab_size
    __lowerCamelCase  							=      hidden_size
    __lowerCamelCase  							=      num_hidden_layers
    __lowerCamelCase  							=      num_attention_heads
    __lowerCamelCase  							=      intermediate_size

    __lowerCamelCase  							=      hidden_dropout_prob
    __lowerCamelCase  							=      attention_probs_dropout_prob
    __lowerCamelCase  							=      max_position_embeddings
    __lowerCamelCase  							=      eos_token_id
    __lowerCamelCase  							=      pad_token_id
    __lowerCamelCase  							=      bos_token_id




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      ids_tensor([self.batch_size, self.seq_length - 1]		,			self.vocab_size			)
    __lowerCamelCase  							=      tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size			)		,			1			)
    __lowerCamelCase  							=      tf.concat([input_ids, eos_tensor]		,			axis=1			)

    __lowerCamelCase  							=      ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size			)

    __lowerCamelCase  							=      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		,			)
    __lowerCamelCase  							=      prepare_marian_inputs_dict(__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase			)
    return config, inputs_dict




 def     lowerCamelCase (  self		,			__UpperCAmelCase		,			__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      TFMarianModel(config=__UpperCAmelCase			).get_decoder()
    __lowerCamelCase  							=      inputs_dict['''input_ids''']

    __lowerCamelCase  							=      input_ids[:1, :]
    __lowerCamelCase  							=      inputs_dict['''attention_mask'''][:1, :]
    __lowerCamelCase  							=      inputs_dict['''head_mask''']
    __lowerCamelCase  							=      1

    # first forward pass
    __lowerCamelCase  							=      model(__UpperCAmelCase		,			attention_mask=__UpperCAmelCase		,			head_mask=__UpperCAmelCase		,			use_cache=__UpperCAmelCase			)

    __lowerCamelCase							,__lowerCamelCase  							=      outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    __lowerCamelCase  							=      ids_tensor((self.batch_size, 3)		,			config.vocab_size			)
    __lowerCamelCase  							=      tf.cast(ids_tensor((self.batch_size, 3)		,			2			)		,			tf.inta			)

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

    __lowerCamelCase  							=      model(__UpperCAmelCase		,			attention_mask=__UpperCAmelCase			)[0]
    __lowerCamelCase  							=      model(__UpperCAmelCase		,			attention_mask=__UpperCAmelCase		,			past_key_values=__UpperCAmelCase			)[0]

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

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

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





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




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      TFMarianModelTester(self			)
    __lowerCamelCase  							=      ConfigTester(self		,			config_class=__UpperCAmelCase			)




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    self.config_tester.run_common_tests()




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*__UpperCAmelCase			)



@require_tf
class 	__lowerCAmelCase       (		unittest.TestCase				):
 lowerCAmelCase__       =					1_0_0_0  # show more chars for failing integration tests




 @classmethod
 def     lowerCamelCase (  cls			):



    '''simple docstring'''
    __lowerCamelCase  							=      F"""Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"""
    return cls




 @cached_property
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    return AutoTokenizer.from_pretrained(self.model_name			)




 @property
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    return self.tokenizer.eos_token_id




 @cached_property
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    warnings.simplefilter('''error'''			)
    __lowerCamelCase  							=      TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name			)
    assert isinstance(__UpperCAmelCase		,			__UpperCAmelCase			)
    __lowerCamelCase  							=      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     lowerCamelCase (  self		,			**__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      self.translate_src_text(**__UpperCAmelCase			)
    self.assertListEqual(self.expected_text		,			__UpperCAmelCase			)




 def     lowerCamelCase (  self		,			**__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      self.tokenizer(self.src_text		,			**__UpperCAmelCase		,			padding=__UpperCAmelCase		,			return_tensors='''tf'''			)
    __lowerCamelCase  							=      self.model.generate(
        model_inputs.input_ids		,			attention_mask=model_inputs.attention_mask		,			num_beams=2		,			max_length=128			)
    __lowerCamelCase  							=      self.tokenizer.batch_decode(generated_ids.numpy()		,			skip_special_tokens=__UpperCAmelCase			)
    return generated_words



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




 @unittest.skip('''Skipping until #12647 is resolved.'''			)
 @slow
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    self._assert_generated_batch_equal_expected()



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




 @unittest.skip('''Skipping until #12647 is resolved.'''			)
 @slow
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    self._assert_generated_batch_equal_expected()



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



    '''simple docstring'''
    self._assert_generated_batch_equal_expected()




 @unittest.skip('''Skipping until #12647 is resolved.'''			)
 @slow
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      TranslationPipeline(self.model		,			self.tokenizer		,			framework='''tf'''			)
    __lowerCamelCase  							=      pipeline(self.src_text			)
    self.assertEqual(self.expected_text		,			[x['''translation_text'''] for x in output]			)











import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  from transformers.models.electra.modeling_flax_electra import (
      FlaxElectraForCausalLM,
      FlaxElectraForMaskedLM,
      FlaxElectraForMultipleChoice,
      FlaxElectraForPreTraining,
      FlaxElectraForQuestionAnswering,
      FlaxElectraForSequenceClassification,
      FlaxElectraForTokenClassification,
      FlaxElectraModel,
  )



class 	__lowerCAmelCase       (		unittest.TestCase				):




 def __init__(  self		,			__UpperCAmelCase		,			__UpperCAmelCase=13		,			__UpperCAmelCase=7		,			__UpperCAmelCase=True		,			__UpperCAmelCase=True		,			__UpperCAmelCase=True		,			__UpperCAmelCase=True		,			__UpperCAmelCase=99		,			__UpperCAmelCase=24		,			__UpperCAmelCase=32		,			__UpperCAmelCase=5		,			__UpperCAmelCase=4		,			__UpperCAmelCase=37		,			__UpperCAmelCase="gelu"		,			__UpperCAmelCase=0.1		,			__UpperCAmelCase=0.1		,			__UpperCAmelCase=512		,			__UpperCAmelCase=16		,			__UpperCAmelCase=2		,			__UpperCAmelCase=0.02		,			__UpperCAmelCase=4		,			):



    '''simple docstring'''
    __lowerCamelCase  							=      parent
    __lowerCamelCase  							=      batch_size
    __lowerCamelCase  							=      seq_length
    __lowerCamelCase  							=      is_training
    __lowerCamelCase  							=      use_attention_mask
    __lowerCamelCase  							=      use_token_type_ids
    __lowerCamelCase  							=      use_labels
    __lowerCamelCase  							=      vocab_size
    __lowerCamelCase  							=      hidden_size
    __lowerCamelCase  							=      embedding_size
    __lowerCamelCase  							=      num_hidden_layers
    __lowerCamelCase  							=      num_attention_heads
    __lowerCamelCase  							=      intermediate_size
    __lowerCamelCase  							=      hidden_act
    __lowerCamelCase  							=      hidden_dropout_prob
    __lowerCamelCase  							=      attention_probs_dropout_prob
    __lowerCamelCase  							=      max_position_embeddings
    __lowerCamelCase  							=      type_vocab_size
    __lowerCamelCase  							=      type_sequence_label_size
    __lowerCamelCase  							=      initializer_range
    __lowerCamelCase  							=      num_choices




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size			)

    __lowerCamelCase  							=      None
    if self.use_attention_mask:
       __lowerCamelCase  							=      random_attention_mask([self.batch_size, self.seq_length]			)

    __lowerCamelCase  							=      None
    if self.use_token_type_ids:
       __lowerCamelCase  							=      ids_tensor([self.batch_size, self.seq_length]		,			self.type_vocab_size			)

    __lowerCamelCase  							=      ElectraConfig(
        vocab_size=self.vocab_size		,			hidden_size=self.hidden_size		,			embedding_size=self.embedding_size		,			num_hidden_layers=self.num_hidden_layers		,			num_attention_heads=self.num_attention_heads		,			intermediate_size=self.intermediate_size		,			hidden_act=self.hidden_act		,			hidden_dropout_prob=self.hidden_dropout_prob		,			attention_probs_dropout_prob=self.attention_probs_dropout_prob		,			max_position_embeddings=self.max_position_embeddings		,			type_vocab_size=self.type_vocab_size		,			initializer_range=self.initializer_range		,			)

    return config, input_ids, token_type_ids, attention_mask




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      self.prepare_config_and_inputs()
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase  							=      config_and_inputs
    __lowerCamelCase  							=      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
    return config, inputs_dict



@require_flax
class 	__lowerCAmelCase       (		lowerCAmelCase__					,		unittest.TestCase				):
 lowerCAmelCase__       =					True

 lowerCAmelCase__       =					(
     (
         FlaxElectraModel,
         FlaxElectraForCausalLM,
         FlaxElectraForMaskedLM,
         FlaxElectraForPreTraining,
         FlaxElectraForTokenClassification,
         FlaxElectraForQuestionAnswering,
         FlaxElectraForMultipleChoice,
         FlaxElectraForSequenceClassification,
     )
     if is_flax_available()
     else ()
 )




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      FlaxElectraModelTester(self			)




 @slow
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    for model_class_name in self.all_model_classes:
       if model_class_name == FlaxElectraForMaskedLM:
          __lowerCamelCase  							=      model_class_name.from_pretrained('''google/electra-small-generator'''			)
       else:
          __lowerCamelCase  							=      model_class_name.from_pretrained('''google/electra-small-discriminator'''			)
       __lowerCamelCase  							=      model(np.ones((1, 1)			)			)
       self.assertIsNotNone(__UpperCAmelCase			)






