




from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase       :	Union[str, Any]									=     logging.get_logger(__name__)

# TODO Update this
lowerCamelCase       :	int									=     {
    """facebook/esm-1b""": """https://huggingface.co/facebook/esm-1b/resolve/main/config.json""",
    # See all ESM models at https://huggingface.co/models?filter=esm
}







class 						A(		UpperCamelCase     ):



 '''simple docstring'''


 UpperCamelCase		     =		"esm"

 def __init__(       self      :						Optional[Any]		,							A_      :						Optional[int]=None		,							A_      :						Optional[Any]=None		,							A_      :						Dict=None		,							A_      :						Any=768		,							A_      :						List[str]=12		,							A_      :						Dict=12		,							A_      :						List[Any]=3072		,							A_      :						Any=0.1		,							A_      :						Tuple=0.1		,							A_      :						Any=1026		,							A_      :						Optional[Any]=0.02		,							A_      :						Optional[int]=1E-12		,							A_      :						List[str]="absolute"		,							A_      :						int=True		,							A_      :						Optional[int]=None		,							A_      :						List[Any]=False		,							A_      :						Optional[int]=False		,							A_      :						Optional[Any]=None		,							A_      :						Dict=None		,							**A_      :						Optional[int]		,							)     ->			Optional[int]:

       """simple docstring"""






       super().__init__(pad_token_id=A_		,							mask_token_id=A_		,							**A_						)

       lowerCamelCase_	  =					vocab_size
       lowerCamelCase_	  =					hidden_size
       lowerCamelCase_	  =					num_hidden_layers
       lowerCamelCase_	  =					num_attention_heads
       lowerCamelCase_	  =					intermediate_size
       lowerCamelCase_	  =					hidden_dropout_prob
       lowerCamelCase_	  =					attention_probs_dropout_prob
       lowerCamelCase_	  =					max_position_embeddings
       lowerCamelCase_	  =					initializer_range
       lowerCamelCase_	  =					layer_norm_eps
       lowerCamelCase_	  =					position_embedding_type
       lowerCamelCase_	  =					use_cache
       lowerCamelCase_	  =					emb_layer_norm_before
       lowerCamelCase_	  =					token_dropout
       lowerCamelCase_	  =					is_folding_model
       if is_folding_model:
             if esmfold_config is None:
                   logger.info('No esmfold_config supplied for folding model, using default values.'						)
                   lowerCamelCase_	  =					EsmFoldConfig()
             elif isinstance(A_		,							A_						):
                   lowerCamelCase_	  =					EsmFoldConfig(**A_						)
             lowerCamelCase_	  =					esmfold_config
             if vocab_list is None:
                   logger.warning('No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'						)
                   lowerCamelCase_	  =					get_default_vocab_list()
             else:
                   lowerCamelCase_	  =					vocab_list
       else:
             lowerCamelCase_	  =					None
             lowerCamelCase_	  =					None
       if self.esmfold_config is not None and getattr(self.esmfold_config		,							'use_esm_attn_map'		,							A_						):
             raise ValueError('The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!'						)





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

       """simple docstring"""






       lowerCamelCase_	  =					super().to_dict()
       if isinstance(self.esmfold_config		,							A_						):
             lowerCamelCase_	  =					self.esmfold_config.to_dict()
       return output







@dataclass
class 						A:



 '''simple docstring'''


 UpperCamelCase		     =		None
 UpperCamelCase		     =		True
 UpperCamelCase		     =		False
 UpperCamelCase		     =		False
 UpperCamelCase		     =		False
 UpperCamelCase		     =		0

 UpperCamelCase		     =		True
 UpperCamelCase		     =		False

 UpperCamelCase		     =		128
 UpperCamelCase		     =		None

 def        a__				(       self      :						Any						)     ->			Optional[int]:

       """simple docstring"""






       if self.trunk is None:
             lowerCamelCase_	  =					TrunkConfig()
       elif isinstance(self.trunk		,							A_						):
             lowerCamelCase_	  =					TrunkConfig(**self.trunk						)





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

       """simple docstring"""






       lowerCamelCase_	  =					asdict(self						)
       lowerCamelCase_	  =					self.trunk.to_dict()
       return output







@dataclass
class 						A:



 '''simple docstring'''


 UpperCamelCase		     =		48
 UpperCamelCase		     =		1024
 UpperCamelCase		     =		128
 UpperCamelCase		     =		32
 UpperCamelCase		     =		32
 UpperCamelCase		     =		32
 UpperCamelCase		     =		0
 UpperCamelCase		     =		0
 UpperCamelCase		     =		False
 UpperCamelCase		     =		4
 UpperCamelCase		     =		128
 UpperCamelCase		     =		None

 def        a__				(       self      :						Optional[int]						)     ->			Union[str, Any]:

       """simple docstring"""






       if self.structure_module is None:
             lowerCamelCase_	  =					StructureModuleConfig()
       elif isinstance(self.structure_module		,							A_						):
             lowerCamelCase_	  =					StructureModuleConfig(**self.structure_module						)

       if self.max_recycles <= 0:
             raise ValueError(f"""`max_recycles` should be positive, got {self.max_recycles}."""						)
       if self.sequence_state_dim % self.sequence_state_dim != 0:
             raise ValueError(
                 '`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got'
                 f""" {self.sequence_state_dim} and {self.sequence_state_dim}."""						)
       if self.pairwise_state_dim % self.pairwise_state_dim != 0:
             raise ValueError(
                 '`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got'
                 f""" {self.pairwise_state_dim} and {self.pairwise_state_dim}."""						)

       lowerCamelCase_	  =					self.sequence_state_dim // self.sequence_head_width
       lowerCamelCase_	  =					self.pairwise_state_dim // self.pairwise_head_width

       if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
             raise ValueError(
                 '`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got'
                 f""" {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}."""						)
       if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
             raise ValueError(
                 '`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got'
                 f""" {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}."""						)
       if self.pairwise_state_dim % 2 != 0:
             raise ValueError(f"""`pairwise_state_dim` should be even, got {self.pairwise_state_dim}."""						)

       if self.dropout >= 0.4:
             raise ValueError(f"""`dropout` should not be greater than 0.4, got {self.dropout}."""						)





 def        a__				(       self      :						List[Any]						)     ->			List[Any]:

       """simple docstring"""






       lowerCamelCase_	  =					asdict(self						)
       lowerCamelCase_	  =					self.structure_module.to_dict()
       return output







@dataclass
class 						A:



 '''simple docstring'''


 UpperCamelCase		     =		384
 UpperCamelCase		     =		128
 UpperCamelCase		     =		16
 UpperCamelCase		     =		128
 UpperCamelCase		     =		12
 UpperCamelCase		     =		4
 UpperCamelCase		     =		8
 UpperCamelCase		     =		0.1
 UpperCamelCase		     =		8
 UpperCamelCase		     =		1
 UpperCamelCase		     =		2
 UpperCamelCase		     =		7
 UpperCamelCase		     =		10
 UpperCamelCase		     =		1e-8
 UpperCamelCase		     =		1e5

 def        a__				(       self      :						Any						)     ->			Optional[Any]:

       """simple docstring"""






       return asdict(self						)

def     _SCREAMING_SNAKE_CASE    (							):


      '''simple docstring'''



      return (
          "<cls>",
          "<pad>",
          "<eos>",
          "<unk>",
          "L",
          "A",
          "G",
          "V",
          "S",
          "E",
          "R",
          "T",
          "I",
          "D",
          "P",
          "K",
          "Q",
          "N",
          "F",
          "Y",
          "M",
          "H",
          "W",
          "C",
          "X",
          "B",
          "U",
          "Z",
          "O",
          ".",
          "-",
          "<null_1>",
          "<mask>",
      )










import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY







@is_pipeline_test
@require_torch_or_tf
class 						A(		unittest.TestCase     ):



 '''simple docstring'''


 UpperCamelCase		     =		MODEL_FOR_CAUSAL_LM_MAPPING
 UpperCamelCase		     =		TF_MODEL_FOR_CAUSAL_LM_MAPPING

 @require_torch
 def        a__				(       self      :						Any						)     ->			Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_	  =					pipeline(task='text-generation'		,							model='sshleifer/tiny-ctrl'		,							framework='pt'						)
       # Using `do_sample=False` to force deterministic output
       lowerCamelCase_	  =					text_generator('This is a test'		,							do_sample=A_						)
       self.assertEqual(
           A_		,							[
               {
                   'generated_text': (
                       'This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'
                       ' oscope. FiliFili@@'
                   )
               }
           ]		,							)

       lowerCamelCase_	  =					text_generator(['This is a test', 'This is a second test']						)
       self.assertEqual(
           A_		,							[
               [
                   {
                       'generated_text': (
                           'This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'
                           ' oscope. FiliFili@@'
                       )
                   }
               ],
               [
                   {
                       'generated_text': (
                           'This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy'
                           ' oscope. oscope. FiliFili@@'
                       )
                   }
               ],
           ]		,							)

       lowerCamelCase_	  =					text_generator('This is a test'		,							do_sample=A_		,							num_return_sequences=2		,							return_tensors=A_						)
       self.assertEqual(
           A_		,							[
               {'generated_token_ids': ANY(A_						)},
               {'generated_token_ids': ANY(A_						)},
           ]		,							)
       lowerCamelCase_	  =					text_generator.model.config.eos_token_id
       lowerCamelCase_	  =					'<pad>'
       lowerCamelCase_	  =					text_generator(
           ['This is a test', 'This is a second test']		,							do_sample=A_		,							num_return_sequences=2		,							batch_size=2		,							return_tensors=A_		,							)
       self.assertEqual(
           A_		,							[
               [
                   {'generated_token_ids': ANY(A_						)},
                   {'generated_token_ids': ANY(A_						)},
               ],
               [
                   {'generated_token_ids': ANY(A_						)},
                   {'generated_token_ids': ANY(A_						)},
               ],
           ]		,							)

 @require_tf
 def        a__				(       self      :						Optional[int]						)     ->			str:

       """simple docstring"""






       lowerCamelCase_	  =					pipeline(task='text-generation'		,							model='sshleifer/tiny-ctrl'		,							framework='tf'						)

       # Using `do_sample=False` to force deterministic output
       lowerCamelCase_	  =					text_generator('This is a test'		,							do_sample=A_						)
       self.assertEqual(
           A_		,							[
               {
                   'generated_text': (
                       'This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'
                       ' please,'
                   )
               }
           ]		,							)

       lowerCamelCase_	  =					text_generator(['This is a test', 'This is a second test']		,							do_sample=A_						)
       self.assertEqual(
           A_		,							[
               [
                   {
                       'generated_text': (
                           'This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'
                           ' please,'
                       )
                   }
               ],
               [
                   {
                       'generated_text': (
                           'This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes'
                           ' Cannes 閲閲Cannes Cannes Cannes 攵 please,'
                       )
                   }
               ],
           ]		,							)

 def        a__				(       self      :						Optional[int]		,							A_      :						Dict		,							A_      :						int		,							A_      :						List[str]						)     ->			str:

       """simple docstring"""






       lowerCamelCase_	  =					TextGenerationPipeline(model=A_		,							tokenizer=A_						)
       return text_generator, ["This is a test", "Another test"]

 def        a__				(       self      :						Dict						)     ->			str:

       """simple docstring"""






       lowerCamelCase_	  =					'Hello I believe in'
       lowerCamelCase_	  =					pipeline('text-generation'		,							model='hf-internal-testing/tiny-random-gpt2'						)
       lowerCamelCase_	  =					text_generator(A_						)
       self.assertEqual(
           A_		,							[{'generated_text': 'Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'}]		,							)

       lowerCamelCase_	  =					text_generator(A_		,							stop_sequence=' fe'						)
       self.assertEqual(A_		,							[{'generated_text': 'Hello I believe in fe'}]						)

 def        a__				(       self      :						Any		,							A_      :						Optional[Any]		,							A_      :						Union[str, Any]						)     ->			Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_	  =					text_generator.model
       lowerCamelCase_	  =					text_generator.tokenizer

       lowerCamelCase_	  =					text_generator('This is a test'						)
       self.assertEqual(A_		,							[{'generated_text': ANY(A_						)}]						)
       self.assertTrue(outputs[0]['generated_text'].startswith('This is a test'						)						)

       lowerCamelCase_	  =					text_generator('This is a test'		,							return_full_text=A_						)
       self.assertEqual(A_		,							[{'generated_text': ANY(A_						)}]						)
       self.assertNotIn('This is a test'		,							outputs[0]['generated_text']						)

       lowerCamelCase_	  =					pipeline(task='text-generation'		,							model=A_		,							tokenizer=A_		,							return_full_text=A_						)
       lowerCamelCase_	  =					text_generator('This is a test'						)
       self.assertEqual(A_		,							[{'generated_text': ANY(A_						)}]						)
       self.assertNotIn('This is a test'		,							outputs[0]['generated_text']						)

       lowerCamelCase_	  =					text_generator('This is a test'		,							return_full_text=A_						)
       self.assertEqual(A_		,							[{'generated_text': ANY(A_						)}]						)
       self.assertTrue(outputs[0]['generated_text'].startswith('This is a test'						)						)

       lowerCamelCase_	  =					text_generator(['This is great !', 'Something else']		,							num_return_sequences=2		,							do_sample=A_						)
       self.assertEqual(
           A_		,							[
               [{'generated_text': ANY(A_						)}, {'generated_text': ANY(A_						)}],
               [{'generated_text': ANY(A_						)}, {'generated_text': ANY(A_						)}],
           ]		,							)

       if text_generator.tokenizer.pad_token is not None:
             lowerCamelCase_	  =					text_generator(
                 ['This is great !', 'Something else']		,							num_return_sequences=2		,							batch_size=2		,							do_sample=A_						)
             self.assertEqual(
                 A_		,							[
                     [{'generated_text': ANY(A_						)}, {'generated_text': ANY(A_						)}],
                     [{'generated_text': ANY(A_						)}, {'generated_text': ANY(A_						)}],
                 ]		,							)

       with self.assertRaises(A_						):
             lowerCamelCase_	  =					text_generator('test'		,							return_full_text=A_		,							return_text=A_						)
       with self.assertRaises(A_						):
             lowerCamelCase_	  =					text_generator('test'		,							return_full_text=A_		,							return_tensors=A_						)
       with self.assertRaises(A_						):
             lowerCamelCase_	  =					text_generator('test'		,							return_text=A_		,							return_tensors=A_						)

       # Empty prompt is slighly special
       # it requires BOS token to exist.
       # Special case for Pegasus which will always append EOS so will
       # work even without BOS.
       if (
           text_generator.tokenizer.bos_token_id is not None
           or "Pegasus" in tokenizer.__class__.__name__
           or "Git" in model.__class__.__name__
       ):
             lowerCamelCase_	  =					text_generator(''						)
             self.assertEqual(A_		,							[{'generated_text': ANY(A_						)}]						)
       else:
             with self.assertRaises((ValueError, AssertionError)						):
                   lowerCamelCase_	  =					text_generator(''						)

       if text_generator.framework == "tf":
             # TF generation does not support max_new_tokens, and it's impossible
             # to control long generation with only max_length without
             # fancy calculation, dismissing tests for now.
             return
       # We don't care about infinite range models.
       # They already work.
       # Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
       lowerCamelCase_	  =					['RwkvForCausalLM', 'XGLMForCausalLM', 'GPTNeoXForCausalLM']
       if (
           tokenizer.model_max_length < 10000
           and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
       ):
             # Handling of large generations
             with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)						):
                   text_generator('This is a test' * 500		,							max_new_tokens=20						)

             lowerCamelCase_	  =					text_generator('This is a test' * 500		,							handle_long_generation='hole'		,							max_new_tokens=20						)
             # Hole strategy cannot work
             with self.assertRaises(A_						):
                   text_generator(
                       'This is a test' * 500		,							handle_long_generation='hole'		,							max_new_tokens=tokenizer.model_max_length + 10		,							)

 @require_torch
 @require_accelerate
 @require_torch_gpu
 def        a__				(       self      :						Union[str, Any]						)     ->			Any:

       """simple docstring"""






       import torch

       # Classic `model_kwargs`
       lowerCamelCase_	  =					pipeline(
           model='hf-internal-testing/tiny-random-bloom'		,							model_kwargs={'device_map': 'auto', 'torch_dtype': torch.bfloataa}		,							)
       self.assertEqual(pipe.model.device		,							torch.device(0						)						)
       self.assertEqual(pipe.model.lm_head.weight.dtype		,							torch.bfloataa						)
       lowerCamelCase_	  =					pipe('This is a test'						)
       self.assertEqual(
           A_		,							[
               {
                   'generated_text': (
                       'This is a test test test test test test test test test test test test test test test test'
                       ' test'
                   )
               }
           ]		,							)

       # Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
       lowerCamelCase_	  =					pipeline(model='hf-internal-testing/tiny-random-bloom'		,							device_map='auto'		,							torch_dtype=torch.bfloataa						)
       self.assertEqual(pipe.model.device		,							torch.device(0						)						)
       self.assertEqual(pipe.model.lm_head.weight.dtype		,							torch.bfloataa						)
       lowerCamelCase_	  =					pipe('This is a test'						)
       self.assertEqual(
           A_		,							[
               {
                   'generated_text': (
                       'This is a test test test test test test test test test test test test test test test test'
                       ' test'
                   )
               }
           ]		,							)

       # torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
       lowerCamelCase_	  =					pipeline(model='hf-internal-testing/tiny-random-bloom'		,							device_map='auto'						)
       self.assertEqual(pipe.model.device		,							torch.device(0						)						)
       self.assertEqual(pipe.model.lm_head.weight.dtype		,							torch.floataa						)
       lowerCamelCase_	  =					pipe('This is a test'						)
       self.assertEqual(
           A_		,							[
               {
                   'generated_text': (
                       'This is a test test test test test test test test test test test test test test test test'
                       ' test'
                   )
               }
           ]		,							)

 @require_torch
 @require_torch_gpu
 def        a__				(       self      :						int						)     ->			str:

       """simple docstring"""






       import torch

       lowerCamelCase_	  =					pipeline(model='hf-internal-testing/tiny-random-bloom'		,							device=0		,							torch_dtype=torch.floataa						)
       pipe('This is a test'						)

 @require_torch
 @require_accelerate
 @require_torch_gpu
 def        a__				(       self      :						List[Any]						)     ->			Dict:

       """simple docstring"""






       import torch

       lowerCamelCase_	  =					pipeline(model='hf-internal-testing/tiny-random-bloom'		,							device_map='auto'		,							torch_dtype=torch.floataa						)
       pipe('This is a test'		,							do_sample=A_		,							top_p=0.5						)





 def        a__				(       self      :						Tuple						)     ->			Dict:

       """simple docstring"""






       lowerCamelCase_	  =					'Hello world'
       lowerCamelCase_	  =					pipeline('text-generation'		,							model='hf-internal-testing/tiny-random-gpt2'						)
       if text_generator.model.framework == "tf":
             lowerCamelCase_	  =					logging.get_logger('transformers.generation.tf_utils'						)
       else:
             lowerCamelCase_	  =					logging.get_logger('transformers.generation.utils'						)
       lowerCamelCase_	  =					'Both `max_new_tokens`'  # The beggining of the message to be checked in this test

       # Both are set by the user -> log warning
       with CaptureLogger(A_						) as cl:
             lowerCamelCase_	  =					text_generator(A_		,							max_length=10		,							max_new_tokens=1						)
       self.assertIn(A_		,							cl.out						)

       # The user only sets one -> no warning
       with CaptureLogger(A_						) as cl:
             lowerCamelCase_	  =					text_generator(A_		,							max_new_tokens=1						)
       self.assertNotIn(A_		,							cl.out						)

       with CaptureLogger(A_						) as cl:
             lowerCamelCase_	  =					text_generator(A_		,							max_length=10						)
       self.assertNotIn(A_		,							cl.out						)



