


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    UpperCAmelCase					:




   def __init__(						self							:		Optional[Any]       ,       __magic_name__							:		Any       ,       __magic_name__							:		Tuple=1_3       ,       __magic_name__							:		List[str]=7       ,       __magic_name__							:		str=False       ,       __magic_name__							:		Tuple=True       ,       __magic_name__							:		List[str]=False       ,       __magic_name__							:		Optional[int]=True       ,       __magic_name__							:		List[str]=3_3       ,       __magic_name__							:		Dict=3_2       ,       __magic_name__							:		Optional[Any]=5       ,       __magic_name__							:		Optional[int]=4       ,       __magic_name__							:		str=3_7       ,       __magic_name__							:		Optional[Any]="gelu"       ,       __magic_name__							:		Optional[int]=0.1       ,       __magic_name__							:		List[Any]=0.1       ,       __magic_name__							:		Tuple=5_1_2       ,       __magic_name__							:		int=1_6       ,       __magic_name__							:		List[str]=2       ,       __magic_name__							:		Tuple=0.02       ,       __magic_name__							:		Union[str, Any]=3       ,       __magic_name__							:		Union[str, Any]=4       ,       __magic_name__							:		List[Any]=None       ,       ):







      """simple docstring"""







      UpperCamelCase         =					parent
      UpperCamelCase         =					batch_size
      UpperCamelCase         =					seq_length
      UpperCamelCase         =					is_training
      UpperCamelCase         =					use_input_mask
      UpperCamelCase         =					use_token_type_ids
      UpperCamelCase         =					use_labels
      UpperCamelCase         =					vocab_size
      UpperCamelCase         =					hidden_size
      UpperCamelCase         =					num_hidden_layers
      UpperCamelCase         =					num_attention_heads
      UpperCamelCase         =					intermediate_size
      UpperCamelCase         =					hidden_act
      UpperCamelCase         =					hidden_dropout_prob
      UpperCamelCase         =					attention_probs_dropout_prob
      UpperCamelCase         =					max_position_embeddings
      UpperCamelCase         =					type_vocab_size
      UpperCamelCase         =					type_sequence_label_size
      UpperCamelCase         =					initializer_range
      UpperCamelCase         =					num_labels
      UpperCamelCase         =					num_choices
      UpperCamelCase         =					scope




   def   lowerCamelCase_					(						self							:		int						):







      """simple docstring"""







      UpperCamelCase         =					ids_tensor([self.batch_size, self.seq_length]       ,       self.vocab_size						)

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

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

      UpperCamelCase         =					self.get_config()

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




   def   lowerCamelCase_					(						self							:		Optional[Any]						):







      """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   lowerCamelCase_					(						self							:		Any       ,       __magic_name__							:		Any       ,       __magic_name__							:		Union[str, Any]       ,       __magic_name__							:		Union[str, Any]       ,       __magic_name__							:		Union[str, Any]       ,       __magic_name__							:		List[Any]       ,       __magic_name__							:		Tuple						):







      """simple docstring"""







      UpperCamelCase         =					EsmModel(config=__magic_name__						)
      model.to(__magic_name__						)
      model.eval()
      UpperCamelCase         =					model(__magic_name__       ,       attention_mask=__magic_name__						)
      UpperCamelCase         =					model(__magic_name__						)
      UpperCamelCase         =					model(__magic_name__						)

      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   lowerCamelCase_					(						self							:		Optional[int]       ,       __magic_name__							:		Optional[int]       ,       __magic_name__							:		List[str]       ,       __magic_name__							:		Optional[int]       ,       __magic_name__							:		Optional[Any]       ,       __magic_name__							:		List[Any]       ,       __magic_name__							:		Union[str, Any]						):







      """simple docstring"""







      UpperCamelCase         =					EsmForMaskedLM(config=__magic_name__						)
      model.to(__magic_name__						)
      model.eval()
      UpperCamelCase         =					model(__magic_name__       ,       attention_mask=__magic_name__       ,       labels=__magic_name__						)
      self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.seq_length, self.vocab_size)						)




   def   lowerCamelCase_					(						self							:		Any       ,       __magic_name__							:		Optional[int]       ,       __magic_name__							:		List[str]       ,       __magic_name__							:		Optional[Any]       ,       __magic_name__							:		Optional[Any]       ,       __magic_name__							:		Any       ,       __magic_name__							:		str						):







      """simple docstring"""







      UpperCamelCase         =					self.num_labels
      UpperCamelCase         =					EsmForTokenClassification(config=__magic_name__						)
      model.to(__magic_name__						)
      model.eval()
      UpperCamelCase         =					model(__magic_name__       ,       attention_mask=__magic_name__       ,       labels=__magic_name__						)
      self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.seq_length, self.num_labels)						)





   def   lowerCamelCase_					(						self							:		Tuple						):







      """simple docstring"""







      UpperCamelCase         =					self.prepare_config_and_inputs()
      (
          (
          UpperCamelCase  
      )  ,  (
          UpperCamelCase  
      )  ,  (
          UpperCamelCase  
      )  ,  (
          UpperCamelCase  
      )  ,  (
          UpperCamelCase  
      )  ,  (
          UpperCamelCase  
      )  ,    
      )         =					config_and_inputs
      UpperCamelCase         =					{"""input_ids""": input_ids, """attention_mask""": input_mask}
      return config, inputs_dict




@require_torch
class    UpperCAmelCase					(       __snake_case      , __snake_case      , unittest.TestCase    ):
   lowercase						= False

   lowercase						= (
       (
           EsmForMaskedLM,
           EsmModel,
           EsmForSequenceClassification,
           EsmForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowercase						= ()
   lowercase						= (
       {
           """feature-extraction""": EsmModel,
           """fill-mask""": EsmForMaskedLM,
           """text-classification""": EsmForSequenceClassification,
           """token-classification""": EsmForTokenClassification,
           """zero-shot""": EsmForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   lowercase						= True




   def   lowerCamelCase_					(						self							:		Optional[Any]						):







      """simple docstring"""







      UpperCamelCase         =					EsmModelTester(self						)
      UpperCamelCase         =					ConfigTester(self       ,       config_class=__magic_name__       ,       hidden_size=3_7						)




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







      """simple docstring"""







      self.config_tester.run_common_tests()




   def   lowerCamelCase_					(						self							:		List[str]						):







      """simple docstring"""







      UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__magic_name__						)




   def   lowerCamelCase_					(						self							:		Optional[int]						):







      """simple docstring"""







      UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
         UpperCamelCase         =					type
         self.model_tester.create_and_check_model(*__magic_name__						)




   def   lowerCamelCase_					(						self							:		Optional[Any]						):







      """simple docstring"""







      UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*__magic_name__						)




   def   lowerCamelCase_					(						self							:		List[Any]						):







      """simple docstring"""







      UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*__magic_name__						)




   @slow
   def   lowerCamelCase_					(						self							:		str						):







      """simple docstring"""







      for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         UpperCamelCase         =					EsmModel.from_pretrained(__magic_name__						)
         self.assertIsNotNone(__magic_name__						)




   def   lowerCamelCase_					(						self							:		List[str]						):







      """simple docstring"""







      UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()[0]
      UpperCamelCase         =					EsmEmbeddings(config=__magic_name__						)

      UpperCamelCase         =					torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]						)
      UpperCamelCase         =					torch.as_tensor(
          [
              [
                  0 + model.padding_idx + 1,
                  1 + model.padding_idx + 1,
                  2 + model.padding_idx + 1,
                  model.padding_idx,
              ]
          ]						)
      UpperCamelCase         =					create_position_ids_from_input_ids(__magic_name__       ,       model.padding_idx						)
      self.assertEqual(position_ids.shape       ,       expected_positions.shape						)
      self.assertTrue(torch.all(torch.eq(__magic_name__       ,       __magic_name__						)						)						)




   def   lowerCamelCase_					(						self							:		Any						):







      """simple docstring"""







      UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()[0]
      UpperCamelCase         =					EsmEmbeddings(config=__magic_name__						)

      UpperCamelCase         =					torch.empty(2       ,       4       ,       3_0						)
      UpperCamelCase         =					[
          0 + embeddings.padding_idx + 1,
          1 + embeddings.padding_idx + 1,
          2 + embeddings.padding_idx + 1,
          3 + embeddings.padding_idx + 1,
      ]
      UpperCamelCase         =					torch.as_tensor([expected_single_positions, expected_single_positions]						)
      UpperCamelCase         =					embeddings.create_position_ids_from_inputs_embeds(__magic_name__						)
      self.assertEqual(position_ids.shape       ,       expected_positions.shape						)
      self.assertTrue(torch.all(torch.eq(__magic_name__       ,       __magic_name__						)						)						)




   @unittest.skip("""Esm does not support embedding resizing"""						)
   def   lowerCamelCase_					(						self							:		Optional[int]						):







      """simple docstring"""







      pass




   @unittest.skip("""Esm does not support embedding resizing"""						)
   def   lowerCamelCase_					(						self							:		List[Any]						):







      """simple docstring"""







      pass





   @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""						)
   def   lowerCamelCase_					(						self							:		Union[str, Any]						):







      """simple docstring"""







      pass




@require_torch
class    UpperCAmelCase					(       __snake_case    ):




   @slow
   def   lowerCamelCase_					(						self							:		List[Any]						):







      """simple docstring"""







      with torch.no_grad():
         UpperCamelCase         =					EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""						)
         model.eval()
         UpperCamelCase         =					torch.tensor([[0, 1, 2, 3, 4, 5]]						)
         UpperCamelCase         =					model(__magic_name__						)[0]

         UpperCamelCase         =					3_3

         UpperCamelCase         =					torch.Size((1, 6, vocab_size)						)
         self.assertEqual(output.shape       ,       __magic_name__						)

         UpperCamelCase         =					torch.tensor(
             [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]						)
         self.assertTrue(torch.allclose(output[:, :3, :3]       ,       __magic_name__       ,       atol=1e-4						)						)





   @slow
   def   lowerCamelCase_					(						self							:		int						):







      """simple docstring"""







      with torch.no_grad():
         UpperCamelCase         =					EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""						)
         model.eval()

         UpperCamelCase         =					torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]						)
         UpperCamelCase         =					model(__magic_name__						)[0]
         # compare the actual values for a slice.
         UpperCamelCase         =					torch.tensor(
             [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]						)
         self.assertTrue(torch.allclose(output[:, :3, :3]       ,       __magic_name__       ,       atol=1e-4						)						)








import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


__snake_case            =   logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
			import torch_xla.core.xla_model as xm
			import torch_xla.debug.metrics as met




class    UpperCAmelCase					(       __snake_case    ):




			def __init__(						self							:		Dict       ,       *__magic_name__							:		Union[str, Any]       ,       __magic_name__							:		Any=None       ,       __magic_name__							:		Union[str, Any]=None       ,       __magic_name__							:		Any=None       ,       **__magic_name__							:		Optional[Any]						):







						"""simple docstring"""







						super().__init__(*__magic_name__       ,       **__magic_name__						)
						UpperCamelCase         =					eval_examples
						UpperCamelCase         =					post_process_function
						UpperCamelCase         =					quant_trainer_args
						UpperCamelCase         =					1_2_8  # default number of calibration samples




			def   lowerCamelCase_					(						self							:		str       ,       __magic_name__							:		List[Any]=None						):







						"""simple docstring"""







						if calib_dataset is None and self.calib_dataset is None:
									raise ValueError("""Trainer: calibration requires an calib_dataset."""						)
						UpperCamelCase         =					calib_dataset if calib_dataset is not None else self.calib_dataset

						UpperCamelCase         =					self._remove_unused_columns(__magic_name__       ,       description="""Calibration"""						)

						return DataLoader(
						    __magic_name__       ,       batch_size=self.args.eval_batch_size       ,       collate_fn=self.data_collator       ,       drop_last=self.args.dataloader_drop_last       ,       num_workers=self.args.dataloader_num_workers       ,       pin_memory=self.args.dataloader_pin_memory       ,       shuffle=__magic_name__       ,       )




			def   lowerCamelCase_					(						self							:		Tuple       ,       __magic_name__							:		List[str]=None						):







						"""simple docstring"""







						UpperCamelCase         =					self.train_dataset if calib_dataset is None else calib_dataset
						UpperCamelCase         =					self.get_calib_dataloader(__magic_name__						)

						UpperCamelCase         =					self.model
						quant_trainer.configure_model(__magic_name__       ,       self.quant_trainer_args       ,       calib=__magic_name__						)
						model.eval()
						quant_trainer.enable_calibration(__magic_name__						)

						logger.info("""***** Running calibration *****"""						)
						logger.info(F'  Num examples = {self.calib_num}'						)
						logger.info(F'  Batch size = {calib_dataloader.batch_size}'						)

						for step, inputs in enumerate(__magic_name__						):
									# Prediction step
									UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase         =					self.prediction_step(__magic_name__       ,       __magic_name__       ,       prediction_loss_only=__magic_name__						)
									if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
												break

						quant_trainer.finish_calibration(__magic_name__       ,       self.quant_trainer_args						)
						UpperCamelCase         =					model




			def   lowerCamelCase_					(						self							:		Any       ,       __magic_name__							:		Optional[int]=None       ,       __magic_name__							:		Optional[Any]=None       ,       __magic_name__							:		Union[str, Any]=None       ,       __magic_name__							:		str = "eval"						):







						"""simple docstring"""







						UpperCamelCase         =					self.eval_dataset if eval_dataset is None else eval_dataset
						UpperCamelCase         =					self.get_eval_dataloader(__magic_name__						)
						UpperCamelCase         =					self.eval_examples if eval_examples is None else eval_examples

						# Temporarily disable metric computation, we will do it in the loop here.
						UpperCamelCase         =					self.compute_metrics
						UpperCamelCase         =					None
						UpperCamelCase         =					self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
						try:
									UpperCamelCase         =					eval_loop(
									    __magic_name__       ,       description="""Evaluation"""       ,       prediction_loss_only=True if compute_metrics is None else None       ,       ignore_keys=__magic_name__       ,       )
						finally:
									UpperCamelCase         =					compute_metrics

						if self.post_process_function is not None and self.compute_metrics is not None:
									UpperCamelCase         =					self.post_process_function(__magic_name__       ,       __magic_name__       ,       output.predictions						)
									UpperCamelCase         =					self.compute_metrics(__magic_name__						)

									# Prefix all keys with metric_key_prefix + '_'
									for key in list(metrics.keys()						):
												if not key.startswith(F'{metric_key_prefix}_'						):
															UpperCamelCase         =					metrics.pop(__magic_name__						)

									self.log(__magic_name__						)
						else:
									UpperCamelCase         =					{}

						if self.args.tpu_metrics_debug or self.args.debug:
									# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
									xm.master_print(met.metrics_report()						)

						UpperCamelCase         =					self.callback_handler.on_evaluate(self.args       ,       self.state       ,       self.control       ,       __magic_name__						)
						return metrics




			def   lowerCamelCase_					(						self							:		Dict       ,       __magic_name__							:		List[str]       ,       __magic_name__							:		List[Any]       ,       __magic_name__							:		List[Any]=None       ,       __magic_name__							:		str = "test"						):







						"""simple docstring"""







						UpperCamelCase         =					self.get_test_dataloader(__magic_name__						)

						# Temporarily disable metric computation, we will do it in the loop here.
						UpperCamelCase         =					self.compute_metrics
						UpperCamelCase         =					None
						UpperCamelCase         =					self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
						try:
									UpperCamelCase         =					eval_loop(
									    __magic_name__       ,       description="""Prediction"""       ,       prediction_loss_only=True if compute_metrics is None else None       ,       ignore_keys=__magic_name__       ,       )
						finally:
									UpperCamelCase         =					compute_metrics

						if self.post_process_function is None or self.compute_metrics is None:
									return output

						UpperCamelCase         =					self.post_process_function(__magic_name__       ,       __magic_name__       ,       output.predictions       ,       """predict"""						)
						UpperCamelCase         =					self.compute_metrics(__magic_name__						)

						# Prefix all keys with metric_key_prefix + '_'
						for key in list(metrics.keys()						):
									if not key.startswith(F'{metric_key_prefix}_'						):
												UpperCamelCase         =					metrics.pop(__magic_name__						)

						return PredictionOutput(predictions=predictions.predictions       ,       label_ids=predictions.label_ids       ,       metrics=__magic_name__						)





			def   lowerCamelCase_					(						self							:		List[str]       ,       __magic_name__							:		Optional[Any]="./"						):







						"""simple docstring"""







						UpperCamelCase         =					self.eval_dataset
						UpperCamelCase         =					self.get_eval_dataloader(__magic_name__						)

						UpperCamelCase         =					next(iter(__magic_name__						)						)

						# saving device - to make it consistent
						UpperCamelCase         =					torch.device("""cuda""" if torch.cuda.is_available() else """cpu"""						)

						# convert to tuple
						UpperCamelCase         =					tuple(v.to(__magic_name__						) for k, v in batch.items()						)

						logger.info("""Converting model to be onnx compatible"""						)
						from pytorch_quantization.nn import TensorQuantizer

						UpperCamelCase         =					True

						UpperCamelCase         =					self.model.to(__magic_name__						)

						model.eval()
						model.float()

						UpperCamelCase         =					model.module if hasattr(__magic_name__       ,       """module"""						) else model
						quant_trainer.configure_model(__magic_name__       ,       self.quant_trainer_args						)

						UpperCamelCase         =					os.path.join(__magic_name__       ,       """model.onnx"""						)
						logger.info(F'exporting model to {output_model_file}'						)

						UpperCamelCase         =					{0: """batch_size""", 1: """seq_len"""}

						torch.onnx.export(
						    __magic_name__       ,       __magic_name__       ,       __magic_name__       ,       export_params=__magic_name__       ,       opset_version=1_3       ,       do_constant_folding=__magic_name__       ,       input_names=["""input_ids""", """attention_mask""", """token_type_ids"""]       ,       output_names=["""output_start_logits""", """output_end_logits"""]       ,       dynamic_axes={
						        """input_ids""": axes,
						        """attention_mask""": axes,
						        """token_type_ids""": axes,
						        """output_start_logits""": axes,
						        """output_end_logits""": axes,
						    }       ,       verbose=__magic_name__       ,       )
						logger.info("""onnx export finished"""						)




