






"""simple docstring"""




import unittest

from transformers import MraConfig, 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, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
     import torch

     from transformers import (
         MraForMaskedLM,
         MraForMultipleChoice,
         MraForQuestionAnswering,
         MraForSequenceClassification,
         MraForTokenClassification,
         MraModel,
     )
     from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST




class    lowercase_    :






      def __init__(    self     ,     a_     ,     a_=2     ,     a_=8     ,     a_=True     ,     a_=True     ,     a_=True     ,     a_=True     ,     a_=9_9     ,     a_=1_6     ,     a_=5     ,     a_=2     ,     a_=3_6     ,     a_="gelu"     ,     a_=0.0     ,     a_=0.0     ,     a_=5_1_2     ,     a_=1_6     ,     a_=2     ,     a_=0.02     ,     a_=3     ,     a_=4     ,     a_=None     ,     )	->List[str]:







        '''simple docstring'''


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






      def 			lowerCamelCase__    (    self	)	->str:







        '''simple docstring'''


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

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

        _a							=							None
        if self.use_token_type_ids:
          _a							=							ids_tensor([self.batch_size, self.seq_length]     ,     self.type_vocab_size	)

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

        _a							=							self.get_config()

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






      def 			lowerCamelCase__    (    self	)	->List[str]:







        '''simple docstring'''


        return MraConfig(
            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=a_     ,     initializer_range=self.initializer_range     ,     )






      def 			lowerCamelCase__    (    self	)	->List[Any]:







        '''simple docstring'''


        _a							=							self.get_config()
        _a							=							3_0_0
        return config






      def 			lowerCamelCase__    (    self	)	->List[str]:







        '''simple docstring'''


        (
            (
            _a	
        )			,	(
            _a	
        )			,	(
            _a	
        )			,	(
            _a	
        )			,	(
            _a	
        )			,	(
            _a	
        )			,	(
            _a	
        )			,		
        )							=							self.prepare_config_and_inputs()

        _a							=							True
        _a							=							floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
        _a							=							ids_tensor([self.batch_size, self.seq_length]     ,     vocab_size=2	)

        return (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
            encoder_hidden_states,
            encoder_attention_mask,
        )






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_	)	->List[Any]:







        '''simple docstring'''


        _a							=							MraModel(config=a_	)
        model.to(a_	)
        model.eval()
        _a							=							model(a_     ,     attention_mask=a_     ,     token_type_ids=a_	)
        _a							=							model(a_     ,     token_type_ids=a_	)
        _a							=							model(a_	)
        self.parent.assertEqual(result.last_hidden_state.shape     ,     (self.batch_size, self.seq_length, self.hidden_size)	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     )	->Optional[Any]:







        '''simple docstring'''


        _a							=							True
        _a							=							MraModel(a_	)
        model.to(a_	)
        model.eval()
        _a							=							model(
            a_     ,     attention_mask=a_     ,     token_type_ids=a_     ,     encoder_hidden_states=a_     ,     encoder_attention_mask=a_     ,     )
        _a							=							model(
            a_     ,     attention_mask=a_     ,     token_type_ids=a_     ,     encoder_hidden_states=a_     ,     )
        _a							=							model(a_     ,     attention_mask=a_     ,     token_type_ids=a_	)
        self.parent.assertEqual(result.last_hidden_state.shape     ,     (self.batch_size, self.seq_length, self.hidden_size)	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_	)	->Optional[Any]:







        '''simple docstring'''


        _a							=							MraForMaskedLM(config=a_	)
        model.to(a_	)
        model.eval()
        _a							=							model(a_     ,     attention_mask=a_     ,     token_type_ids=a_     ,     labels=a_	)
        self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.seq_length, self.vocab_size)	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_	)	->List[Any]:







        '''simple docstring'''


        _a							=							MraForQuestionAnswering(config=a_	)
        model.to(a_	)
        model.eval()
        _a							=							model(
            a_     ,     attention_mask=a_     ,     token_type_ids=a_     ,     start_positions=a_     ,     end_positions=a_     ,     )
        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 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_	)	->int:







        '''simple docstring'''


        _a							=							self.num_labels
        _a							=							MraForSequenceClassification(a_	)
        model.to(a_	)
        model.eval()
        _a							=							model(a_     ,     attention_mask=a_     ,     token_type_ids=a_     ,     labels=a_	)
        self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.num_labels)	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_	)	->List[str]:







        '''simple docstring'''


        _a							=							self.num_labels
        _a							=							MraForTokenClassification(config=a_	)
        model.to(a_	)
        model.eval()
        _a							=							model(a_     ,     attention_mask=a_     ,     token_type_ids=a_     ,     labels=a_	)
        self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.seq_length, self.num_labels)	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_     ,     a_	)	->List[str]:







        '''simple docstring'''


        _a							=							self.num_choices
        _a							=							MraForMultipleChoice(config=a_	)
        model.to(a_	)
        model.eval()
        _a							=							input_ids.unsqueeze(1	).expand(-1     ,     self.num_choices     ,     -1	).contiguous()
        _a							=							token_type_ids.unsqueeze(1	).expand(-1     ,     self.num_choices     ,     -1	).contiguous()
        _a							=							input_mask.unsqueeze(1	).expand(-1     ,     self.num_choices     ,     -1	).contiguous()
        _a							=							model(
            a_     ,     attention_mask=a_     ,     token_type_ids=a_     ,     labels=a_     ,     )
        self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.num_choices)	)






      def 			lowerCamelCase__    (    self	)	->Any:







        '''simple docstring'''


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







@require_torch
class    lowercase_    (_UpperCAmelCase,    unittest.TestCase     ):
      A__	:		List[Any]       			=				(
          (
              MraModel,
              MraForMaskedLM,
              MraForMultipleChoice,
              MraForQuestionAnswering,
              MraForSequenceClassification,
              MraForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      A__	:		Union[str, Any]       			=				False
      A__	:		Any       			=				False
      A__	:		Optional[int]       			=				False
      A__	:		Union[str, Any]       			=				False

      A__	:		Union[str, Any]       			=				()






      def 			lowerCamelCase__    (    self	)	->Dict:







        '''simple docstring'''


        _a							=							MraModelTester(self	)
        _a							=							ConfigTester(self     ,     config_class=a_     ,     hidden_size=3_7	)






      def 			lowerCamelCase__    (    self	)	->Optional[int]:







        '''simple docstring'''


        self.config_tester.run_common_tests()






      def 			lowerCamelCase__    (    self	)	->int:







        '''simple docstring'''


        _a							=							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*a_	)






      def 			lowerCamelCase__    (    self	)	->List[Any]:







        '''simple docstring'''


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






      def 			lowerCamelCase__    (    self	)	->List[Any]:







        '''simple docstring'''


        _a							=							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*a_	)






      def 			lowerCamelCase__    (    self	)	->Optional[int]:







        '''simple docstring'''


        _a							=							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*a_	)






      def 			lowerCamelCase__    (    self	)	->List[str]:







        '''simple docstring'''


        _a							=							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*a_	)






      def 			lowerCamelCase__    (    self	)	->Any:







        '''simple docstring'''


        _a							=							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*a_	)






      def 			lowerCamelCase__    (    self	)	->Optional[int]:







        '''simple docstring'''


        _a							=							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*a_	)






      @slow
      def 			lowerCamelCase__    (    self	)	->Any:







        '''simple docstring'''


        for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          _a							=							MraModel.from_pretrained(a_	)
          self.assertIsNotNone(a_	)






      @unittest.skip(reason="MRA does not output attentions"	)
      def 			lowerCamelCase__    (    self	)	->Optional[int]:







        '''simple docstring'''


        return







@require_torch
class    lowercase_    (unittest.TestCase     ):






      @slow
      def 			lowerCamelCase__    (    self	)	->Optional[int]:







        '''simple docstring'''


        _a							=							MraModel.from_pretrained("uw-madison/mra-base-512-4"	)
        _a							=							torch.arange(2_5_6	).unsqueeze(0	)

        with torch.no_grad():
          _a							=							model(a_	)[0]

        _a							=							torch.Size((1, 2_5_6, 7_6_8)	)
        self.assertEqual(output.shape     ,     a_	)

        _a							=							torch.tensor(
            [[[-0.0_140, 0.0_830, -0.0_381], [0.1_546, 0.1_402, 0.0_220], [0.1_162, 0.0_851, 0.0_165]]]	)

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






      @slow
      def 			lowerCamelCase__    (    self	)	->str:







        '''simple docstring'''


        _a							=							MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4"	)
        _a							=							torch.arange(2_5_6	).unsqueeze(0	)

        with torch.no_grad():
          _a							=							model(a_	)[0]

        _a							=							5_0_2_6_5

        _a							=							torch.Size((1, 2_5_6, vocab_size)	)
        self.assertEqual(output.shape     ,     a_	)

        _a							=							torch.tensor(
            [[[9.2_595, -3.6_038, 11.8_819], [9.3_869, -3.2_693, 11.0_956], [11.8_524, -3.4_938, 13.1_210]]]	)

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






      @slow
      def 			lowerCamelCase__    (    self	)	->str:







        '''simple docstring'''


        _a							=							MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3"	)
        _a							=							torch.arange(4_0_9_6	).unsqueeze(0	)

        with torch.no_grad():
          _a							=							model(a_	)[0]

        _a							=							5_0_2_6_5

        _a							=							torch.Size((1, 4_0_9_6, vocab_size)	)
        self.assertEqual(output.shape     ,     a_	)

        _a							=							torch.tensor(
            [[[5.4_789, -2.3_564, 7.5_064], [7.9_067, -1.3_369, 9.9_668], [9.0_712, -1.8_106, 7.0_380]]]	)

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










"""simple docstring"""




import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch






def        lowerCAmelCase (      UpperCamelCase_:							Dict					)       ->							Any:







  '''simple docstring'''







  _a							=	os.path.join(args.tf_model_dir     , "parameters.json"					)
  _a							=	json.loads(open(UpperCamelCase_					).read()					)
  if not params:
    raise ValueError(
        f'''It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file.'''					)
  if not args.output.endswith(".pt"					):
    _a							=	args.output + ".pt"
  _a							=	OrderedDict()
  with tf.device("/CPU:0"					):
    _a							=	tf.train.load_checkpoint(args.tf_model_dir					)
    _a							=	reader.get_variable_to_shape_map()
    for key_name in shapes.keys():
      _a							=	reader.get_tensor(UpperCamelCase_					).astype(np.floataa					)
      if key_name.endswith("/adam_m"					) or key_name.endswith("/adam_v"					):
        continue
      if key_name.startswith("pasts/"					):
        if key_name.startswith("pasts/mlp"					):
          _a							=	int(key_name[9]					)
        elif key_name.startswith("pasts/out"					):
          _a							=	8
        _a							=	"model.sqout.%d.weight" % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
        _a							=	vnp.transpose([1, 0]					).copy()  # Mesh-Tensorflow is a diagonal matrix
        _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name.startswith("model/moe"					):
        _a							=	int(key_name[9:].split("/"					)[0]					)
        if key_name.endswith("/switch_gating/kernel"					):
          _a							=	"model.blocks.%d.feed_forward.mlp.router.classifier.weight" % player
          _a							=	vnp.transpose([1, 0]					).copy()  # Mesh-Tensorflow is a diagonal matrix
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/softmlp/kernel"					):
          _a							=	"model.blocks.%d.feed_forward.soft_bypass_mlp.weight" % player
          _a							=	vnp.transpose([1, 0]					).copy()  # Mesh-Tensorflow is a diagonal matrix
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/wo/kernel"					) or key_name.endswith("/wi/kernel"					):
          _a							=	key_name[-9:-7]
          for i in range(16					):
            _a							=	"model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight" % (player, i, nlayer)
            _a							=	(
                vnp[i].transpose([1, 0]					).copy()
            )  # In Mesh-Tensorflow, it is one array, so it is divided
            _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name.startswith("model/mlp"					):
        _a							=	int(key_name[9:].split("/"					)[0]					)
        if key_name.endswith("/p1/kernel"					):
          _a							=	"model.blocks.%d.feed_forward.mlp.wi.weight" % player
          _a							=	vnp.transpose([1, 0]					).copy()  # Mesh-Tensorflow is a diagonal matrix
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/p1/bias"					):
          _a							=	"model.blocks.%d.feed_forward.mlp.wi.bias" % player
          _a							=	vnp.copy()  # same because it is one dimensional
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/p2/kernel"					):
          _a							=	"model.blocks.%d.feed_forward.mlp.wo.weight" % player
          _a							=	vnp.transpose([1, 0]					).copy()  # Mesh-Tensorflow is a diagonal matrix
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/p2/bias"					):
          _a							=	"model.blocks.%d.feed_forward.mlp.wo.bias" % player
          _a							=	vnp.copy()  # same because it is one dimensional
          _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name.startswith("model/ln"					):
        _a							=	int(key_name[8:].split("/"					)[0]					)
        if key_name.endswith("/b"					):
          _a							=	"model.blocks.%d.feed_forward.norm.bias" % player
          _a							=	vnp.copy()  # same because it is one dimensional
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/g"					):
          _a							=	"model.blocks.%d.feed_forward.norm.weight" % player
          _a							=	vnp.copy()  # same because it is one dimensional
          _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name.startswith("model/att"					):
        _a							=	int(key_name[9:].split("/"					)[0]					)
        if key_name.endswith("/qkv/kernel"					):
          _a							=	vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
          _a							=	state[:, 0, :, :]
          _a							=	state[:, 1, :, :]
          _a							=	state[:, 2, :, :]
          _a							=	(
              state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]]					)
              .transpose([1, 0]					)
              .copy()
          )  # Mesh-Tensorflow is a diagonal matrix
          _a							=	(
              state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]]					)
              .transpose([1, 0]					)
              .copy()
          )  # Mesh-Tensorflow is a diagonal matrix
          _a							=	(
              state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]]					)
              .transpose([1, 0]					)
              .copy()
          )  # Mesh-Tensorflow is a diagonal matrix
          _a							=	"model.blocks.%d.self_attn.self_attn.q_proj.weight" % player
          _a							=	torch.tensor(UpperCamelCase_					)
          _a							=	"model.blocks.%d.self_attn.self_attn.k_proj.weight" % player
          _a							=	torch.tensor(UpperCamelCase_					)
          _a							=	"model.blocks.%d.self_attn.self_attn.v_proj.weight" % player
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/o/kernel"					):
          _a							=	"model.blocks.%d.self_attn.self_attn.out_proj.weight" % player
          _a							=	(
              vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]					).transpose([1, 0]					).copy()
          )  # Mesh-Tensorflow is a diagonal matrix
          _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name.startswith("model/an"					):
        _a							=	int(key_name[8:].split("/"					)[0]					)
        if key_name.endswith("/b"					):
          _a							=	"model.blocks.%d.self_attn.norm.bias" % player
          _a							=	vnp.copy()  # same because it is one dimensional
          _a							=	torch.tensor(UpperCamelCase_					)
        elif key_name.endswith("/g"					):
          _a							=	"model.blocks.%d.self_attn.norm.weight" % player
          _a							=	vnp.copy()  # same because it is one dimensional
          _a							=	torch.tensor(UpperCamelCase_					)
      elif (
          key_name.startswith("model/wte"					)
          or key_name.startswith("model/wpe"					)
          or key_name.startswith("model/ete"					)
      ):
        _a							=	{"wte": "embed_tokens", "wpe": "position_embeddings", "ete": "extra_position_embeddings"}[
            key_name[-3:]
        ]
        _a							=	"model.%s.weight" % nlayer
        _a							=	vnp.copy()  # same in embedded
        _a							=	torch.tensor(UpperCamelCase_					)
        if key_name.startswith("model/wte"					):
          _a							=	"lm_head.weight"
          _a							=	vnp.copy()  # same in embedded
          _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name.startswith("model/wob"					):
        _a							=	"final_logits_bias"
        _a							=	vnp.copy()  # same in embedded
        _a							=	state.reshape((1, -1)					)
        _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name == "model/dense/kernel":
        _a							=	"model.last_project.weight"
        _a							=	vnp.transpose([1, 0]					).copy()  # Mesh-Tensorflow is a diagonal matrix
        _a							=	torch.tensor(UpperCamelCase_					)
      elif key_name == "model/dense_1/bias":
        _a							=	"model.last_project.bias"
        _a							=	vnp.copy()  # same because it is one dimensional
        _a							=	torch.tensor(UpperCamelCase_					)
  torch.save(UpperCamelCase_     , args.output					)


if __name__ == "__main__":
     UpperCamelCase							       =    argparse.ArgumentParser(
         description="""model converter.""", formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
     parser.add_argument("""--tf_model_dir""", metavar="""PATH""", type=str, required=True, help="""import model""")
     parser.add_argument("""--output""", metavar="""PATH""", type=str, required=True, help="""output model""")
     UpperCamelCase							       =    parser.parse_args()
     convert_tf_gptsan_to_pt(args)


