


import enum
import warnings

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
							import tensorflow as tf

							from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

if is_torch_available():
							from ..models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

a					  =     logging.get_logger(__name__)





class        UpperCamelCase__   (	enum.Enum   ):
					__SCREAMING_SNAKE_CASE       :  List[Any]						=   0
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   1





@add_end_docstrings(__magic_name__   )
class        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  List[str]						=   'generated'

					def __init__(     self       :   Any							,					*UpperCamelCase__       :   Dict							,					**UpperCamelCase__       :   Tuple       ):






										'''simple docstring'''






										super().__init__(*lowercase_							,					**lowercase_       )

										self.check_model_type(
										    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
										    if self.framework == """tf"""
										    else MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING       )

					def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   Union[str, Any]=None							,					UpperCamelCase__       :   Optional[Any]=None							,					UpperCamelCase__       :   Optional[int]=None							,					UpperCamelCase__       :   Optional[Any]=None							,					UpperCamelCase__       :   Any=None							,					UpperCamelCase__       :   Union[str, Any]=None							,					**UpperCamelCase__       :   Optional[Any]							,					):






										'''simple docstring'''






										lowercase_  							=			{}
										if truncation is not None:
															lowercase_  							=			truncation

										lowercase_  							=			generate_kwargs

										lowercase_  							=			{}
										if return_tensors is not None and return_type is None:
															lowercase_  							=			ReturnType.TENSORS if return_tensors else ReturnType.TEXT
										if return_type is not None:
															lowercase_  							=			return_type

										if clean_up_tokenization_spaces is not None:
															lowercase_  							=			clean_up_tokenization_spaces

										if stop_sequence is not None:
															lowercase_  							=			self.tokenizer.encode(lowercase_							,					add_special_tokens=lowercase_       )
															if len(lowercase_       ) > 1:
																				warnings.warn(
																				    """Stopping on a multiple token sequence is not yet supported on transformers. The first token of"""
																				    """ the stop sequence will be used as the stop sequence string in the interim."""       )
															lowercase_  							=			stop_sequence_ids[0]

										return preprocess_params, forward_params, postprocess_params

					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										return True

					def   UpperCAmelCase__       (     self       :   Dict							,					*UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.model.config.prefix if self.model.config.prefix is not None else """"""
										if isinstance(args[0]							,					lowercase_       ):
															if self.tokenizer.pad_token_id is None:
																				raise ValueError("""Please make sure that the tokenizer has a pad_token_id when using a batch input"""       )
															lowercase_  							=			([prefix + arg for arg in args[0]],)
															lowercase_  							=			True

										elif isinstance(args[0]							,					lowercase_       ):
															lowercase_  							=			(prefix + args[0],)
															lowercase_  							=			False
										else:
															raise ValueError(
															    F''' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`'''       )
										lowercase_  							=			self.tokenizer(*lowercase_							,					padding=lowercase_							,					truncation=lowercase_							,					return_tensors=self.framework       )
										# This is produced by tokenizers but is an invalid generate kwargs
										if "token_type_ids" in inputs:
															del inputs["token_type_ids"]
										return inputs

					def __call__(     self       :   List[Any]							,					*UpperCamelCase__       :   Any							,					**UpperCamelCase__       :   int       ):






										'''simple docstring'''






										lowercase_  							=			super().__call__(*lowercase_							,					**lowercase_       )
										if (
										    isinstance(args[0]							,					lowercase_       )
										    and all(isinstance(lowercase_							,					lowercase_       ) for el in args[0]       )
										    and all(len(lowercase_       ) == 1 for res in result       )
										):
															return [res[0] for res in result]
										return result

					def   UpperCAmelCase__       (     self       :   Tuple							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   str=TruncationStrategy.DO_NOT_TRUNCATE							,					**UpperCamelCase__       :   Dict       ):






										'''simple docstring'''






										lowercase_  							=			self._parse_and_tokenize(lowercase_							,					truncation=lowercase_							,					**lowercase_       )
										return inputs

					def   UpperCAmelCase__       (     self       :   str							,					UpperCamelCase__       :   str							,					**UpperCamelCase__       :   str       ):






										'''simple docstring'''






										if self.framework == "pt":
															lowercase_							,					lowercase_  							=			model_inputs["""input_ids"""].shape
										elif self.framework == "tf":
															lowercase_							,					lowercase_  							=			tf.shape(model_inputs["""input_ids"""]       ).numpy()

										lowercase_  							=			generate_kwargs.get("""min_length"""							,					self.model.config.min_length       )
										lowercase_  							=			generate_kwargs.get("""max_length"""							,					self.model.config.max_length       )
										self.check_inputs(lowercase_							,					generate_kwargs["""min_length"""]							,					generate_kwargs["""max_length"""]       )
										lowercase_  							=			self.model.generate(**lowercase_							,					**lowercase_       )
										lowercase_  							=			output_ids.shape[0]
										if self.framework == "pt":
															lowercase_  							=			output_ids.reshape(lowercase_							,					out_b // in_b							,					*output_ids.shape[1:]       )
										elif self.framework == "tf":
															lowercase_  							=			tf.reshape(lowercase_							,					(in_b, out_b // in_b, *output_ids.shape[1:])       )
										return {"output_ids": output_ids}







					def   UpperCAmelCase__       (     self       :   Dict							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   int=ReturnType.TEXT							,					UpperCamelCase__       :   int=False       ):






										'''simple docstring'''






										lowercase_  							=			[]
										for output_ids in model_outputs["output_ids"][0]:
															if return_type == ReturnType.TENSORS:
																				lowercase_  							=			{F'''{self.return_name}_token_ids''': output_ids}
															elif return_type == ReturnType.TEXT:
																				lowercase_  							=			{
																				    F'''{self.return_name}_text''': self.tokenizer.decode(
																				        lowercase_							,					skip_special_tokens=lowercase_							,					clean_up_tokenization_spaces=lowercase_							,					)
																				}
															records.append(lowercase_       )
										return records





@add_end_docstrings(__magic_name__   )
class        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  List[str]						=   'summary'

					def __call__(     self       :   Optional[Any]							,					*UpperCamelCase__       :   int							,					**UpperCamelCase__       :   Dict       ):






										'''simple docstring'''






										return super().__call__(*lowercase_							,					**lowercase_       )







					def   UpperCAmelCase__       (     self       :   List[str]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										if max_length < min_length:
															logger.warning(F'''Your min_length={min_length} must be inferior than your max_length={max_length}.'''       )

										if input_length < max_length:
															logger.warning(
															    F'''Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is '''
															    """a summarization task, where outputs shorter than the input are typically wanted, you might """
															    F'''consider decreasing max_length manually, e.g. summarizer(\'...\', max_length={input_length//2})'''       )





@add_end_docstrings(__magic_name__   )
class        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   'translation'

					def   UpperCAmelCase__       (     self       :   Dict							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										if input_length > 0.9 * max_length:
															logger.warning(
															    F'''Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider '''
															    """increasing your max_length manually, e.g. translator('...', max_length=400)"""       )
										return True

					def   UpperCAmelCase__       (     self       :   Tuple							,					*UpperCamelCase__       :   Any							,					UpperCamelCase__       :   List[Any]=TruncationStrategy.DO_NOT_TRUNCATE							,					UpperCamelCase__       :   Any=None							,					UpperCamelCase__       :   Optional[Any]=None       ):






										'''simple docstring'''






										if getattr(self.tokenizer							,					"""_build_translation_inputs"""							,					lowercase_       ):
															return self.tokenizer._build_translation_inputs(
															    *lowercase_							,					return_tensors=self.framework							,					truncation=lowercase_							,					src_lang=lowercase_							,					tgt_lang=lowercase_       )
										else:
															return super()._parse_and_tokenize(*lowercase_							,					truncation=lowercase_       )

					def   UpperCAmelCase__       (     self       :   List[str]							,					UpperCamelCase__       :   Dict=None							,					UpperCamelCase__       :   str=None							,					**UpperCamelCase__       :   List[Any]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_							,					lowercase_  							=			super()._sanitize_parameters(**lowercase_       )
										if src_lang is not None:
															lowercase_  							=			src_lang
										if tgt_lang is not None:
															lowercase_  							=			tgt_lang
										if src_lang is None and tgt_lang is None:
															# Backward compatibility, direct arguments use is preferred.
															lowercase_  							=			kwargs.get("""task"""							,					self.task       )
															lowercase_  							=			task.split("""_"""       )
															if task and len(lowercase_       ) == 4:
																				# translation, XX, to YY
																				lowercase_  							=			items[1]
																				lowercase_  							=			items[3]
										return preprocess_params, forward_params, postprocess_params







					def __call__(     self       :   List[str]							,					*UpperCamelCase__       :   List[str]							,					**UpperCamelCase__       :   str       ):






										'''simple docstring'''






										return super().__call__(*lowercase_							,					**lowercase_       )






from dataclasses import asdict, dataclass
from typing import Optional

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


a					  =     logging.get_logger(__name__)

# TODO Update this
a					  =     {
    '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        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  Tuple						=   'esm'

					def __init__(     self       :   Optional[Any]							,					UpperCamelCase__       :   Dict=None							,					UpperCamelCase__       :   Tuple=None							,					UpperCamelCase__       :   List[Any]=None							,					UpperCamelCase__       :   List[Any]=768							,					UpperCamelCase__       :   Tuple=12							,					UpperCamelCase__       :   Optional[int]=12							,					UpperCamelCase__       :   Dict=3_072							,					UpperCamelCase__       :   Union[str, Any]=0.1							,					UpperCamelCase__       :   Dict=0.1							,					UpperCamelCase__       :   Optional[int]=1_026							,					UpperCamelCase__       :   Any=0.02							,					UpperCamelCase__       :   Dict=1e-12							,					UpperCamelCase__       :   List[str]="absolute"							,					UpperCamelCase__       :   Union[str, Any]=True							,					UpperCamelCase__       :   Tuple=None							,					UpperCamelCase__       :   Union[str, Any]=False							,					UpperCamelCase__       :   Optional[int]=False							,					UpperCamelCase__       :   Dict=None							,					UpperCamelCase__       :   str=None							,					**UpperCamelCase__       :   str							,					):






										'''simple docstring'''






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

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







					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			super().to_dict()
										if isinstance(self.esmfold_config							,					UpperCamelCase__       ):
															lowercase_  							=			self.esmfold_config.to_dict()
										return output







@dataclass
class        UpperCamelCase__   :
					__SCREAMING_SNAKE_CASE       :  str						=   None
					__SCREAMING_SNAKE_CASE       :  bool						=   True
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  float						=   0

					__SCREAMING_SNAKE_CASE       :  bool						=   True
					__SCREAMING_SNAKE_CASE       :  bool						=   False

					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  "TrunkConfig"						=   None

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										if self.trunk is None:
															lowercase_  							=			TrunkConfig()
										elif isinstance(self.trunk							,					UpperCamelCase__       ):
															lowercase_  							=			TrunkConfig(**self.trunk       )







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






										'''simple docstring'''






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







@dataclass
class        UpperCamelCase__   :
					__SCREAMING_SNAKE_CASE       :  int						=   48
					__SCREAMING_SNAKE_CASE       :  int						=   1024
					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  int						=   32
					__SCREAMING_SNAKE_CASE       :  int						=   32
					__SCREAMING_SNAKE_CASE       :  int						=   32
					__SCREAMING_SNAKE_CASE       :  float						=   0
					__SCREAMING_SNAKE_CASE       :  float						=   0
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  int						=   4
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   128
					__SCREAMING_SNAKE_CASE       :  "StructureModuleConfig"						=   None

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






										'''simple docstring'''






										if self.structure_module is None:
															lowercase_  							=			StructureModuleConfig()
										elif isinstance(self.structure_module							,					UpperCamelCase__       ):
															lowercase_  							=			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}.'''       )

										lowercase_  							=			self.sequence_state_dim // self.sequence_head_width
										lowercase_  							=			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   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






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







@dataclass
class        UpperCamelCase__   :
					__SCREAMING_SNAKE_CASE       :  int						=   384
					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  int						=   16
					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  int						=   12
					__SCREAMING_SNAKE_CASE       :  int						=   4
					__SCREAMING_SNAKE_CASE       :  int						=   8
					__SCREAMING_SNAKE_CASE       :  float						=   0.1
					__SCREAMING_SNAKE_CASE       :  int						=   8
					__SCREAMING_SNAKE_CASE       :  int						=   1
					__SCREAMING_SNAKE_CASE       :  int						=   2
					__SCREAMING_SNAKE_CASE       :  int						=   7
					__SCREAMING_SNAKE_CASE       :  int						=   10
					__SCREAMING_SNAKE_CASE       :  float						=   1e-8
					__SCREAMING_SNAKE_CASE       :  float						=   1e5

					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''simple docstring'''






										return asdict(self       )



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


