






"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
  import numpy
  import tensorflow as tf

  from transformers.models.esm.modeling_tf_esm import (
      TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFEsmForMaskedLM,
      TFEsmForSequenceClassification,
      TFEsmForTokenClassification,
      TFEsmModel,
  )







class  lowerCamelCase__ :






    def __init__(			self     ,A     ,):
          UpperCAmelCase          =	parent
          UpperCAmelCase          =	13
          UpperCAmelCase          =	7
          UpperCAmelCase          =	True
          UpperCAmelCase          =	True
          UpperCAmelCase          =	True
          UpperCAmelCase          =	99
          UpperCAmelCase          =	32
          UpperCAmelCase          =	2
          UpperCAmelCase          =	4
          UpperCAmelCase          =	37
          UpperCAmelCase          =	'''gelu'''
          UpperCAmelCase          =	0.1
          UpperCAmelCase          =	0.1
          UpperCAmelCase          =	512
          UpperCAmelCase          =	16
          UpperCAmelCase          =	2
          UpperCAmelCase          =	0.02
          UpperCAmelCase          =	3
          UpperCAmelCase          =	4
          UpperCAmelCase          =	None






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size )

          UpperCAmelCase          =	None
          if self.use_input_mask:
                UpperCAmelCase          =	random_attention_mask([self.batch_size, self.seq_length] )

          UpperCAmelCase          =	None
          UpperCAmelCase          =	None
          UpperCAmelCase          =	None
          if self.use_labels:
                UpperCAmelCase          =	ids_tensor([self.batch_size]     ,self.type_sequence_label_size )
                UpperCAmelCase          =	ids_tensor([self.batch_size, self.seq_length]     ,self.num_labels )
                UpperCAmelCase          =	ids_tensor([self.batch_size]     ,self.num_choices )

          UpperCAmelCase          =	EsmConfig(
              vocab_size=self.vocab_size     ,hidden_size=self.hidden_size     ,num_hidden_layers=self.num_hidden_layers     ,pad_token_id=1     ,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, input_mask, sequence_labels, token_labels, choice_labels






    def     _UpperCamelCase							(			self ):
          (
              UpperCAmelCase	
          )          =	self.prepare_config_and_inputs()

          UpperCAmelCase          =	True
          UpperCAmelCase          =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
          UpperCAmelCase          =	ids_tensor([self.batch_size, self.seq_length]     ,vocab_size=2 )

          return (
              config,
              input_ids,
              input_mask,
              sequence_labels,
              token_labels,
              choice_labels,
              encoder_hidden_states,
              encoder_attention_mask,
          )






    def     _UpperCamelCase							(			self     ,A     ,A     ,A     ,A     ,A     ,A ):
          UpperCAmelCase          =	TFEsmModel(config=_SCREAMING_SNAKE_CASE )
          UpperCAmelCase          =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE )

          UpperCAmelCase          =	[input_ids, input_mask]
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE )

          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE )

          self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size) )






    def     _UpperCamelCase							(			self     ,A     ,A     ,A     ,A     ,A     ,A     ,A     ,A     ,):
          UpperCAmelCase          =	True

          UpperCAmelCase          =	TFEsmModel(config=_SCREAMING_SNAKE_CASE )
          UpperCAmelCase          =	{
              '''input_ids''': input_ids,
              '''attention_mask''': input_mask,
              '''encoder_hidden_states''': encoder_hidden_states,
              '''encoder_attention_mask''': encoder_attention_mask,
          }
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE )

          UpperCAmelCase          =	[input_ids, input_mask]
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE     ,encoder_hidden_states=_SCREAMING_SNAKE_CASE )

          # Also check the case where encoder outputs are not passed
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE     ,attention_mask=_SCREAMING_SNAKE_CASE )

          self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size) )






    def     _UpperCamelCase							(			self     ,A     ,A     ,A     ,A     ,A     ,A ):
          UpperCAmelCase          =	TFEsmForMaskedLM(config=_SCREAMING_SNAKE_CASE )
          UpperCAmelCase          =	model([input_ids, input_mask] )
          self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length, self.vocab_size) )






    def     _UpperCamelCase							(			self     ,A     ,A     ,A     ,A     ,A     ,A ):
          UpperCAmelCase          =	self.num_labels
          UpperCAmelCase          =	TFEsmForTokenClassification(config=_SCREAMING_SNAKE_CASE )
          UpperCAmelCase          =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE )
          self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length, self.num_labels) )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	self.prepare_config_and_inputs()
          (
              UpperCAmelCase	
          )          =	config_and_inputs
          UpperCAmelCase          =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
          return config, inputs_dict







@require_tf
class  lowerCamelCase__ (	snake_case   ,					snake_case   ,					unittest.TestCase			):
    SCREAMING_SNAKE_CASE		=      (
        (
            TFEsmModel,
            TFEsmForMaskedLM,
            TFEsmForSequenceClassification,
            TFEsmForTokenClassification,
        )
        if is_tf_available()
        else ()
    )
    SCREAMING_SNAKE_CASE		=      (
        {
            '''feature-extraction''': TFEsmModel,
            '''fill-mask''': TFEsmForMaskedLM,
            '''text-classification''': TFEsmForSequenceClassification,
            '''token-classification''': TFEsmForTokenClassification,
            '''zero-shot''': TFEsmForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    SCREAMING_SNAKE_CASE		=      False
    SCREAMING_SNAKE_CASE		=      False






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	TFEsmModelTester(self )
          UpperCAmelCase          =	ConfigTester(self     ,config_class=_SCREAMING_SNAKE_CASE     ,hidden_size=37 )






    def     _UpperCamelCase							(			self ):
          self.config_tester.run_common_tests()






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(*_SCREAMING_SNAKE_CASE )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*_SCREAMING_SNAKE_CASE )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*_SCREAMING_SNAKE_CASE )






    @slow
    def     _UpperCamelCase							(			self ):





          for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCAmelCase          =	TFEsmModel.from_pretrained(_SCREAMING_SNAKE_CASE )
                self.assertIsNotNone(_SCREAMING_SNAKE_CASE )






    @unittest.skip("""Protein models do not support embedding resizing.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""Protein models do not support embedding resizing.""" )
    def     _UpperCamelCase							(			self ):
          pass






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                UpperCAmelCase          =	model_class(_SCREAMING_SNAKE_CASE )
                assert isinstance(model.get_input_embeddings()     ,tf.keras.layers.Layer )
                if model_class is TFEsmForMaskedLM:
                      # Output embedding test differs from the main test because they're a matrix, not a layer
                      UpperCAmelCase          =	model.get_bias()
                      assert isinstance(_SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE )
                      for k, v in name.items():
                            assert isinstance(_SCREAMING_SNAKE_CASE     ,tf.Variable )
                else:
                      UpperCAmelCase          =	model.get_output_embeddings()
                      assert x is None
                      UpperCAmelCase          =	model.get_bias()
                      assert name is None







@require_tf
class  lowerCamelCase__ (	unittest.TestCase			):






    @slow
    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D""" )

          UpperCAmelCase          =	tf.constant([[0, 1, 2, 3, 4, 5]] )
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE )[0]
          UpperCAmelCase          =	[1, 6, 33]
          self.assertEqual(list(output.numpy().shape )     ,_SCREAMING_SNAKE_CASE )
          # compare the actual values for a slice.
          UpperCAmelCase          =	tf.constant(
              [
                  [
                      [8.921518, -10.589_814, -6.4671307],
                      [-6.3967156, -13.911_377, -1.1211915],
                      [-7.781247, -13.951_557, -3.740592],
                  ]
              ] )
          self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()     ,expected_slice.numpy()     ,atol=1e-2 ) )






    @slow
    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D""" )

          UpperCAmelCase          =	tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
          UpperCAmelCase          =	model(_SCREAMING_SNAKE_CASE )[0]
          # compare the actual values for a slice.
          UpperCAmelCase          =	tf.constant(
              [
                  [
                      [0.14443092, 0.54125327, 0.3247739],
                      [0.30340484, 0.00526676, 0.31077722],
                      [0.32278043, -0.24987096, 0.3414628],
                  ]
              ] )
          self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()     ,expected_slice.numpy()     ,atol=1e-4 ) )











"""simple docstring"""







import io
import json
import unittest

from parameterized import parameterized

from transformers import FSMTForConditionalGeneration, FSMTTokenizer
from transformers.testing_utils import get_tests_dir, require_torch, slow, torch_device
from utils import calculate_bleu


_UpperCamelCase					      =  get_tests_dir() + """/test_data/fsmt/fsmt_val_data.json"""
with io.open(filename, """r""", encoding="""utf-8""") as f:
  _UpperCamelCase					      =  json.load(f)







@require_torch
class  lowerCamelCase__ (	unittest.TestCase			):






    def     _UpperCamelCase							(			self     ,A ):
          return FSMTTokenizer.from_pretrained(A )






    def     _UpperCamelCase							(			self     ,A ):
          UpperCAmelCase          =	FSMTForConditionalGeneration.from_pretrained(A ).to(A )
          if torch_device == "cuda":
                model.half()
          return model






    @parameterized.expand(
        [
            ["""en-ru""", 26.0],
            ["""ru-en""", 22.0],
            ["""en-de""", 22.0],
            ["""de-en""", 29.0],
        ] )
    @slow
    def     _UpperCamelCase							(			self     ,A     ,A ):
          # note: this test is not testing the best performance since it only evals a small batch
          # but it should be enough to detect a regression in the output quality
          UpperCAmelCase          =	F'''facebook/wmt19-{pair}'''
          UpperCAmelCase          =	self.get_tokenizer(A )
          UpperCAmelCase          =	self.get_model(A )

          UpperCAmelCase          =	bleu_data[pair]["""src"""]
          UpperCAmelCase          =	bleu_data[pair]["""tgt"""]

          UpperCAmelCase          =	tokenizer(A     ,return_tensors="""pt"""     ,truncation=A     ,padding="""longest""" ).to(A )
          UpperCAmelCase          =	model.generate(
              input_ids=batch.input_ids     ,num_beams=8     ,)
          UpperCAmelCase          =	tokenizer.batch_decode(
              A     ,skip_special_tokens=A     ,clean_up_tokenization_spaces=A )
          UpperCAmelCase          =	calculate_bleu(A     ,A )
          print(A )
          self.assertGreaterEqual(scores["""bleu"""]     ,A )



