
import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM


@require_torch
class 			lowercase_		:


       def __init__(    self						,       __A						,       __A=99						,       __A=13						,       __A=16						,       __A=7						,       __A=True						,       __A=True						,       __A=True						,       __A=False						,       __A=True						,       __A=2						,       __A=32						,       __A=4						,       __A=4						,       __A=30						,       __A=0						,       __A=1						,       __A=2						,       __A=None						,       )						->    Union[str, Any]:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=parent
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=batch_size
              SCREAMING_SNAKE_CASE_							:				str		=decoder_seq_length
              # For common tests
              SCREAMING_SNAKE_CASE_							:				Tuple		=self.decoder_seq_length
              SCREAMING_SNAKE_CASE_							:				Tuple		=is_training
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=use_attention_mask
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=use_labels

              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=vocab_size
              SCREAMING_SNAKE_CASE_							:				Dict		=d_model
              SCREAMING_SNAKE_CASE_							:				Dict		=d_model
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=decoder_layers
              SCREAMING_SNAKE_CASE_							:				str		=decoder_layers
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=decoder_ffn_dim
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=decoder_attention_heads
              SCREAMING_SNAKE_CASE_							:				str		=decoder_attention_heads
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=eos_token_id
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=bos_token_id
              SCREAMING_SNAKE_CASE_							:				Tuple		=pad_token_id
              SCREAMING_SNAKE_CASE_							:				Any		=decoder_start_token_id
              SCREAMING_SNAKE_CASE_							:				str		=use_cache
              SCREAMING_SNAKE_CASE_							:				Any		=max_position_embeddings

              SCREAMING_SNAKE_CASE_							:				Any		=None
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=decoder_seq_length
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=2
              SCREAMING_SNAKE_CASE_							:				int		=1


       def 	_snake_case				(    self				)						->    List[str]:
              SCREAMING_SNAKE_CASE_							:				Any		=ids_tensor([self.batch_size, self.decoder_seq_length]						,       self.vocab_size				)

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=None
              if self.use_attention_mask:
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=ids_tensor([self.batch_size, self.decoder_seq_length]						,       vocab_size=2				)

              SCREAMING_SNAKE_CASE_							:				Any		=None
              if self.use_labels:
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=ids_tensor([self.batch_size, self.decoder_seq_length]						,       self.vocab_size				)

              SCREAMING_SNAKE_CASE_							:				int		=TrOCRConfig(
                  vocab_size=self.vocab_size						,       d_model=self.d_model						,       decoder_layers=self.decoder_layers						,       decoder_ffn_dim=self.decoder_ffn_dim						,       decoder_attention_heads=self.decoder_attention_heads						,       eos_token_id=self.eos_token_id						,       bos_token_id=self.bos_token_id						,       use_cache=self.use_cache						,       pad_token_id=self.pad_token_id						,       decoder_start_token_id=self.decoder_start_token_id						,       max_position_embeddings=self.max_position_embeddings						,       )

              return (config, input_ids, attention_mask, lm_labels)


       def 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A						,       )						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=True
              SCREAMING_SNAKE_CASE_							:				int		=TrOCRDecoder(config=__A				).to(__A				).eval()
              SCREAMING_SNAKE_CASE_							:				Tuple		=input_ids[:2]

              input_ids[input_ids == 0] += 1
              # first forward pass
              SCREAMING_SNAKE_CASE_							:				Tuple		=model(__A						,       use_cache=__A				)
              SCREAMING_SNAKE_CASE_							:				Any		=model(__A				)
              SCREAMING_SNAKE_CASE_							:				List[str]		=model(__A						,       use_cache=__A				)

              self.parent.assertTrue(len(__A				) == len(__A				)				)
              self.parent.assertTrue(len(__A				) == len(__A				) + 1				)

              SCREAMING_SNAKE_CASE_							:				str		=outputs['''past_key_values''']

              # create hypothetical next token and extent to next_input_ids
              SCREAMING_SNAKE_CASE_							:				str		=ids_tensor((2, 1)						,       config.vocab_size - 1				) + 1

              # append to next input_ids and
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=torch.cat([input_ids, next_tokens]						,       dim=-1				)

              SCREAMING_SNAKE_CASE_							:				str		=model(__A				)['''last_hidden_state''']
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=model(__A						,       past_key_values=__A				)['''last_hidden_state''']

              # select random slice
              SCREAMING_SNAKE_CASE_							:				str		=ids_tensor((1,)						,       output_from_past.shape[-1]				).item()
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=output_from_past[:, 0, random_slice_idx].detach()

              # test that outputs are equal for slice
              assert torch.allclose(__A						,       __A						,       atol=1e-3				)





       def 	_snake_case				(    self				)						->    List[str]:
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.prepare_config_and_inputs()
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Optional[int]		=config_and_inputs

              SCREAMING_SNAKE_CASE_							:				Tuple		={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
              return config, inputs_dict






@require_torch
class 			lowercase_		(      A  ,      A  ,      A  ,      unittest.TestCase       ):
       __lowerCamelCase				       =						(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
       __lowerCamelCase				       =						(TrOCRForCausalLM,) if is_torch_available() else ()
       __lowerCamelCase				       =						{"text-generation": TrOCRForCausalLM} if is_torch_available() else {}
       __lowerCamelCase				       =						True
       __lowerCamelCase				       =						False


       def 	_snake_case				(    self				)						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_							:				Dict		=TrOCRStandaloneDecoderModelTester(self						,       is_training=__A				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=ConfigTester(self						,       config_class=__A				)


       def 	_snake_case				(    self				)						->    int:
              pass


       def 	_snake_case				(    self				)						->    str:
              pass


       def 	_snake_case				(    self				)						->    Tuple:
              pass


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


       def 	_snake_case				(    self				)						->    int:
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_decoder_model_past(*__A				)


       def 	_snake_case				(    self				)						->    List[str]:
              return





       @unittest.skip('''The model doesn\'t support left padding'''				)  # and it's not used enough to be worth fixing :)
       def 	_snake_case				(    self				)						->    str:
              pass


from dataclasses import asdict, dataclass
from typing import Optional

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


_lowercase											=    logging.get_logger(__name__)

# TODO Update this
_lowercase											=    {
    """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 			lowercase_		(      A       ):
       __lowerCamelCase				       =						"esm"


       def __init__(    self						,       __A=None						,       __A=None						,       __A=None						,       __A=768						,       __A=12						,       __A=12						,       __A=3_072						,       __A=0.1						,       __A=0.1						,       __A=1_026						,       __A=0.02						,       __A=1e-1_2						,       __A="absolute"						,       __A=True						,       __A=None						,       __A=False						,       __A=False						,       __A=None						,       __A=None						,       **__A						,       )						->    int:
              super().__init__(pad_token_id=__A						,       mask_token_id=__A						,       **__A				)

              SCREAMING_SNAKE_CASE_							:				Any		=vocab_size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=hidden_size
              SCREAMING_SNAKE_CASE_							:				int		=num_hidden_layers
              SCREAMING_SNAKE_CASE_							:				int		=num_attention_heads
              SCREAMING_SNAKE_CASE_							:				List[str]		=intermediate_size
              SCREAMING_SNAKE_CASE_							:				int		=hidden_dropout_prob
              SCREAMING_SNAKE_CASE_							:				int		=attention_probs_dropout_prob
              SCREAMING_SNAKE_CASE_							:				List[Any]		=max_position_embeddings
              SCREAMING_SNAKE_CASE_							:				List[Any]		=initializer_range
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=layer_norm_eps
              SCREAMING_SNAKE_CASE_							:				str		=position_embedding_type
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=use_cache
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=emb_layer_norm_before
              SCREAMING_SNAKE_CASE_							:				Tuple		=token_dropout
              SCREAMING_SNAKE_CASE_							:				List[Any]		=is_folding_model
              if is_folding_model:
                     if esmfold_config is None:
                            logger.info('''No esmfold_config supplied for folding model, using default values.'''				)
                            SCREAMING_SNAKE_CASE_							:				Optional[int]		=EsmFoldConfig()
                     elif isinstance(__A						,       __A				):
                            SCREAMING_SNAKE_CASE_							:				List[Any]		=EsmFoldConfig(**__A				)
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=esmfold_config
                     if vocab_list is None:
                            logger.warning('''No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'''				)
                            SCREAMING_SNAKE_CASE_							:				Tuple		=get_default_vocab_list()
                     else:
                            SCREAMING_SNAKE_CASE_							:				Optional[Any]		=vocab_list
              else:
                     SCREAMING_SNAKE_CASE_							:				List[str]		=None
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=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 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				str		=super().to_dict()
              if isinstance(self.esmfold_config						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.esmfold_config.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						None
       __lowerCamelCase				       =						True
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						0

       __lowerCamelCase				       =						True
       __lowerCamelCase				       =						False

       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						None


       def 	_snake_case				(    self				)						->    Optional[Any]:
              if self.trunk is None:
                     SCREAMING_SNAKE_CASE_							:				Tuple		=TrunkConfig()
              elif isinstance(self.trunk						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				List[str]		=TrunkConfig(**self.trunk				)





       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				int		=asdict(self				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=self.trunk.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						4_8
       __lowerCamelCase				       =						1_0_2_4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						0
       __lowerCamelCase				       =						0
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						None


       def 	_snake_case				(    self				)						->    int:
              if self.structure_module is None:
                     SCREAMING_SNAKE_CASE_							:				Dict		=StructureModuleConfig()
              elif isinstance(self.structure_module						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=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}.'				)

              SCREAMING_SNAKE_CASE_							:				Tuple		=self.sequence_state_dim // self.sequence_head_width
              SCREAMING_SNAKE_CASE_							:				str		=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 	_snake_case				(    self				)						->    int:
              SCREAMING_SNAKE_CASE_							:				str		=asdict(self				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=self.structure_module.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						3_8_4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						1_6
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						1_2
       __lowerCamelCase				       =						4
       __lowerCamelCase				       =						8
       __lowerCamelCase				       =						0.1
       __lowerCamelCase				       =						8
       __lowerCamelCase				       =						1
       __lowerCamelCase				       =						2
       __lowerCamelCase				       =						7
       __lowerCamelCase				       =						1_0
       __lowerCamelCase				       =						1E-8
       __lowerCamelCase				       =						1E5


       def 	_snake_case				(    self				)						->    str:
              return asdict(self				)






def 							SCREAMING_SNAKE_CASE_	(					)			->						int:
       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>",
       )
