




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     SCREAMING_SNAKE_CASE__		:
     def __init__(						self	:		Dict       ,						__lowerCamelCase	:		Optional[int]       ,						__lowerCamelCase	:		List[Any]=13       ,						__lowerCamelCase	:		Dict=7       ,						__lowerCamelCase	:		str=False       ,						__lowerCamelCase	:		Union[str, Any]=True       ,						__lowerCamelCase	:		List[Any]=False       ,						__lowerCamelCase	:		int=True       ,						__lowerCamelCase	:		int=33       ,						__lowerCamelCase	:		str=32       ,						__lowerCamelCase	:		Any=5       ,						__lowerCamelCase	:		List[Any]=4       ,						__lowerCamelCase	:		Tuple=37       ,						__lowerCamelCase	:		str="gelu"       ,						__lowerCamelCase	:		str=0.1       ,						__lowerCamelCase	:		Dict=0.1       ,						__lowerCamelCase	:		Union[str, Any]=5_12       ,						__lowerCamelCase	:		Tuple=16       ,						__lowerCamelCase	:		List[str]=2       ,						__lowerCamelCase	:		Optional[int]=0.02       ,						__lowerCamelCase	:		Tuple=3       ,						__lowerCamelCase	:		Dict=4       ,						__lowerCamelCase	:		List[Any]=None       ,						):



       """simple docstring"""
       lowerCAmelCase__										=      parent
       lowerCAmelCase__										=      batch_size
       lowerCAmelCase__										=      seq_length
       lowerCAmelCase__										=      is_training
       lowerCAmelCase__										=      use_input_mask
       lowerCAmelCase__										=      use_token_type_ids
       lowerCAmelCase__										=      use_labels
       lowerCAmelCase__										=      vocab_size
       lowerCAmelCase__										=      hidden_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_labels
       lowerCAmelCase__										=      num_choices
       lowerCAmelCase__										=      scope
     def        A__      (						self	:		Tuple						):



       """simple docstring"""
       lowerCAmelCase__										=      ids_tensor([self.batch_size, self.seq_length]       ,						self.vocab_size						)

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

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

       lowerCAmelCase__										=      self.get_config()

       return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
     def        A__      (						self	:		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        A__      (						self	:		str       ,						__lowerCamelCase	:		Optional[int]       ,						__lowerCamelCase	:		Any       ,						__lowerCamelCase	:		List[str]       ,						__lowerCamelCase	:		Dict       ,						__lowerCamelCase	:		Dict       ,						__lowerCamelCase	:		Optional[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      EsmModel(config=__lowerCamelCase						)
       model.to(__lowerCamelCase						)
       model.eval()
       lowerCAmelCase__										=      model(__lowerCamelCase       ,						attention_mask=__lowerCamelCase						)
       lowerCAmelCase__										=      model(__lowerCamelCase						)
       lowerCAmelCase__										=      model(__lowerCamelCase						)

       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        A__      (						self	:		Optional[int]       ,						__lowerCamelCase	:		Tuple       ,						__lowerCamelCase	:		Any       ,						__lowerCamelCase	:		Any       ,						__lowerCamelCase	:		Tuple       ,						__lowerCamelCase	:		str       ,						__lowerCamelCase	:		Tuple						):



       """simple docstring"""
       lowerCAmelCase__										=      EsmForMaskedLM(config=__lowerCamelCase						)
       model.to(__lowerCamelCase						)
       model.eval()
       lowerCAmelCase__										=      model(__lowerCamelCase       ,						attention_mask=__lowerCamelCase       ,						labels=__lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.seq_length, self.vocab_size)						)
     def        A__      (						self	:		List[str]       ,						__lowerCamelCase	:		Dict       ,						__lowerCamelCase	:		Union[str, Any]       ,						__lowerCamelCase	:		List[str]       ,						__lowerCamelCase	:		List[Any]       ,						__lowerCamelCase	:		List[Any]       ,						__lowerCamelCase	:		Union[str, Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.num_labels
       lowerCAmelCase__										=      EsmForTokenClassification(config=__lowerCamelCase						)
       model.to(__lowerCamelCase						)
       model.eval()
       lowerCAmelCase__										=      model(__lowerCamelCase       ,						attention_mask=__lowerCamelCase       ,						labels=__lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.seq_length, self.num_labels)						)



     def        A__      (						self	:		Union[str, Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.prepare_config_and_inputs()
       (
           (
           lowerCAmelCase__
       )				,      (
           lowerCAmelCase__
       )				,      (
           lowerCAmelCase__
       )				,      (
           lowerCAmelCase__
       )				,      (
           lowerCAmelCase__
       )				,      (
           lowerCAmelCase__
       )				,      
       )										=      config_and_inputs
       lowerCAmelCase__										=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict





@require_torch
class     SCREAMING_SNAKE_CASE__		(_a    ,					_a    ,					unittest.TestCase			):
     lowercase_ :      str				= False

     lowercase_ :      Optional[int]				= (
         (
             EsmForMaskedLM,
             EsmModel,
             EsmForSequenceClassification,
             EsmForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     lowercase_ :      Dict				= ()
     lowercase_ :      Union[str, Any]				= (
         {
             "feature-extraction": EsmModel,
             "fill-mask": EsmForMaskedLM,
             "text-classification": EsmForSequenceClassification,
             "token-classification": EsmForTokenClassification,
             "zero-shot": EsmForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     lowercase_ :      Dict				= True
     def        A__      (						self	:		str						):



       """simple docstring"""
       lowerCAmelCase__										=      EsmModelTester(self						)
       lowerCAmelCase__										=      ConfigTester(self       ,						config_class=__lowerCamelCase       ,						hidden_size=37						)
     def        A__      (						self	:		str						):



       """simple docstring"""
       self.config_tester.run_common_tests()
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__lowerCamelCase						)
     def        A__      (						self	:		Any						):



       """simple docstring"""
       lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
         lowerCAmelCase__										=      type
         self.model_tester.create_and_check_model(*__lowerCamelCase						)
     def        A__      (						self	:		Dict						):



       """simple docstring"""
       lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase						)
     def        A__      (						self	:		Dict						):



       """simple docstring"""
       lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase						)
     @slow
     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCAmelCase__										=      EsmModel.from_pretrained(__lowerCamelCase						)
         self.assertIsNotNone(__lowerCamelCase						)
     def        A__      (						self	:		List[str]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs()[0]
       lowerCAmelCase__										=      EsmEmbeddings(config=__lowerCamelCase						)

       lowerCAmelCase__										=      torch.as_tensor([[12, 31, 13, model.padding_idx]]						)
       lowerCAmelCase__										=      torch.as_tensor(
           [
               [
                   0 + model.padding_idx + 1,
                   1 + model.padding_idx + 1,
                   2 + model.padding_idx + 1,
                   model.padding_idx,
               ]
           ]						)
       lowerCAmelCase__										=      create_position_ids_from_input_ids(__lowerCamelCase       ,						model.padding_idx						)
       self.assertEqual(position_ids.shape       ,						expected_positions.shape						)
       self.assertTrue(torch.all(torch.eq(__lowerCamelCase       ,						__lowerCamelCase						)						)						)
     def        A__      (						self	:		List[str]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs()[0]
       lowerCAmelCase__										=      EsmEmbeddings(config=__lowerCamelCase						)

       lowerCAmelCase__										=      torch.empty(2       ,						4       ,						30						)
       lowerCAmelCase__										=      [
           0 + embeddings.padding_idx + 1,
           1 + embeddings.padding_idx + 1,
           2 + embeddings.padding_idx + 1,
           3 + embeddings.padding_idx + 1,
       ]
       lowerCAmelCase__										=      torch.as_tensor([expected_single_positions, expected_single_positions]						)
       lowerCAmelCase__										=      embeddings.create_position_ids_from_inputs_embeds(__lowerCamelCase						)
       self.assertEqual(position_ids.shape       ,						expected_positions.shape						)
       self.assertTrue(torch.all(torch.eq(__lowerCamelCase       ,						__lowerCamelCase						)						)						)
     @unittest.skip('''Esm does not support embedding resizing'''						)
     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       pass
     @unittest.skip('''Esm does not support embedding resizing'''						)
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       pass



     @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''						)
     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       pass





@require_torch
class     SCREAMING_SNAKE_CASE__		(_a			):
     @slow
     def        A__      (						self	:		Tuple						):



       """simple docstring"""
       with torch.no_grad():
         lowerCAmelCase__										=      EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''						)
         model.eval()
         lowerCAmelCase__										=      torch.tensor([[0, 1, 2, 3, 4, 5]]						)
         lowerCAmelCase__										=      model(__lowerCamelCase						)[0]

         lowerCAmelCase__										=      33

         lowerCAmelCase__										=      torch.Size((1, 6, vocab_size)						)
         self.assertEqual(output.shape       ,						__lowerCamelCase						)

         lowerCAmelCase__										=      torch.tensor(
             [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]						)
         self.assertTrue(torch.allclose(output[:, :3, :3]       ,						__lowerCamelCase       ,						atol=1e-4						)						)



     @slow
     def        A__      (						self	:		List[str]						):



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

         lowerCAmelCase__										=      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
         lowerCAmelCase__										=      model(__lowerCamelCase						)[0]
         # compare the actual values for a slice.
         lowerCAmelCase__										=      torch.tensor(
             [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]						)
         self.assertTrue(torch.allclose(output[:, :3, :3]       ,						__lowerCamelCase       ,						atol=1e-4						)						)







import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
__magic_name__  :  Tuple					=       data_utils.TransfoXLTokenizer
__magic_name__  :  int					=       data_utils.TransfoXLCorpus
__magic_name__  :  List[str]					=       data_utils
__magic_name__  :  str					=       data_utils
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				):
  if transfo_xl_dataset_file:
    # Convert a pre-processed corpus (see original TensorFlow repo)
    with open(__lowerCAmelCase   , '''rb'''				) as fp:
      lowerCAmelCase__										=      pickle.load(__lowerCAmelCase   , encoding='''latin1'''				)
    # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
    lowerCAmelCase__										=      pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
    print(F"""Save vocabulary to {pytorch_vocab_dump_path}"""				)
    lowerCAmelCase__										=      corpus.vocab.__dict__
    torch.save(__lowerCAmelCase   , __lowerCAmelCase				)

    lowerCAmelCase__										=      corpus.__dict__
    corpus_dict_no_vocab.pop('''vocab'''   , __lowerCAmelCase				)
    lowerCAmelCase__										=      pytorch_dump_folder_path + '''/''' + CORPUS_NAME
    print(F"""Save dataset to {pytorch_dataset_dump_path}"""				)
    torch.save(__lowerCAmelCase   , __lowerCAmelCase				)

  if tf_checkpoint_path:
    # Convert a pre-trained TensorFlow model
    lowerCAmelCase__										=      os.path.abspath(__lowerCAmelCase				)
    lowerCAmelCase__										=      os.path.abspath(__lowerCAmelCase				)

    print(F"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""				)
    # Initialise PyTorch model
    if transfo_xl_config_file == "":
      lowerCAmelCase__										=      TransfoXLConfig()
    else:
      lowerCAmelCase__										=      TransfoXLConfig.from_json_file(__lowerCAmelCase				)
    print(F"""Building PyTorch model from configuration: {config}"""				)
    lowerCAmelCase__										=      TransfoXLLMHeadModel(__lowerCAmelCase				)

    lowerCAmelCase__										=      load_tf_weights_in_transfo_xl(__lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				)
    # Save pytorch-model
    lowerCAmelCase__										=      os.path.join(__lowerCAmelCase   , __lowerCAmelCase				)
    lowerCAmelCase__										=      os.path.join(__lowerCAmelCase   , __lowerCAmelCase				)
    print(F"""Save PyTorch model to {os.path.abspath(__lowerCAmelCase				)}"""				)
    torch.save(model.state_dict()   , __lowerCAmelCase				)
    print(F"""Save configuration file to {os.path.abspath(__lowerCAmelCase				)}"""				)
    with open(__lowerCAmelCase   , '''w'''   , encoding='''utf-8'''				) as f:
      f.write(config.to_json_string()				)


if __name__ == "__main__":
      __magic_name__  :  int					=       argparse.ArgumentParser()
      parser.add_argument(
          """--pytorch_dump_folder_path""",
          default=None,
          type=str,
          required=True,
          help="""Path to the folder to store the PyTorch model or dataset/vocab.""",
      )
      parser.add_argument(
          """--tf_checkpoint_path""",
          default="""""",
          type=str,
          help="""An optional path to a TensorFlow checkpoint path to be converted.""",
      )
      parser.add_argument(
          """--transfo_xl_config_file""",
          default="""""",
          type=str,
          help=(
              """An optional config json file corresponding to the pre-trained BERT model. \n"""
              """This specifies the model architecture."""
          ),
      )
      parser.add_argument(
          """--transfo_xl_dataset_file""",
          default="""""",
          type=str,
          help="""An optional dataset file to be converted in a vocabulary.""",
      )
      __magic_name__  :  Optional[Any]					=       parser.parse_args()
      convert_transfo_xl_checkpoint_to_pytorch(
          args.tf_checkpoint_path,
          args.transfo_xl_config_file,
          args.pytorch_dump_folder_path,
          args.transfo_xl_dataset_file,
      )

