



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    __UpperCamelCase			:

 """simple docstring"""



 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=1_3		,				SCREAMING_SNAKE_CASE=7		,				SCREAMING_SNAKE_CASE=False		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=False		,				SCREAMING_SNAKE_CASE=False		,				SCREAMING_SNAKE_CASE=1_9		,				SCREAMING_SNAKE_CASE=3_2		,				SCREAMING_SNAKE_CASE=5		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=3_7		,				SCREAMING_SNAKE_CASE="gelu"		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=5_1_2		,				SCREAMING_SNAKE_CASE=1_6		,				SCREAMING_SNAKE_CASE=2		,				SCREAMING_SNAKE_CASE=0.02		,				SCREAMING_SNAKE_CASE=3		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=None		,				) ->			Dict:
  a__												= parent
  a__												= batch_size
  a__												= seq_length
  a__												= is_training
  a__												= use_input_mask
  a__												= use_token_type_ids
  a__												= use_labels
  a__												= vocab_size
  a__												= hidden_size
  a__												= num_hidden_layers
  a__												= num_attention_heads
  a__												= intermediate_size
  a__												= hidden_act
  a__												= hidden_dropout_prob
  a__												= attention_probs_dropout_prob
  a__												= max_position_embeddings
  a__												= type_vocab_size
  a__												= type_sequence_label_size
  a__												= initializer_range
  a__												= num_labels
  a__												= num_choices
  a__												= scope

 def 						_UpperCAmelCase       (	self      ) ->			List[str]:
  a__												= ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

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

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

  a__												= self.get_config()

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

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= EsmConfig(
      vocab_size=3_3		,				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=SCREAMING_SNAKE_CASE		,				esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}		,				)
  return config

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Dict:
  a__												= EsmForProteinFolding(config=SCREAMING_SNAKE_CASE      ).float()
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE      )

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



 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  a__												= self.prepare_config_and_inputs()
  (
      (
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			
  )												= config_and_inputs
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
  return config, inputs_dict





@require_torch
class    __UpperCamelCase			(       _lowercase	,	_lowercase	,	unittest.TestCase				):

 """simple docstring"""



 _lowercase   :    str      						=					False

 _lowercase   :    List[str]      						=					(EsmForProteinFolding,) if is_torch_available() else ()
 _lowercase   :    str      						=					()
 _lowercase   :    Optional[Any]      						=					{} if is_torch_available() else {}
 _lowercase   :    Optional[Any]      						=					False

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= EsmFoldModelTester(self      )
  a__												= ConfigTester(self		,				config_class=SCREAMING_SNAKE_CASE		,				hidden_size=3_7      )

 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  self.config_tester.run_common_tests()

 def 						_UpperCAmelCase       (	self      ) ->			Any:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE      )

 @unittest.skip('''Does not support attention outputs'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  pass

 @unittest.skip
 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  pass

 @unittest.skip('''Esm does not support embedding resizing'''      )
 def 						_UpperCAmelCase       (	self      ) ->			str:
  pass

 @unittest.skip('''Esm does not support embedding resizing'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  pass

 @unittest.skip('''ESMFold does not support passing input embeds!'''      )
 def 						_UpperCAmelCase       (	self      ) ->			List[str]:
  pass

 @unittest.skip('''ESMFold does not support head pruning.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			str:
  pass

 @unittest.skip('''ESMFold does not support head pruning.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  pass

 @unittest.skip('''ESMFold does not support head pruning.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			int:
  pass

 @unittest.skip('''ESMFold does not support head pruning.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			List[Any]:
  pass

 @unittest.skip('''ESMFold does not support head pruning.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  pass

 @unittest.skip('''ESMFold does not output hidden states in the normal way.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  pass

 @unittest.skip('''ESMfold does not output hidden states in the normal way.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  pass

 @unittest.skip('''ESMFold only has one output format.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			int:
  pass

 @unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  pass

 @unittest.skip('''ESMFold does not support input chunking.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  pass

 @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  pass

 @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			List[Any]:
  pass

 @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Any:
  pass

 @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			int:
  pass

 @unittest.skip('''ESMFold doesn\'t support data parallel.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  pass



 @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  pass





@require_torch
class    __UpperCamelCase			(       _lowercase				):

 """simple docstring"""



 @slow
 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  a__												= EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''      ).float()
  model.eval()
  a__												= torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]      )
  a__												= model(SCREAMING_SNAKE_CASE      )['''positions''']
  a__												= torch.tensor([2.58_28, 0.79_93, -10.93_34]		,				dtype=torch.floataa      )
  self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]		,				SCREAMING_SNAKE_CASE		,				atol=1e-4      )      )










import json
import os
from typing import Dict, List, Optional, Tuple

import regex as re

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


a_      :      Dict    =     logging.get_logger(__name__)


a_      :      Any    =     {
    'vocab_file': 'vocab.json',
    'merges_file': 'merges.txt',
    'tokenizer_config_file': 'tokenizer_config.json',
}

a_      :      Union[str, Any]    =     {
    'vocab_file': {
        'facebook/blenderbot_small-90M': 'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/vocab.json'
    },
    'merges_file': {
        'facebook/blenderbot_small-90M': 'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/merges.txt'
    },
    'tokenizer_config_file': {
        'facebook/blenderbot_small-90M': (
            'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/tokenizer_config.json'
        )
    },
}

a_      :      List[Any]    =     {'facebook/blenderbot_small-90M': 5_12}






def     __a					(		__UpperCAmelCase       ):
 a__												= set()
 a__												= word[0]
 for char in word[1:]:
  pairs.add((prev_char, char)       )
  a__												= char

 a__												= set(__UpperCAmelCase       )
 return pairs

class    __UpperCamelCase			(       _lowercase				):

 """simple docstring"""



 _lowercase   :    Union[str, Any]      						=					VOCAB_FILES_NAMES
 _lowercase   :    List[Any]      						=					PRETRAINED_VOCAB_FILES_MAP
 _lowercase   :    Tuple      						=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 _lowercase   :    int      						=					['''input_ids''', '''attention_mask''']

 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE="__start__"		,				SCREAMING_SNAKE_CASE="__end__"		,				SCREAMING_SNAKE_CASE="__unk__"		,				SCREAMING_SNAKE_CASE="__null__"		,				**SCREAMING_SNAKE_CASE		,				) ->			List[str]:
  super().__init__(unk_token=SCREAMING_SNAKE_CASE		,				bos_token=SCREAMING_SNAKE_CASE		,				eos_token=SCREAMING_SNAKE_CASE		,				pad_token=SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

  with open(SCREAMING_SNAKE_CASE		,				encoding='''utf-8'''      ) as vocab_handle:
   a__												= json.load(SCREAMING_SNAKE_CASE      )
  a__												= {v: k for k, v in self.encoder.items()}
  with open(SCREAMING_SNAKE_CASE		,				encoding='''utf-8'''      ) as merges_handle:
   a__												= merges_handle.read().split('''\n'''      )[1:-1]
  a__												= [tuple(merge.split()      ) for merge in merges]
  a__												= dict(zip(SCREAMING_SNAKE_CASE		,				range(len(SCREAMING_SNAKE_CASE      )      )      )      )
  a__												= {}

 @property
 def 						_UpperCAmelCase       (	self      ) ->			int:
  return len(self.encoder      )

 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  return dict(self.encoder		,				**self.added_tokens_encoder      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			str:
  if token in self.cache:
   return self.cache[token]
  a__												= re.sub('''([.,!?()])'''		,				R''' \1'''		,				SCREAMING_SNAKE_CASE      )
  a__												= re.sub('''(\')'''		,				R''' \1 '''		,				SCREAMING_SNAKE_CASE      )
  a__												= re.sub(R'''\s{2,}'''		,				''' '''		,				SCREAMING_SNAKE_CASE      )
  if "\n" in token:
   a__												= token.replace('''\n'''		,				''' __newln__'''      )

  a__												= token.split(''' '''      )
  a__												= []
  for token in tokens:
   if not len(SCREAMING_SNAKE_CASE      ):
    continue

   a__												= token.lower()
   a__												= tuple(SCREAMING_SNAKE_CASE      )
   a__												= tuple(list(word[:-1]      ) + [word[-1] + '''</w>''']      )
   a__												= get_pairs(SCREAMING_SNAKE_CASE      )

   if not pairs:
    words.append(SCREAMING_SNAKE_CASE      )
    continue

   while True:
    a__												= min(SCREAMING_SNAKE_CASE		,				key=lambda SCREAMING_SNAKE_CASE      : self.bpe_ranks.get(SCREAMING_SNAKE_CASE		,				float('''inf'''      )      )      )
    if bigram not in self.bpe_ranks:
     break
    a__							,			a__												= bigram
    a__												= []
    a__												= 0

    while i < len(SCREAMING_SNAKE_CASE      ):
     try:
      a__												= word.index(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
      new_word.extend(word[i:j]      )
      a__												= j
     except ValueError:
      new_word.extend(word[i:]      )
      break

     if word[i] == first and i < len(SCREAMING_SNAKE_CASE      ) - 1 and word[i + 1] == second:
      new_word.append(first + second      )
      i += 2
     else:
      new_word.append(word[i]      )
      i += 1
    a__												= tuple(SCREAMING_SNAKE_CASE      )
    a__												= new_word
    if len(SCREAMING_SNAKE_CASE      ) == 1:
     break
    else:
     a__												= get_pairs(SCREAMING_SNAKE_CASE      )
   a__												= '''@@ '''.join(SCREAMING_SNAKE_CASE      )
   a__												= word[:-4]

   a__												= word
   words.append(SCREAMING_SNAKE_CASE      )
  return " ".join(SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			List[str]:
  a__												= []

  a__												= re.findall(R'''\S+\n?'''		,				SCREAMING_SNAKE_CASE      )

  for token in words:
   split_tokens.extend(list(self.bpe(SCREAMING_SNAKE_CASE      ).split(''' '''      )      )      )
  return split_tokens

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			int:
  a__												= token.lower()
  return self.encoder.get(SCREAMING_SNAKE_CASE		,				self.encoder.get(self.unk_token      )      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			str:
  return self.decoder.get(SCREAMING_SNAKE_CASE		,				self.unk_token      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			str:
  a__												= ''' '''.join(SCREAMING_SNAKE_CASE      ).replace('''@@ '''		,				''''''      ).strip()
  return out_string



 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = None      ) ->			Tuple[str]:
  if not os.path.isdir(SCREAMING_SNAKE_CASE      ):
   logger.error(f"Vocabulary path ({save_directory}) should be a directory"      )
   return
  a__												= os.path.join(
      SCREAMING_SNAKE_CASE		,				(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']      )
  a__												= os.path.join(
      SCREAMING_SNAKE_CASE		,				(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']      )

  with open(SCREAMING_SNAKE_CASE		,				'''w'''		,				encoding='''utf-8'''      ) as f:
   f.write(json.dumps(self.encoder		,				indent=2		,				sort_keys=SCREAMING_SNAKE_CASE		,				ensure_ascii=SCREAMING_SNAKE_CASE      ) + '''\n'''      )

  a__												= 0
  with open(SCREAMING_SNAKE_CASE		,				'''w'''		,				encoding='''utf-8'''      ) as writer:
   writer.write('''#version: 0.2\n'''      )
   for bpe_tokens, token_index in sorted(self.bpe_ranks.items()		,				key=lambda SCREAMING_SNAKE_CASE      : kv[1]      ):
    if index != token_index:
     logger.warning(
         f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
         ''' Please check that the tokenizer is not corrupted!'''      )
     a__												= token_index
    writer.write(''' '''.join(SCREAMING_SNAKE_CASE      ) + '''\n'''      )
    index += 1

  return vocab_file, merge_file





