





import logging
import os
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

from filelock import FileLock

from transformers import PreTrainedTokenizer, is_tf_available, is_torch_available


_SCREAMING_SNAKE_CASE							=						logging.getLogger(__name__)
@dataclass
class SCREAMING_SNAKE_CASE_  :
						__magic_name__:	str
						__magic_name__:	List[str]
						__magic_name__:	Optional[List[str]]
@dataclass
class SCREAMING_SNAKE_CASE_  :
						__magic_name__:	List[int]
						__magic_name__:	List[int]
						__magic_name__:	Optional[List[int]]       					=				None
						__magic_name__:	Optional[List[int]]       					=				None
class SCREAMING_SNAKE_CASE_  (    __lowercase    ):
						__magic_name__:	Union[str, Any]       					=				"train"
						__magic_name__:	Dict       					=				"dev"
						__magic_name__:	List[Any]       					=				"test"
class SCREAMING_SNAKE_CASE_  :






						@staticmethod
						def UpperCAmelCase_ (					_A					:					int						,      _A					:					Union[str, Any]		)							->			List[InputExample]:







							"""simple docstring"""


							raise NotImplementedError






						@staticmethod
						def UpperCAmelCase_ (					_A					:					List[str]		)							->			List[str]:







							"""simple docstring"""


							raise NotImplementedError






						@staticmethod
						def UpperCAmelCase_ (					_A					:					Tuple						,      _A					:					str						,      _A					:					int						,      _A					:					List[str]						,      _A					:					Any=False						,      _A					:					List[str]="[CLS]"						,      _A					:					List[Any]=1						,      _A					:					Dict="[SEP]"						,      _A					:					List[Any]=False						,      _A					:					Optional[int]=False						,      _A					:					Tuple=0						,      _A					:					Any=0						,      _A					:					Optional[Any]=-100						,      _A					:					str=0						,      _A					:					Optional[Any]=True						,      )							->			List[InputFeatures]:







							"""simple docstring"""


							snake_case_						:	Union[str, Any]      =		{label: i for i, label in enumerate(_a		)}

							snake_case_						:	Optional[Any]      =		[]
							for ex_index, example in enumerate(_a		):
								if ex_index % 10000 == 0:
									logger.info('Writing example %d of %d'						,      _a						,      len(_a		)		)

								snake_case_						:	List[str]      =		[]
								snake_case_						:	int      =		[]
								for word, label in zip(example.words						,      example.labels		):
									snake_case_						:	List[Any]      =		tokenizer.tokenize(_a		)

									# bert-base-multilingual-cased sometimes output "nothing ([]) when calling tokenize with just a space.
									if len(_a		) > 0:
										tokens.extend(_a		)
										# Use the real label id for the first token of the word, and padding ids for the remaining tokens
										label_ids.extend([label_map[label]] + [pad_token_label_id] * (len(_a		) - 1)		)

            # Account for [CLS] and [SEP] with "- 2" and with "- 3" for RoBERTa.
								snake_case_						:	Optional[int]      =		tokenizer.num_special_tokens_to_add()
								if len(_a		) > max_seq_length - special_tokens_count:
									snake_case_						:	str      =		tokens[: (max_seq_length - special_tokens_count)]
									snake_case_						:	int      =		label_ids[: (max_seq_length - special_tokens_count)]

								# The convention in BERT is:
								# (a) For sequence pairs:
								#  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
								#  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1
								# (b) For single sequences:
								#  tokens:   [CLS] the dog is hairy . [SEP]
								#  type_ids:   0   0   0   0  0     0   0
								#
								# Where "type_ids" are used to indicate whether this is the first
								# sequence or the second sequence. The embedding vectors for `type=0` and
								# `type=1` were learned during pre-training and are added to the wordpiece
								# embedding vector (and position vector). This is not *strictly* necessary
								# since the [SEP] token unambiguously separates the sequences, but it makes
								# it easier for the model to learn the concept of sequences.
								#
								# For classification tasks, the first vector (corresponding to [CLS]) is
								# used as the "sentence vector". Note that this only makes sense because
								# the entire model is fine-tuned.
								tokens += [sep_token]
								label_ids += [pad_token_label_id]
								if sep_token_extra:
									# roberta uses an extra separator b/w pairs of sentences
									tokens += [sep_token]
									label_ids += [pad_token_label_id]
								snake_case_						:	List[str]      =		[sequence_a_segment_id] * len(_a		)

								if cls_token_at_end:
									tokens += [cls_token]
									label_ids += [pad_token_label_id]
									segment_ids += [cls_token_segment_id]
								else:
									snake_case_						:	int      =		[cls_token] + tokens
									snake_case_						:	Optional[int]      =		[pad_token_label_id] + label_ids
									snake_case_						:	Dict      =		[cls_token_segment_id] + segment_ids

								snake_case_						:	List[Any]      =		tokenizer.convert_tokens_to_ids(_a		)

								# The mask has 1 for real tokens and 0 for padding tokens. Only real
								# tokens are attended to.
								snake_case_						:	Optional[int]      =		[1 if mask_padding_with_zero else 0] * len(_a		)

								# Zero-pad up to the sequence length.
								snake_case_						:	List[str]      =		max_seq_length - len(_a		)
								if pad_on_left:
									snake_case_						:	Union[str, Any]      =		([pad_token] * padding_length) + input_ids
									snake_case_						:	Dict      =		([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
									snake_case_						:	int      =		([pad_token_segment_id] * padding_length) + segment_ids
									snake_case_						:	Optional[int]      =		([pad_token_label_id] * padding_length) + label_ids
								else:
									input_ids += [pad_token] * padding_length
									input_mask += [0 if mask_padding_with_zero else 1] * padding_length
									segment_ids += [pad_token_segment_id] * padding_length
									label_ids += [pad_token_label_id] * padding_length

								assert len(_a		) == max_seq_length
								assert len(_a		) == max_seq_length
								assert len(_a		) == max_seq_length
								assert len(_a		) == max_seq_length

								if ex_index < 5:
									logger.info('*** Example ***'		)
									logger.info('guid: %s'						,      example.guid		)
									logger.info('tokens: %s'						,      ' '.join([str(_a		) for x in tokens]		)		)
									logger.info('input_ids: %s'						,      ' '.join([str(_a		) for x in input_ids]		)		)
									logger.info('input_mask: %s'						,      ' '.join([str(_a		) for x in input_mask]		)		)
									logger.info('segment_ids: %s'						,      ' '.join([str(_a		) for x in segment_ids]		)		)
									logger.info('label_ids: %s'						,      ' '.join([str(_a		) for x in label_ids]		)		)

								if "token_type_ids" not in tokenizer.model_input_names:
									snake_case_						:	Dict      =		None

								features.append(
								    InputFeatures(
								        input_ids=_a						,      attention_mask=_a						,      token_type_ids=_a						,      label_ids=_a		)		)
							return features


if is_torch_available():
						import torch
						from torch import nn
						from torch.utils.data import Dataset
						class SCREAMING_SNAKE_CASE_  (    __lowercase    ):
												__magic_name__:	List[InputFeatures]
												__magic_name__:	int       					=				nn.CrossEntropyLoss().ignore_index






												def __init__(					self					:					Tuple						,      _A					:					Dict						,      _A					:					Optional[Any]						,      _A					:					List[str]						,      _A					:					Optional[Any]						,      _A					:					List[Any]						,      _A					:					Any = None						,      _A					:					List[str]=False						,      _A					:					Any = Split.train						,      )							->			Tuple:







													"""simple docstring"""


													snake_case_						:	List[Any]      =		os.path.join(
													    _a						,      'cached_{}_{}_{}'.format(mode.value						,      tokenizer.__class__.__name__						,      str(_a		)		)						,      )

													# Make sure only the first process in distributed training processes the dataset,
													# and the others will use the cache.
													snake_case_						:	List[Any]      =		cached_features_file + '''.lock'''
													with FileLock(_a		):
														if os.path.exists(_a		) and not overwrite_cache:
															logger.info(F"""Loading features from cached file {cached_features_file}"""		)
															snake_case_						:	Optional[Any]      =		torch.load(_a		)
														else:
															logger.info(F"""Creating features from dataset file at {data_dir}"""		)
															snake_case_						:	List[Any]      =		token_classification_task.read_examples_from_file(_a						,      _a		)
															# TODO clean up all this to leverage built-in features of tokenizers
															snake_case_						:	List[Any]      =		token_classification_task.convert_examples_to_features(
															    _a						,      _a						,      _a						,      _a						,      cls_token_at_end=bool(model_type in ['xlnet']		)						,      cls_token=tokenizer.cls_token						,      cls_token_segment_id=2 if model_type in ['xlnet'] else 0						,      sep_token=tokenizer.sep_token						,      sep_token_extra=_a						,      pad_on_left=bool(tokenizer.padding_side == 'left'		)						,      pad_token=tokenizer.pad_token_id						,      pad_token_segment_id=tokenizer.pad_token_type_id						,      pad_token_label_id=self.pad_token_label_id						,      )
															logger.info(F"""Saving features into cached file {cached_features_file}"""		)
															torch.save(self.features						,      _a		)






												def __len__(					self					:					Dict		)							->			Optional[int]:







													"""simple docstring"""


													return len(self.features		)






												def __getitem__(					self					:					Any						,      _A					:					Optional[int]		)							->			InputFeatures:







													"""simple docstring"""


													return self.features[i]


if is_tf_available():
						import tensorflow as tf
						class SCREAMING_SNAKE_CASE_  :
												__magic_name__:	List[InputFeatures]
												__magic_name__:	int       					=				-100






												def __init__(					self					:					Dict						,      _A					:					List[str]						,      _A					:					Dict						,      _A					:					Optional[int]						,      _A					:					Union[str, Any]						,      _A					:					Optional[Any]						,      _A					:					str = None						,      _A					:					Any=False						,      _A					:					Any = Split.train						,      )							->			Dict:







													"""simple docstring"""


													snake_case_						:	str      =		token_classification_task.read_examples_from_file(_a						,      _a		)
													# TODO clean up all this to leverage built-in features of tokenizers
													snake_case_						:	Optional[int]      =		token_classification_task.convert_examples_to_features(
													    _a						,      _a						,      _a						,      _a						,      cls_token_at_end=bool(model_type in ['xlnet']		)						,      cls_token=tokenizer.cls_token						,      cls_token_segment_id=2 if model_type in ['xlnet'] else 0						,      sep_token=tokenizer.sep_token						,      sep_token_extra=_a						,      pad_on_left=bool(tokenizer.padding_side == 'left'		)						,      pad_token=tokenizer.pad_token_id						,      pad_token_segment_id=tokenizer.pad_token_type_id						,      pad_token_label_id=self.pad_token_label_id						,      )

													def gen():
														for ex in self.features:
															if ex.token_type_ids is None:
																yield (
																    {"input_ids": ex.input_ids, "attention_mask": ex.attention_mask},
																    ex.label_ids,
																)
															else:
																yield (
																    {
																        "input_ids": ex.input_ids,
																        "attention_mask": ex.attention_mask,
																        "token_type_ids": ex.token_type_ids,
																    },
																    ex.label_ids,
																)

													if "token_type_ids" not in tokenizer.model_input_names:
														snake_case_						:	Union[str, Any]      =		tf.data.Dataset.from_generator(
														    _a						,      ({'input_ids': tf.intaa, 'attention_mask': tf.intaa}, tf.intaa)						,      (
														        {'input_ids': tf.TensorShape([None]		), 'attention_mask': tf.TensorShape([None]		)},
														        tf.TensorShape([None]		),
														    )						,      )
													else:
														snake_case_						:	Union[str, Any]      =		tf.data.Dataset.from_generator(
														    _a						,      ({'input_ids': tf.intaa, 'attention_mask': tf.intaa, 'token_type_ids': tf.intaa}, tf.intaa)						,      (
														        {
														            'input_ids': tf.TensorShape([None]		),
														            'attention_mask': tf.TensorShape([None]		),
														            'token_type_ids': tf.TensorShape([None]		),
														        },
														        tf.TensorShape([None]		),
														    )						,      )






												def UpperCAmelCase_ (					self					:					Optional[Any]		)							->			str:







													"""simple docstring"""


													snake_case_						:	Dict      =		self.dataset.apply(tf.data.experimental.assert_cardinality(len(self.features		)		)		)

													return self.dataset






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







													"""simple docstring"""


													return len(self.features		)






												def __getitem__(					self					:					List[Any]						,      _A					:					str		)							->			InputFeatures:







													"""simple docstring"""


													return self.features[i]




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       ,							_a       ,							_a=1_3       ,							_a=7       ,							_a=False       ,							_a=True       ,							_a=False       ,							_a=False       ,							_a=1_9       ,							_a=3_2       ,							_a=5       ,							_a=4       ,							_a=3_7       ,							_a="gelu"       ,							_a=0.1       ,							_a=0.1       ,							_a=5_1_2       ,							_a=1_6       ,							_a=2       ,							_a=0.02       ,							_a=3       ,							_a=4       ,							_a=None       ,							)    ->    Union[str, Any]:
												_a	:    Optional[Any]         =		parent
												_a	:    Union[str, Any]         =		batch_size
												_a	:    List[Any]         =		seq_length
												_a	:    Dict         =		is_training
												_a	:    int         =		use_input_mask
												_a	:    str         =		use_token_type_ids
												_a	:    Any         =		use_labels
												_a	:    List[Any]         =		vocab_size
												_a	:    Any         =		hidden_size
												_a	:    int         =		num_hidden_layers
												_a	:    str         =		num_attention_heads
												_a	:    Dict         =		intermediate_size
												_a	:    List[str]         =		hidden_act
												_a	:    Optional[Any]         =		hidden_dropout_prob
												_a	:    Optional[Any]         =		attention_probs_dropout_prob
												_a	:    int         =		max_position_embeddings
												_a	:    Tuple         =		type_vocab_size
												_a	:    str         =		type_sequence_label_size
												_a	:    Any         =		initializer_range
												_a	:    Union[str, Any]         =		num_labels
												_a	:    Dict         =		num_choices
												_a	:    Union[str, Any]         =		scope


							def 				__lowercase				(			self       )    ->    List[Any]:
												_a	:    str         =		ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size       )

												_a	:    Dict         =		None
												if self.use_input_mask:
																	_a	:    int         =		random_attention_mask([self.batch_size, self.seq_length]       )

												_a	:    List[Any]         =		None
												_a	:    Tuple         =		None
												_a	:    Any         =		None
												if self.use_labels:
																	_a	:    Tuple         =		ids_tensor([self.batch_size]       ,							self.type_sequence_label_size       )
																	_a	:    Union[str, Any]         =		ids_tensor([self.batch_size, self.seq_length]       ,							self.num_labels       )
																	_a	:    List[Any]         =		ids_tensor([self.batch_size]       ,							self.num_choices       )

												_a	:    str         =		self.get_config()

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


							def 				__lowercase				(			self       )    ->    str:
												_a	:    Optional[int]         =		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=_a       ,							esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}       ,							)
												return config


							def 				__lowercase				(			self       ,							_a       ,							_a       ,							_a       ,							_a       ,							_a       ,							_a       )    ->    str:
												_a	:    Union[str, Any]         =		EsmForProteinFolding(config=_a       ).float()
												model.to(_a       )
												model.eval()
												_a	:    str         =		model(_a       ,							attention_mask=_a       )
												_a	:    Union[str, Any]         =		model(_a       )
												_a	:    Optional[int]         =		model(_a       )

												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 				__lowercase				(			self       )    ->    str:
												_a	:    List[str]         =		self.prepare_config_and_inputs()
												(
												    (
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					
												)	:    Optional[Any]         =		config_and_inputs
												_a	:    List[Any]         =		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
												return config, inputs_dict



@require_torch
class 				UpperCAmelCase_		(			__lowercase				,       __lowercase				,       unittest.TestCase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Any	    =    False

							UpperCAmelCase__       :						Any	    =    (EsmForProteinFolding,) if is_torch_available() else ()
							UpperCAmelCase__       :						Union[str, Any]	    =    ()
							UpperCAmelCase__       :						int	    =    {} if is_torch_available() else {}
							UpperCAmelCase__       :						Optional[int]	    =    False


							def 				__lowercase				(			self       )    ->    List[Any]:
												_a	:    Optional[int]         =		EsmFoldModelTester(self       )
												_a	:    Dict         =		ConfigTester(self       ,							config_class=_a       ,							hidden_size=3_7       )


							def 				__lowercase				(			self       )    ->    List[str]:
												self.config_tester.run_common_tests()


							def 				__lowercase				(			self       )    ->    str:
												_a	:    str         =		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_a       )


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


							@unittest.skip
							def 				__lowercase				(			self       )    ->    List[str]:
												pass


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


							@unittest.skip('''Esm does not support embedding resizing'''       )
							def 				__lowercase				(			self       )    ->    Optional[Any]:
												pass


							@unittest.skip('''ESMFold does not support passing input embeds!'''       )
							def 				__lowercase				(			self       )    ->    int:
												pass


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


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


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


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


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


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


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


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


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


							@unittest.skip('''ESMFold does not support input chunking.'''       )
							def 				__lowercase				(			self       )    ->    Tuple:
												pass


							@unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''       )
							def 				__lowercase				(			self       )    ->    Optional[Any]:
												pass


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


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


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


							@unittest.skip('''ESMFold doesn\'t support data parallel.'''       )
							def 				__lowercase				(			self       )    ->    Union[str, Any]:
												pass




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



@require_torch
class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							@slow
							def 				__lowercase				(			self       )    ->    Optional[int]:
												_a	:    Dict         =		EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''       ).float()
												model.eval()
												_a	:    Tuple         =		torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]       )
												_a	:    Optional[Any]         =		model(_a       )['''positions''']
												_a	:    Union[str, Any]         =		torch.tensor([2.5828, 0.7993, -10.9334]       ,							dtype=torch.floataa       )
												self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]       ,							_a       ,							atol=1e-4       )       )



