


'''simple docstring'''







import os
from pathlib import Path

import numpy as np
import pytest
from pack_dataset import pack_data_dir
from parameterized import parameterized
from save_len_file import save_len_file
from torch.utils.data import DataLoader

from transformers import AutoTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right
from transformers.testing_utils import TestCasePlus, slow
from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeqaSeqDataset, SeqaSeqDataset


A__:    Dict					    =       '''bert-base-cased'''
A__:    Union[str, Any]					    =       '''google/pegasus-xsum'''
A__:    Tuple					    =       [''' Sam ate lunch today.''', '''Sams lunch ingredients.''']
A__:    Optional[int]					    =       ['''A very interesting story about what I ate for lunch.''', '''Avocado, celery, turkey, coffee''']
A__:    Tuple					    =       '''patrickvonplaten/t5-tiny-random'''
A__:    List[str]					    =       '''sshleifer/bart-tiny-random'''
A__:    List[str]					    =       '''sshleifer/tiny-mbart'''
A__:    List[str]					    =       '''sshleifer/tiny-marian-en-de'''


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Path    ,_UpperCAmelCase	:	list  )  ->			Tuple:
       _a    : str     					="""\n""".join(_UpperCAmelCase  )
       Path(_UpperCAmelCase  ).open("""w"""  ).writelines(_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Any  )  ->			Dict:
       for split in ["train", "val", "test"]:
              _dump_articles(os.path.join(_UpperCAmelCase    ,F"{split}.source"  )    ,_UpperCAmelCase  )
              _dump_articles(os.path.join(_UpperCAmelCase    ,F"{split}.target"  )    ,_UpperCAmelCase  )
       return tmp_dir







class A__       (						UpperCAmelCase__      ):






      @parameterized.expand(
          [
              MBART_TINY,
              MARIAN_TINY,
              T5_TINY,
              BART_TINY,
              PEGASUS_XSUM,
          ]     ,						)
      @slow
      def __UpperCAmelCase       ( self			:Any     ,						SCREAMING_SNAKE_CASE			:Optional[Any]	)       ->					Optional[int]:




             '''simple docstring'''
             _a    : List[Any]     					=AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
             _a    : Any     					=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)
             _a    : str     					=max(len(tokenizer.encode(SCREAMING_SNAKE_CASE	)	) for a in ARTICLES	)
             _a    : List[str]     					=max(len(tokenizer.encode(SCREAMING_SNAKE_CASE	)	) for a in SUMMARIES	)
             _a    : str     					=4
             _a    : int     					=8
             assert max_len_target > max_src_len  # Will be truncated
             assert max_len_source > max_src_len  # Will be truncated
             _a					,    _a    : str     					="""ro_RO""", """de_DE"""  # ignored for all but mbart, but never causes error.
             _a    : Tuple     					=SeqaSeqDataset(
                 SCREAMING_SNAKE_CASE     ,						data_dir=SCREAMING_SNAKE_CASE     ,						type_path="""train"""     ,						max_source_length=SCREAMING_SNAKE_CASE     ,						max_target_length=SCREAMING_SNAKE_CASE     ,						src_lang=SCREAMING_SNAKE_CASE     ,						tgt_lang=SCREAMING_SNAKE_CASE     ,						)
             _a    : Optional[Any]     					=DataLoader(SCREAMING_SNAKE_CASE     ,						batch_size=2     ,						collate_fn=train_dataset.collate_fn	)
             for batch in dataloader:
                    assert isinstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)
                    assert batch["attention_mask"].shape == batch["input_ids"].shape
                    # show that articles were trimmed.
                    assert batch["input_ids"].shape[1] == max_src_len
                    # show that targets are the same len
                    assert batch["labels"].shape[1] == max_tgt_len
                    if tok_name != MBART_TINY:
                           continue
                    # check language codes in correct place
                    _a    : List[str]     					=shift_tokens_right(batch["""labels"""]     ,						tokenizer.pad_token_id	)
                    assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
                    assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
                    assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
                    assert batch["input_ids"][0, -1].item() == tokenizer.lang_code_to_id[src_lang]

                    break  # No need to test every batch






      @parameterized.expand([BART_TINY, BERT_BASE_CASED]	)
      def __UpperCAmelCase       ( self			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:Optional[Any]	)       ->					List[str]:




             '''simple docstring'''
             _a    : Optional[Any]     					=AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
             _a    : Optional[Any]     					=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)
             _a    : str     					=max(len(tokenizer.encode(SCREAMING_SNAKE_CASE	)	) for a in ARTICLES	)
             _a    : List[str]     					=max(len(tokenizer.encode(SCREAMING_SNAKE_CASE	)	) for a in SUMMARIES	)
             _a    : List[Any]     					=4
             _a    : Any     					=LegacySeqaSeqDataset(
                 SCREAMING_SNAKE_CASE     ,						data_dir=SCREAMING_SNAKE_CASE     ,						type_path="""train"""     ,						max_source_length=2_0     ,						max_target_length=SCREAMING_SNAKE_CASE     ,						)
             _a    : List[str]     					=DataLoader(SCREAMING_SNAKE_CASE     ,						batch_size=2     ,						collate_fn=train_dataset.collate_fn	)
             for batch in dataloader:
                    assert batch["attention_mask"].shape == batch["input_ids"].shape
                    # show that articles were trimmed.
                    assert batch["input_ids"].shape[1] == max_len_source
                    assert 2_0 >= batch["input_ids"].shape[1]  # trimmed significantly
                    # show that targets were truncated
                    assert batch["labels"].shape[1] == trunc_target  # Truncated
                    assert max_len_target > trunc_target  # Truncated
                    break  # No need to test every batch






      def __UpperCAmelCase       ( self			:Union[str, Any]	)       ->					Dict:




             '''simple docstring'''
             _a    : List[str]     					=AutoTokenizer.from_pretrained("""facebook/mbart-large-cc25"""	)

             _a    : Any     					=Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)	)
             _a    : str     					=tmp_dir.joinpath("""train.source"""	).open().readlines()
             _a    : Optional[Any]     					=Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)	)
             pack_data_dir(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						1_2_8     ,						SCREAMING_SNAKE_CASE	)
             _a    : int     					={x.name for x in tmp_dir.iterdir()}
             _a    : List[str]     					={x.name for x in save_dir.iterdir()}
             _a    : Any     					=save_dir.joinpath("""train.source"""	).open().readlines()
             # orig: [' Sam ate lunch today.\n', 'Sams lunch ingredients.']
             # desired_packed: [' Sam ate lunch today.\n Sams lunch ingredients.']
             assert len(SCREAMING_SNAKE_CASE	) < len(SCREAMING_SNAKE_CASE	)
             assert len(SCREAMING_SNAKE_CASE	) == 1
             assert len(packed_examples[0]	) == sum(len(SCREAMING_SNAKE_CASE	) for x in orig_examples	)
             assert orig_paths == new_paths






      @pytest.mark.skipif(not FAIRSEQ_AVAILABLE     ,						reason="""This test requires fairseq"""	)
      def __UpperCAmelCase       ( self			:str	)       ->					Optional[int]:




             '''simple docstring'''
             if not FAIRSEQ_AVAILABLE:
                    return
             _a					,    _a					,    _a    : str     					=self._get_dataset(max_len=6_4	)
             _a    : List[Any]     					=6_4
             _a    : List[str]     					=ds.make_dynamic_sampler(SCREAMING_SNAKE_CASE     ,						required_batch_size_multiple=SCREAMING_SNAKE_CASE	)
             _a    : Optional[Any]     					=[len(SCREAMING_SNAKE_CASE	) for x in batch_sampler]
             assert len(set(SCREAMING_SNAKE_CASE	)	) > 1  # it's not dynamic batch size if every batch is the same length
             assert sum(SCREAMING_SNAKE_CASE	) == len(SCREAMING_SNAKE_CASE	)  # no dropped or added examples
             _a    : Tuple     					=DataLoader(SCREAMING_SNAKE_CASE     ,						batch_sampler=SCREAMING_SNAKE_CASE     ,						collate_fn=ds.collate_fn     ,						num_workers=2	)
             _a    : Union[str, Any]     					=[]
             _a    : str     					=[]
             for batch in data_loader:
                    _a    : int     					=batch["""input_ids"""].shape
                    _a    : Optional[int]     					=src_shape[0]
                    assert bs % required_batch_size_multiple == 0 or bs < required_batch_size_multiple
                    _a    : Optional[int]     					=np.product(batch["""input_ids"""].shape	)
                    num_src_per_batch.append(SCREAMING_SNAKE_CASE	)
                    if num_src_tokens > (max_tokens * 1.1):
                           failures.append(SCREAMING_SNAKE_CASE	)
             assert num_src_per_batch[0] == max(SCREAMING_SNAKE_CASE	)
             if failures:
                    raise AssertionError(f"too many tokens in {len(SCREAMING_SNAKE_CASE	)} batches"	)






      def __UpperCAmelCase       ( self			:Dict	)       ->					List[Any]:




             '''simple docstring'''
             _a					,    _a					,    _a    : Optional[int]     					=self._get_dataset(max_len=5_1_2	)
             _a    : List[str]     					=2
             _a    : Dict     					=ds.make_sortish_sampler(SCREAMING_SNAKE_CASE     ,						shuffle=SCREAMING_SNAKE_CASE	)

             _a    : List[Any]     					=DataLoader(SCREAMING_SNAKE_CASE     ,						batch_size=SCREAMING_SNAKE_CASE     ,						collate_fn=ds.collate_fn     ,						num_workers=2	)
             _a    : Any     					=DataLoader(SCREAMING_SNAKE_CASE     ,						batch_size=SCREAMING_SNAKE_CASE     ,						collate_fn=ds.collate_fn     ,						num_workers=2     ,						sampler=SCREAMING_SNAKE_CASE	)

             _a    : List[Any]     					=tokenizer.pad_token_id

             def count_pad_tokens(SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:Any="input_ids"	):
                    return [batch[k].eq(SCREAMING_SNAKE_CASE	).sum().item() for batch in data_loader]

             assert sum(count_pad_tokens(SCREAMING_SNAKE_CASE     ,						k="""labels"""	)	) < sum(count_pad_tokens(SCREAMING_SNAKE_CASE     ,						k="""labels"""	)	)
             assert sum(count_pad_tokens(SCREAMING_SNAKE_CASE	)	) < sum(count_pad_tokens(SCREAMING_SNAKE_CASE	)	)
             assert len(SCREAMING_SNAKE_CASE	) == len(SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:str     ,						SCREAMING_SNAKE_CASE			:Tuple=1_0_0_0     ,						SCREAMING_SNAKE_CASE			:Tuple=1_2_8	)       ->					List[Any]:




             '''simple docstring'''
             if os.getenv("""USE_REAL_DATA"""     ,						SCREAMING_SNAKE_CASE	):
                    _a    : Optional[Any]     					="""examples/seq2seq/wmt_en_ro"""
                    _a    : Optional[Any]     					=max_len * 2 * 6_4
                    if not Path(SCREAMING_SNAKE_CASE	).joinpath("""train.len"""	).exists():
                           save_len_file(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)
             else:
                    _a    : Union[str, Any]     					="""examples/seq2seq/test_data/wmt_en_ro"""
                    _a    : List[Any]     					=max_len * 4
                    save_len_file(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

             _a    : Optional[Any]     					=AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
             _a    : str     					=SeqaSeqDataset(
                 SCREAMING_SNAKE_CASE     ,						data_dir=SCREAMING_SNAKE_CASE     ,						type_path="""train"""     ,						max_source_length=SCREAMING_SNAKE_CASE     ,						max_target_length=SCREAMING_SNAKE_CASE     ,						n_obs=SCREAMING_SNAKE_CASE     ,						)
             return ds, max_tokens, tokenizer






      def __UpperCAmelCase       ( self			:int	)       ->					Any:




             '''simple docstring'''
             _a					,    _a					,    _a    : Optional[Any]     					=self._get_dataset()
             _a    : int     					=set(DistributedSortishSampler(SCREAMING_SNAKE_CASE     ,						2_5_6     ,						num_replicas=2     ,						rank=0     ,						add_extra_examples=SCREAMING_SNAKE_CASE	)	)
             _a    : List[str]     					=set(DistributedSortishSampler(SCREAMING_SNAKE_CASE     ,						2_5_6     ,						num_replicas=2     ,						rank=1     ,						add_extra_examples=SCREAMING_SNAKE_CASE	)	)
             assert idsa.intersection(SCREAMING_SNAKE_CASE	) == set()






      @parameterized.expand(
          [
              MBART_TINY,
              MARIAN_TINY,
              T5_TINY,
              BART_TINY,
              PEGASUS_XSUM,
          ]     ,						)
      def __UpperCAmelCase       ( self			:Optional[int]     ,						SCREAMING_SNAKE_CASE			:List[Any]	)       ->					Tuple:




             '''simple docstring'''
             _a    : str     					=AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE     ,						use_fast=SCREAMING_SNAKE_CASE	)
             if tok_name == MBART_TINY:
                    _a    : Tuple     					=SeqaSeqDataset(
                        SCREAMING_SNAKE_CASE     ,						data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)     ,						type_path="""train"""     ,						max_source_length=4     ,						max_target_length=8     ,						src_lang="""EN"""     ,						tgt_lang="""FR"""     ,						)
                    _a    : int     					=train_dataset.dataset_kwargs
                    assert "src_lang" in kwargs and "tgt_lang" in kwargs
             else:
                    _a    : str     					=SeqaSeqDataset(
                        SCREAMING_SNAKE_CASE     ,						data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)     ,						type_path="""train"""     ,						max_source_length=4     ,						max_target_length=8     ,						)
                    _a    : Union[str, Any]     					=train_dataset.dataset_kwargs
                    assert "add_prefix_space" not in kwargs if tok_name != BART_TINY else "add_prefix_space" in kwargs
                    assert len(SCREAMING_SNAKE_CASE	) == 1 if tok_name == BART_TINY else len(SCREAMING_SNAKE_CASE	) == 0






'''simple docstring'''







import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
   from transformers.models.esm.modeling_esm import (
       ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
       EsmEmbeddings,
       create_position_ids_from_input_ids,
   )







class A__       :






      def __init__( self			:Tuple     ,						SCREAMING_SNAKE_CASE			:int     ,						SCREAMING_SNAKE_CASE			:Optional[int]=1_3     ,						SCREAMING_SNAKE_CASE			:Optional[int]=7     ,						SCREAMING_SNAKE_CASE			:Tuple=False     ,						SCREAMING_SNAKE_CASE			:Dict=True     ,						SCREAMING_SNAKE_CASE			:Optional[int]=False     ,						SCREAMING_SNAKE_CASE			:Optional[Any]=True     ,						SCREAMING_SNAKE_CASE			:List[str]=3_3     ,						SCREAMING_SNAKE_CASE			:Tuple=3_2     ,						SCREAMING_SNAKE_CASE			:Tuple=5     ,						SCREAMING_SNAKE_CASE			:int=4     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]=3_7     ,						SCREAMING_SNAKE_CASE			:List[str]="gelu"     ,						SCREAMING_SNAKE_CASE			:Optional[Any]=0.1     ,						SCREAMING_SNAKE_CASE			:Tuple=0.1     ,						SCREAMING_SNAKE_CASE			:str=5_1_2     ,						SCREAMING_SNAKE_CASE			:Dict=1_6     ,						SCREAMING_SNAKE_CASE			:Dict=2     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]=0.02     ,						SCREAMING_SNAKE_CASE			:str=3     ,						SCREAMING_SNAKE_CASE			:List[str]=4     ,						SCREAMING_SNAKE_CASE			:List[str]=None     ,						)       ->					Union[str, Any]:




             '''simple docstring'''
             _a    : Union[str, Any]     					=parent
             _a    : List[Any]     					=batch_size
             _a    : Optional[int]     					=seq_length
             _a    : Union[str, Any]     					=is_training
             _a    : List[Any]     					=use_input_mask
             _a    : Optional[int]     					=use_token_type_ids
             _a    : int     					=use_labels
             _a    : List[str]     					=vocab_size
             _a    : List[Any]     					=hidden_size
             _a    : int     					=num_hidden_layers
             _a    : Tuple     					=num_attention_heads
             _a    : Any     					=intermediate_size
             _a    : str     					=hidden_act
             _a    : Union[str, Any]     					=hidden_dropout_prob
             _a    : Union[str, Any]     					=attention_probs_dropout_prob
             _a    : str     					=max_position_embeddings
             _a    : Dict     					=type_vocab_size
             _a    : Tuple     					=type_sequence_label_size
             _a    : Dict     					=initializer_range
             _a    : List[str]     					=num_labels
             _a    : Tuple     					=num_choices
             _a    : int     					=scope






      def __UpperCAmelCase       ( self			:Union[str, Any]	)       ->					str:




             '''simple docstring'''
             _a    : Optional[int]     					=ids_tensor([self.batch_size, self.seq_length]     ,						self.vocab_size	)

             _a    : List[Any]     					=None
             if self.use_input_mask:
                    _a    : Any     					=random_attention_mask([self.batch_size, self.seq_length]	)

             _a    : Optional[int]     					=None
             _a    : str     					=None
             _a    : Dict     					=None
             if self.use_labels:
                    _a    : Dict     					=ids_tensor([self.batch_size]     ,						self.type_sequence_label_size	)
                    _a    : str     					=ids_tensor([self.batch_size, self.seq_length]     ,						self.num_labels	)
                    _a    : List[str]     					=ids_tensor([self.batch_size]     ,						self.num_choices	)

             _a    : List[Any]     					=self.get_config()

             return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels






      def __UpperCAmelCase       ( self			:str	)       ->					Optional[int]:




             '''simple docstring'''
             return EsmConfig(
                 vocab_size=self.vocab_size     ,						hidden_size=self.hidden_size     ,						pad_token_id=1     ,						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     ,						)






      def __UpperCAmelCase       ( self			:List[str]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:List[str]     ,						SCREAMING_SNAKE_CASE			:int	)       ->					Tuple:




             '''simple docstring'''
             _a    : Any     					=EsmModel(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Optional[Any]     					=model(SCREAMING_SNAKE_CASE     ,						attention_mask=SCREAMING_SNAKE_CASE	)
             _a    : Union[str, Any]     					=model(SCREAMING_SNAKE_CASE	)
             _a    : str     					=model(SCREAMING_SNAKE_CASE	)

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






      def __UpperCAmelCase       ( self			:str     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:List[str]     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:Dict     ,						SCREAMING_SNAKE_CASE			:Dict     ,						SCREAMING_SNAKE_CASE			:Optional[Any]	)       ->					Dict:




             '''simple docstring'''
             _a    : str     					=EsmForMaskedLM(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Union[str, Any]     					=model(SCREAMING_SNAKE_CASE     ,						attention_mask=SCREAMING_SNAKE_CASE     ,						labels=SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length, self.vocab_size)	)






      def __UpperCAmelCase       ( self			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:Dict     ,						SCREAMING_SNAKE_CASE			:List[Any]     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : int     					=self.num_labels
             _a    : Tuple     					=EsmForTokenClassification(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Tuple     					=model(SCREAMING_SNAKE_CASE     ,						attention_mask=SCREAMING_SNAKE_CASE     ,						labels=SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length, self.num_labels)	)






      def __UpperCAmelCase       ( self			:Dict	)       ->					List[str]:




             '''simple docstring'''
             _a    : Optional[Any]     					=self.prepare_config_and_inputs()
             (
                 (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    
             )    : Any     					=config_and_inputs
             _a    : List[Any]     					={"""input_ids""": input_ids, """attention_mask""": input_mask}
             return config, inputs_dict







@require_torch
class A__       (						UpperCAmelCase__ ,     UpperCAmelCase__ ,     unittest.TestCase      ):
      __UpperCamelCase       :  Any	       =							False

      __UpperCamelCase       :  Any	       =							(
          (
              EsmForMaskedLM,
              EsmModel,
              EsmForSequenceClassification,
              EsmForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      __UpperCamelCase       :  str	       =							()
      __UpperCamelCase       :  List[str]	       =							(
          {
              "feature-extraction": EsmModel,
              "fill-mask": EsmForMaskedLM,
              "text-classification": EsmForSequenceClassification,
              "token-classification": EsmForTokenClassification,
              "zero-shot": EsmForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __UpperCamelCase       :  Union[str, Any]	       =							True






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					int:




             '''simple docstring'''
             _a    : Dict     					=EsmModelTester(self	)
             _a    : Optional[Any]     					=ConfigTester(self     ,						config_class=SCREAMING_SNAKE_CASE     ,						hidden_size=3_7	)






      def __UpperCAmelCase       ( self			:Tuple	)       ->					List[Any]:




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






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					str:




             '''simple docstring'''
             _a    : List[Any]     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:List[Any]	)       ->					Dict:




             '''simple docstring'''
             _a    : List[str]     					=self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                    _a    : Dict     					=type
                    self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Dict	)       ->					List[str]:




             '''simple docstring'''
             _a    : Tuple     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:List[Any]	)       ->					List[str]:




             '''simple docstring'''
             _a    : str     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE	)






      @slow
      def __UpperCAmelCase       ( self			:str	)       ->					Dict:




             '''simple docstring'''
             for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : Union[str, Any]     					=EsmModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Tuple	)       ->					int:




             '''simple docstring'''
             _a    : Optional[Any]     					=self.model_tester.prepare_config_and_inputs()[0]
             _a    : Dict     					=EsmEmbeddings(config=SCREAMING_SNAKE_CASE	)

             _a    : Tuple     					=torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]	)
             _a    : Optional[Any]     					=torch.as_tensor(
                 [
                     [
                         0 + model.padding_idx + 1,
                         1 + model.padding_idx + 1,
                         2 + model.padding_idx + 1,
                         model.padding_idx,
                     ]
                 ]	)
             _a    : Any     					=create_position_ids_from_input_ids(SCREAMING_SNAKE_CASE     ,						model.padding_idx	)
             self.assertEqual(position_ids.shape     ,						expected_positions.shape	)
             self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)	)	)






      def __UpperCAmelCase       ( self			:Optional[Any]	)       ->					Tuple:




             '''simple docstring'''
             _a    : List[Any]     					=self.model_tester.prepare_config_and_inputs()[0]
             _a    : Optional[int]     					=EsmEmbeddings(config=SCREAMING_SNAKE_CASE	)

             _a    : Tuple     					=torch.empty(2     ,						4     ,						3_0	)
             _a    : str     					=[
                 0 + embeddings.padding_idx + 1,
                 1 + embeddings.padding_idx + 1,
                 2 + embeddings.padding_idx + 1,
                 3 + embeddings.padding_idx + 1,
             ]
             _a    : int     					=torch.as_tensor([expected_single_positions, expected_single_positions]	)
             _a    : Any     					=embeddings.create_position_ids_from_inputs_embeds(SCREAMING_SNAKE_CASE	)
             self.assertEqual(position_ids.shape     ,						expected_positions.shape	)
             self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)	)	)






      @unittest.skip("""Esm does not support embedding resizing"""	)
      def __UpperCAmelCase       ( self			:Tuple	)       ->					List[str]:




             '''simple docstring'''
             pass






      @unittest.skip("""Esm does not support embedding resizing"""	)
      def __UpperCAmelCase       ( self			:str	)       ->					Any:




             '''simple docstring'''
             pass






      @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""	)
      def __UpperCAmelCase       ( self			:Dict	)       ->					Any:




             '''simple docstring'''
             pass







@require_torch
class A__       (						UpperCAmelCase__      ):






      @slow
      def __UpperCAmelCase       ( self			:List[Any]	)       ->					str:




             '''simple docstring'''
             with torch.no_grad():
                    _a    : Optional[int]     					=EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""	)
                    model.eval()
                    _a    : Any     					=torch.tensor([[0, 1, 2, 3, 4, 5]]	)
                    _a    : Tuple     					=model(SCREAMING_SNAKE_CASE	)[0]

                    _a    : int     					=3_3

                    _a    : Tuple     					=torch.Size((1, 6, vocab_size)	)
                    self.assertEqual(output.shape     ,						SCREAMING_SNAKE_CASE	)

                    _a    : Union[str, Any]     					=torch.tensor(
                        [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]	)
                    self.assertTrue(torch.allclose(output[:, :3, :3]     ,						SCREAMING_SNAKE_CASE     ,						atol=1e-4	)	)






      @slow
      def __UpperCAmelCase       ( self			:Union[str, Any]	)       ->					Tuple:




             '''simple docstring'''
             with torch.no_grad():
                    _a    : Any     					=EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""	)
                    model.eval()

                    _a    : Any     					=torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]	)
                    _a    : int     					=model(SCREAMING_SNAKE_CASE	)[0]
                    # compare the actual values for a slice.
                    _a    : str     					=torch.tensor(
                        [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]	)
                    self.assertTrue(torch.allclose(output[:, :3, :3]     ,						SCREAMING_SNAKE_CASE     ,						atol=1e-4	)	)


