






"""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.models.esm.modeling_esmfold import EsmForProteinFolding






class  lowerCamelCase__ :






    def __init__(			self     ,A     ,A=13     ,A=7     ,A=False     ,A=True     ,A=False     ,A=False     ,A=19     ,A=32     ,A=5     ,A=4     ,A=37     ,A="gelu"     ,A=0.1     ,A=0.1     ,A=512     ,A=16     ,A=2     ,A=0.02     ,A=3     ,A=4     ,A=None     ,):
          UpperCAmelCase          =	parent
          UpperCAmelCase          =	batch_size
          UpperCAmelCase          =	seq_length
          UpperCAmelCase          =	is_training
          UpperCAmelCase          =	use_input_mask
          UpperCAmelCase          =	use_token_type_ids
          UpperCAmelCase          =	use_labels
          UpperCAmelCase          =	vocab_size
          UpperCAmelCase          =	hidden_size
          UpperCAmelCase          =	num_hidden_layers
          UpperCAmelCase          =	num_attention_heads
          UpperCAmelCase          =	intermediate_size
          UpperCAmelCase          =	hidden_act
          UpperCAmelCase          =	hidden_dropout_prob
          UpperCAmelCase          =	attention_probs_dropout_prob
          UpperCAmelCase          =	max_position_embeddings
          UpperCAmelCase          =	type_vocab_size
          UpperCAmelCase          =	type_sequence_label_size
          UpperCAmelCase          =	initializer_range
          UpperCAmelCase          =	num_labels
          UpperCAmelCase          =	num_choices
          UpperCAmelCase          =	scope






    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          =	self.get_config()

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






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	EsmConfig(
              vocab_size=33     ,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     ,is_folding_model=A     ,esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}     ,)
          return config






    def     _UpperCamelCase							(			self     ,A     ,A     ,A     ,A     ,A     ,A ):
          UpperCAmelCase          =	EsmForProteinFolding(config=A ).float()
          model.to(A )
          model.eval()
          UpperCAmelCase          =	model(A     ,attention_mask=A )
          UpperCAmelCase          =	model(A )
          UpperCAmelCase          =	model(A )

          self.parent.assertEqual(result.positions.shape     ,(8, self.batch_size, self.seq_length, 14, 3) )
          self.parent.assertEqual(result.angles.shape     ,(8, self.batch_size, self.seq_length, 7, 2) )






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







@require_torch
class  lowerCamelCase__ (	snake_case   ,					snake_case   ,					unittest.TestCase			):
    SCREAMING_SNAKE_CASE		=      False

    SCREAMING_SNAKE_CASE		=      (EsmForProteinFolding,) if is_torch_available() else ()
    SCREAMING_SNAKE_CASE		=      ()
    SCREAMING_SNAKE_CASE		=      {} if is_torch_available() else {}
    SCREAMING_SNAKE_CASE		=      False






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	EsmFoldModelTester(self )
          UpperCAmelCase          =	ConfigTester(self     ,config_class=A     ,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(*A )






    @unittest.skip("""Does not support attention outputs""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""Esm does not support embedding resizing""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""Esm does not support embedding resizing""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not support passing input embeds!""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not support head pruning.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not support head pruning.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not support head pruning.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not support head pruning.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not support head pruning.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not output hidden states in the normal way.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMfold does not output hidden states in the normal way.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold only has one output format.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold does not support input chunking.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold doesn't support torchscript compilation.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold doesn't support torchscript compilation.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold doesn't support torchscript compilation.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""ESMFold doesn't support data parallel.""" )
    def     _UpperCamelCase							(			self ):
          pass






    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""" )
    def     _UpperCamelCase							(			self ):
          pass







@require_torch
class  lowerCamelCase__ (	snake_case			):






    @slow
    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1""" ).float()
          model.eval()
          UpperCAmelCase          =	torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
          UpperCAmelCase          =	model(A )["""positions"""]
          UpperCAmelCase          =	torch.tensor([2.5828, 0.7993, -10.9334]     ,dtype=torch.floataa )
          self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]     ,A     ,atol=1e-4 ) )











"""simple docstring"""







import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
  from .tokenization_pegasus import PegasusTokenizer
else:
  _UpperCamelCase					      =  None


_UpperCamelCase					      =  logging.get_logger(__name__)


_UpperCamelCase					      =  """▁"""

_UpperCamelCase					      =  {"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

_UpperCamelCase					      =  {
    """vocab_file""": {"""google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"""},
    """tokenizer_file""": {
        """google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/tokenizer.json"""
    },
}

_UpperCamelCase					      =  {
    """google/pegasus-xsum""": 512,
}







class  lowerCamelCase__ (	snake_case			):
    SCREAMING_SNAKE_CASE		=      VOCAB_FILES_NAMES
    SCREAMING_SNAKE_CASE		=      PRETRAINED_VOCAB_FILES_MAP
    SCREAMING_SNAKE_CASE		=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    SCREAMING_SNAKE_CASE		=      PegasusTokenizer
    SCREAMING_SNAKE_CASE		=      ['''input_ids''', '''attention_mask''']






    def __init__(			self     ,A=None     ,A=None     ,A="<pad>"     ,A="</s>"     ,A="<unk>"     ,A="<mask_2>"     ,A="<mask_1>"     ,A=None     ,A=103     ,**A     ,):
          UpperCAmelCase          =	offset

          if additional_special_tokens is not None:
                if not isinstance(A     ,A ):
                      raise TypeError(
                          F'''additional_special_tokens should be of type {type(A )}, but is'''
                          F''' {type(A )}''' )

                UpperCAmelCase          =	(
                    ([mask_token_sent] + additional_special_tokens)
                    if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                    else additional_special_tokens
                )
                # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
                additional_special_tokens_extended += [
                    F'''<unk_{i}>''' for i in range(len(A )     ,self.offset - 1 )
                ]

                if len(set(A ) ) != len(A ):
                      raise ValueError(
                          """Please make sure that the provided additional_special_tokens do not contain an incorrectly"""
                          F''' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.''' )
                UpperCAmelCase          =	additional_special_tokens_extended
          else:
                UpperCAmelCase          =	[mask_token_sent] if mask_token_sent is not None else []
                additional_special_tokens += [F'''<unk_{i}>''' for i in range(2     ,self.offset )]

          super().__init__(
              A     ,tokenizer_file=A     ,pad_token=A     ,eos_token=A     ,unk_token=A     ,mask_token=A     ,mask_token_sent=A     ,offset=A     ,additional_special_tokens=A     ,**A     ,)
          UpperCAmelCase          =	vocab_file
          UpperCAmelCase          =	False if not self.vocab_file else True






    def     _UpperCamelCase							(			self     ,A ):
          UpperCAmelCase          =	set(self.all_special_ids )  # call it once instead of inside list comp
          all_special_ids.remove(self.unk_token_id )  # <unk> is only sometimes special

          if all_special_ids != set(range(len(self.additional_special_tokens ) + 3 ) ):
                raise ValueError(
                    """There should be 3 special tokens: mask_token, pad_token, and eos_token +"""
                    F''' {len(self.additional_special_tokens )} additional_special_tokens, but got {all_special_ids}''' )

          return [1 if x in all_special_ids else 0 for x in seq]






    def     _UpperCamelCase							(			self     ,A     ,A = None     ,A = False ):
          if already_has_special_tokens:
                return self._special_token_mask(A )
          elif token_ids_a is None:
                return self._special_token_mask(A ) + [1]
          else:
                return self._special_token_mask(token_ids_a + token_ids_a ) + [1]






    def     _UpperCamelCase							(			self     ,A     ,A=None ):
          if token_ids_a is None:
                return token_ids_a + [self.eos_token_id]
          # We don't expect to process pairs, but leave the pair logic for API consistency
          return token_ids_a + token_ids_a + [self.eos_token_id]






    def     _UpperCamelCase							(			self     ,A     ,A = None ):
          if not self.can_save_slow_tokenizer:
                raise ValueError(
                    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
                    """tokenizer.""" )

          if not os.path.isdir(A ):
                logger.error(F'''Vocabulary path ({save_directory}) should be a directory''' )
                return
          UpperCAmelCase          =	os.path.join(
              A     ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""] )

          if os.path.abspath(self.vocab_file ) != os.path.abspath(A ):
                copyfile(self.vocab_file     ,A )

          return (out_vocab_file,)



