


"""simple docstring"""






import unittest

from transformers import AutoTokenizer, NystromformerConfig, is_torch_available
from transformers.testing_utils import 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 (
      NystromformerForMaskedLM,
      NystromformerForMultipleChoice,
      NystromformerForQuestionAnswering,
      NystromformerForSequenceClassification,
      NystromformerForTokenClassification,
      NystromformerModel,
  )
  from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST




class 							A_      :


  def __init__(			self:						Optional[Any]							,__lowerCAmelCase:						Any							,__lowerCAmelCase:						List[Any]=13							,__lowerCAmelCase:						Tuple=7							,__lowerCAmelCase:						List[str]=True							,__lowerCAmelCase:						Optional[Any]=True							,__lowerCAmelCase:						List[str]=True							,__lowerCAmelCase:						Optional[Any]=True							,__lowerCAmelCase:						Optional[Any]=99							,__lowerCAmelCase:						Optional[int]=32							,__lowerCAmelCase:						Dict=5							,__lowerCAmelCase:						Any=4							,__lowerCAmelCase:						int=37							,__lowerCAmelCase:						Tuple="gelu"							,__lowerCAmelCase:						Optional[int]=0.1							,__lowerCAmelCase:						List[str]=0.1							,__lowerCAmelCase:						Union[str, Any]=512							,__lowerCAmelCase:						Optional[int]=16							,__lowerCAmelCase:						List[Any]=2							,__lowerCAmelCase:						Optional[Any]=0.02							,__lowerCAmelCase:						int=3							,__lowerCAmelCase:						Tuple=4							,__lowerCAmelCase:						List[str]=None							,):







     '''simple docstring'''







     _lowerCamelCase   :				int							=     parent
     _lowerCamelCase   :				Union[str, Any]							=     batch_size
     _lowerCamelCase   :				Optional[Any]							=     seq_length
     _lowerCamelCase   :				Union[str, Any]							=     is_training
     _lowerCamelCase   :				Any							=     use_input_mask
     _lowerCamelCase   :				Tuple							=     use_token_type_ids
     _lowerCamelCase   :				Dict							=     use_labels
     _lowerCamelCase   :				Any							=     vocab_size
     _lowerCamelCase   :				int							=     hidden_size
     _lowerCamelCase   :				Optional[Any]							=     num_hidden_layers
     _lowerCamelCase   :				Union[str, Any]							=     num_attention_heads
     _lowerCamelCase   :				List[str]							=     intermediate_size
     _lowerCamelCase   :				str							=     hidden_act
     _lowerCamelCase   :				Union[str, Any]							=     hidden_dropout_prob
     _lowerCamelCase   :				Union[str, Any]							=     attention_probs_dropout_prob
     _lowerCamelCase   :				Optional[Any]							=     max_position_embeddings
     _lowerCamelCase   :				Optional[int]							=     type_vocab_size
     _lowerCamelCase   :				Dict							=     type_sequence_label_size
     _lowerCamelCase   :				List[Any]							=     initializer_range
     _lowerCamelCase   :				int							=     num_labels
     _lowerCamelCase   :				List[str]							=     num_choices
     _lowerCamelCase   :				Any							=     scope


  def 							_lowercase					(			self:						List[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				str							=     ids_tensor([self.batch_size, self.seq_length]							,self.vocab_size    )

     _lowerCamelCase   :				int							=     None
     if self.use_input_mask:
        _lowerCamelCase   :				Optional[Any]							=     random_attention_mask([self.batch_size, self.seq_length]    )

     _lowerCamelCase   :				Optional[int]							=     None
     if self.use_token_type_ids:
        _lowerCamelCase   :				Any							=     ids_tensor([self.batch_size, self.seq_length]							,self.type_vocab_size    )

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

     _lowerCamelCase   :				Tuple							=     self.get_config()

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


  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     return NystromformerConfig(
         vocab_size=self.vocab_size							,hidden_size=self.hidden_size							,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							,is_decoder=__lowerCAmelCase							,initializer_range=self.initializer_range							,)


  def 							_lowercase					(			self:						int							,__lowerCAmelCase:						Optional[int]							,__lowerCAmelCase:						str							,__lowerCAmelCase:						List[Any]							,__lowerCAmelCase:						Tuple							,__lowerCAmelCase:						List[Any]							,__lowerCAmelCase:						List[str]							,__lowerCAmelCase:						Dict    ):







     '''simple docstring'''







     _lowerCamelCase   :				Dict							=     NystromformerModel(config=__lowerCAmelCase    )
     model.to(__lowerCAmelCase    )
     model.eval()
     _lowerCamelCase   :				Optional[int]							=     model(__lowerCAmelCase							,attention_mask=__lowerCAmelCase							,token_type_ids=__lowerCAmelCase    )
     _lowerCamelCase   :				Union[str, Any]							=     model(__lowerCAmelCase							,token_type_ids=__lowerCAmelCase    )
     _lowerCamelCase   :				Dict							=     model(__lowerCAmelCase    )
     self.parent.assertEqual(result.last_hidden_state.shape							,(self.batch_size, self.seq_length, self.hidden_size)    )


  def 							_lowercase					(			self:						Dict							,__lowerCAmelCase:						int							,__lowerCAmelCase:						List[str]							,__lowerCAmelCase:						Union[str, Any]							,__lowerCAmelCase:						int							,__lowerCAmelCase:						Dict							,__lowerCAmelCase:						Optional[Any]							,__lowerCAmelCase:						Tuple    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[Any]							=     NystromformerForMaskedLM(config=__lowerCAmelCase    )
     model.to(__lowerCAmelCase    )
     model.eval()
     _lowerCamelCase   :				List[Any]							=     model(__lowerCAmelCase							,attention_mask=__lowerCAmelCase							,token_type_ids=__lowerCAmelCase							,labels=__lowerCAmelCase    )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)    )


  def 							_lowercase					(			self:						Union[str, Any]							,__lowerCAmelCase:						List[str]							,__lowerCAmelCase:						Optional[int]							,__lowerCAmelCase:						Optional[int]							,__lowerCAmelCase:						str							,__lowerCAmelCase:						Tuple							,__lowerCAmelCase:						Optional[int]							,__lowerCAmelCase:						Optional[int]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[Any]							=     NystromformerForQuestionAnswering(config=__lowerCAmelCase    )
     model.to(__lowerCAmelCase    )
     model.eval()
     _lowerCamelCase   :				Any							=     model(
         __lowerCAmelCase							,attention_mask=__lowerCAmelCase							,token_type_ids=__lowerCAmelCase							,start_positions=__lowerCAmelCase							,end_positions=__lowerCAmelCase							,)
     self.parent.assertEqual(result.start_logits.shape							,(self.batch_size, self.seq_length)    )
     self.parent.assertEqual(result.end_logits.shape							,(self.batch_size, self.seq_length)    )


  def 							_lowercase					(			self:						Union[str, Any]							,__lowerCAmelCase:						str							,__lowerCAmelCase:						str							,__lowerCAmelCase:						int							,__lowerCAmelCase:						Optional[Any]							,__lowerCAmelCase:						Dict							,__lowerCAmelCase:						int							,__lowerCAmelCase:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[Any]							=     self.num_labels
     _lowerCamelCase   :				int							=     NystromformerForSequenceClassification(__lowerCAmelCase    )
     model.to(__lowerCAmelCase    )
     model.eval()
     _lowerCamelCase   :				int							=     model(__lowerCAmelCase							,attention_mask=__lowerCAmelCase							,token_type_ids=__lowerCAmelCase							,labels=__lowerCAmelCase    )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)    )


  def 							_lowercase					(			self:						str							,__lowerCAmelCase:						List[str]							,__lowerCAmelCase:						List[Any]							,__lowerCAmelCase:						str							,__lowerCAmelCase:						Dict							,__lowerCAmelCase:						str							,__lowerCAmelCase:						int							,__lowerCAmelCase:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[str]							=     self.num_labels
     _lowerCamelCase   :				int							=     NystromformerForTokenClassification(config=__lowerCAmelCase    )
     model.to(__lowerCAmelCase    )
     model.eval()
     _lowerCamelCase   :				Tuple							=     model(__lowerCAmelCase							,attention_mask=__lowerCAmelCase							,token_type_ids=__lowerCAmelCase							,labels=__lowerCAmelCase    )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.num_labels)    )


  def 							_lowercase					(			self:						str							,__lowerCAmelCase:						List[Any]							,__lowerCAmelCase:						List[Any]							,__lowerCAmelCase:						List[Any]							,__lowerCAmelCase:						Optional[int]							,__lowerCAmelCase:						int							,__lowerCAmelCase:						List[Any]							,__lowerCAmelCase:						str    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[int]							=     self.num_choices
     _lowerCamelCase   :				List[Any]							=     NystromformerForMultipleChoice(config=__lowerCAmelCase    )
     model.to(__lowerCAmelCase    )
     model.eval()
     _lowerCamelCase   :				int							=     input_ids.unsqueeze(1    ).expand(-1							,self.num_choices							,-1    ).contiguous()
     _lowerCamelCase   :				Dict							=     token_type_ids.unsqueeze(1    ).expand(-1							,self.num_choices							,-1    ).contiguous()
     _lowerCamelCase   :				Any							=     input_mask.unsqueeze(1    ).expand(-1							,self.num_choices							,-1    ).contiguous()
     _lowerCamelCase   :				Optional[Any]							=     model(
         __lowerCAmelCase							,attention_mask=__lowerCAmelCase							,token_type_ids=__lowerCAmelCase							,labels=__lowerCAmelCase							,)
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_choices)    )






  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[str]							=     self.prepare_config_and_inputs()
     (
         (
         _lowerCamelCase
     ),					(
         _lowerCamelCase
     ),					(
         _lowerCamelCase
     ),					(
         _lowerCamelCase
     ),					(
         _lowerCamelCase
     ),					(
         _lowerCamelCase
     ),					(
         _lowerCamelCase
     ),					
     )   :				str							=     config_and_inputs
     _lowerCamelCase   :				Optional[int]							=     {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict




@require_torch
class 							A_      (       _a     ,	_a     ,	unittest.TestCase ):
  lowerCAmelCase__      			=	(
      (
          NystromformerModel,
          NystromformerForMaskedLM,
          NystromformerForMultipleChoice,
          NystromformerForQuestionAnswering,
          NystromformerForSequenceClassification,
          NystromformerForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  lowerCAmelCase__      			=	(
      {
          'feature-extraction': NystromformerModel,
          'fill-mask': NystromformerForMaskedLM,
          'question-answering': NystromformerForQuestionAnswering,
          'text-classification': NystromformerForSequenceClassification,
          'token-classification': NystromformerForTokenClassification,
          'zero-shot': NystromformerForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  lowerCAmelCase__      			=	False
  lowerCAmelCase__      			=	False


  def 							_lowercase					(			self:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Union[str, Any]							=     NystromformerModelTester(self    )
     _lowerCamelCase   :				Optional[Any]							=     ConfigTester(self							,config_class=__lowerCAmelCase							,hidden_size=37    )


  def 							_lowercase					(			self:						Optional[int]    ):







     '''simple docstring'''







     self.config_tester.run_common_tests()


  def 							_lowercase					(			self:						List[str]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Dict							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__lowerCAmelCase    )


  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				int							=     self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
        _lowerCamelCase   :				Tuple							=     type
        self.model_tester.create_and_check_model(*__lowerCAmelCase    )


  def 							_lowercase					(			self:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[str]							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase    )


  def 							_lowercase					(			self:						List[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[Any]							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*__lowerCAmelCase    )


  def 							_lowercase					(			self:						int    ):







     '''simple docstring'''







     _lowerCamelCase   :				Union[str, Any]							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase    )


  def 							_lowercase					(			self:						int    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[Any]							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase    )


  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[int]							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase    )






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







     '''simple docstring'''







     for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _lowerCamelCase   :				str							=     NystromformerModel.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )




@require_torch
class 							A_      (       unittest.TestCase ):


  @slow
  def 							_lowercase					(			self:						List[str]    ):







     '''simple docstring'''







     _lowerCamelCase   :				str							=     NystromformerModel.from_pretrained("uw-madison/nystromformer-512"    )
     _lowerCamelCase   :				Optional[Any]							=     torch.tensor([[0, 1, 2, 3, 4, 5]]    )

     with torch.no_grad():
        _lowerCamelCase   :				List[str]							=     model(__lowerCAmelCase    )[0]

     _lowerCamelCase   :				Dict							=     torch.Size((1, 6, 768)    )
     self.assertEqual(output.shape							,__lowerCAmelCase    )

     _lowerCamelCase   :				Optional[Any]							=     torch.tensor(
         [[[-0.45_32, -0.09_36, 0.51_37], [-0.26_76, 0.06_28, 0.61_86], [-0.36_29, -0.17_26, 0.47_16]]]    )

     self.assertTrue(torch.allclose(output[:, :3, :3]							,__lowerCAmelCase							,atol=1e-4    )    )






  @slow
  def 							_lowercase					(			self:						Tuple    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[Any]							=     "the [MASK] of Belgium is Brussels"

     _lowerCamelCase   :				str							=     AutoTokenizer.from_pretrained("uw-madison/nystromformer-512"    )
     _lowerCamelCase   :				List[Any]							=     NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512"    )

     _lowerCamelCase   :				Optional[int]							=     tokenizer(__lowerCAmelCase							,return_tensors="pt"    )

     with torch.no_grad():
        _lowerCamelCase   :				Dict							=     model(encoding.input_ids    ).logits

     _lowerCamelCase   :				Tuple							=     token_logits[:, 2, :].argmax(-1    )[0]

     self.assertEqual(tokenizer.decode(__lowerCAmelCase    )							,"capital"    )




import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel






def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase       ):
 a__												= ('''dense.weight''', '''attention.self.query''', '''attention.self.key''', '''attention.self.value''')

 a__												= (
     ('''layer.''', '''layer_'''),
     ('''word_embeddings.weight''', '''word_embeddings'''),
     ('''position_embeddings.weight''', '''position_embeddings'''),
     ('''token_type_embeddings.weight''', '''token_type_embeddings'''),
     ('''.''', '''/'''),
     ('''LayerNorm/weight''', '''LayerNorm/gamma'''),
     ('''LayerNorm/bias''', '''LayerNorm/beta'''),
     ('''weight''', '''kernel'''),
 )

 if not os.path.isdir(__UpperCAmelCase       ):
  os.makedirs(__UpperCAmelCase       )

 a__												= model.state_dict()

 def to_tf_var_name(__UpperCAmelCase       ):
  for patt, repl in iter(__UpperCAmelCase       ):
   a__												= name.replace(__UpperCAmelCase							,  __UpperCAmelCase       )
  return f"bert/{name}"

 def create_tf_var(__UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase       ):
  a__												= tf.dtypes.as_dtype(tensor.dtype       )
  a__												= tf.get_variable(dtype=__UpperCAmelCase							,  shape=tensor.shape							,  name=__UpperCAmelCase							,  initializer=tf.zeros_initializer()       )
  session.run(tf.variables_initializer([tf_var]       )       )
  session.run(__UpperCAmelCase       )
  return tf_var

 tf.reset_default_graph()
 with tf.Session() as session:
  for var_name in state_dict:
   a__												= to_tf_var_name(__UpperCAmelCase       )
   a__												= state_dict[var_name].numpy()
   if any(x in var_name for x in tensors_to_transpose       ):
    a__												= torch_tensor.T
   a__												= create_tf_var(tensor=__UpperCAmelCase							,  name=__UpperCAmelCase							,  session=__UpperCAmelCase       )
   tf.keras.backend.set_value(__UpperCAmelCase							,  __UpperCAmelCase       )
   a__												= session.run(__UpperCAmelCase       )
   print(f"Successfully created {tf_name}: {np.allclose(__UpperCAmelCase							,  __UpperCAmelCase       )}"       )

  a__												= tf.train.Saver(tf.trainable_variables()       )
  saver.save(__UpperCAmelCase							,  os.path.join(__UpperCAmelCase							,  model_name.replace('''-'''							,  '''_'''       ) + '''.ckpt'''       )       )






def     __a					(		__UpperCAmelCase=None       ):
 a__												= argparse.ArgumentParser()
 parser.add_argument('''--model_name'''							,  type=__UpperCAmelCase							,  required=__UpperCAmelCase							,  help='''model name e.g. bert-base-uncased'''       )
 parser.add_argument(
     '''--cache_dir'''							,  type=__UpperCAmelCase							,  default=__UpperCAmelCase							,  required=__UpperCAmelCase							,  help='''Directory containing pytorch model'''       )
 parser.add_argument('''--pytorch_model_path'''							,  type=__UpperCAmelCase							,  required=__UpperCAmelCase							,  help='''/path/to/<pytorch-model-name>.bin'''       )
 parser.add_argument('''--tf_cache_dir'''							,  type=__UpperCAmelCase							,  required=__UpperCAmelCase							,  help='''Directory in which to save tensorflow model'''       )
 a__												= parser.parse_args(__UpperCAmelCase       )

 a__												= BertModel.from_pretrained(
     pretrained_model_name_or_path=args.model_name							,  state_dict=torch.load(args.pytorch_model_path       )							,  cache_dir=args.cache_dir							,  )

 convert_pytorch_checkpoint_to_tf(model=__UpperCAmelCase							,  ckpt_dir=args.tf_cache_dir							,  model_name=args.model_name       )


if __name__ == "__main__":
     main()





