



import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__ ):


      """simple docstring"""



      return params[F"""{prefix}/{prefix}/relpos_bias/rel_embedding"""][:, i, :]

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__="attention" ):


      """simple docstring"""



      lowerCamelCase     :   Dict								=		np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/key/kernel"""][:, i, :, :] )
      lowerCamelCase     :   List[Any]								=		k_tmp.reshape(k_tmp.shape[0]   ,   k_tmp.shape[1] * k_tmp.shape[2] )
      lowerCamelCase     :   Any								=		np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/out/kernel"""][:, i, :, :] )
      lowerCamelCase     :   Dict								=		o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]   ,   o_tmp.shape[2] )
      lowerCamelCase     :   Optional[int]								=		np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/query/kernel"""][:, i, :, :] )
      lowerCamelCase     :   str								=		q_tmp.reshape(q_tmp.shape[0]   ,   q_tmp.shape[1] * q_tmp.shape[2] )
      lowerCamelCase     :   Dict								=		np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/value/kernel"""][:, i, :, :] )
      lowerCamelCase     :   Optional[int]								=		v_tmp.reshape(v_tmp.shape[0]   ,   v_tmp.shape[1] * v_tmp.shape[2] )
      return k, o, q, v

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__=False ):


      """simple docstring"""



      if split_mlp_wi:
            lowerCamelCase     :   List[Any]								=		params[F"""{prefix}/{prefix}/mlp/wi_0/kernel"""][:, i, :]
            lowerCamelCase     :   Union[str, Any]								=		params[F"""{prefix}/{prefix}/mlp/wi_1/kernel"""][:, i, :]
            lowerCamelCase     :   int								=		(wi_a, wi_a)
      else:
            lowerCamelCase     :   List[str]								=		params[F"""{prefix}/{prefix}/mlp/wi/kernel"""][:, i, :]

      lowerCamelCase     :   Tuple								=		params[F"""{prefix}/{prefix}/mlp/wo/kernel"""][:, i, :]
      return wi, wo

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__ ):


      """simple docstring"""



      return params[F"""{prefix}/{prefix}/{layer_name}/scale"""][:, i]

def     __lowercase(	UpperCAmelCase__   ,   *, UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__ = False ):


      """simple docstring"""



      lowerCamelCase     :   Any								=		traverse_util.flatten_dict(variables["target"] )
      lowerCamelCase     :   Optional[int]								=		{"/".join(lowerCamelCase__ ): v for k, v in old.items()}

      # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
      lowerCamelCase     :   Tuple								=		"encoder/encoder/mlp/wi_0/kernel" in old
      print("Split MLP:"   ,   lowerCamelCase__ )

      lowerCamelCase     :   List[Any]								=		collections.OrderedDict()

      # Shared embeddings.
      lowerCamelCase     :   Optional[Any]								=		old["token_embedder/embedding"]

      # Encoder.
      for i in range(lowerCamelCase__ ):
            # Block i, layer 0 (Self Attention).
            lowerCamelCase     :   Dict								=		tax_layer_norm_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "encoder"   ,   "pre_attention_layer_norm" )
            lowerCamelCase     :   Dict								=		tax_attention_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "encoder"   ,   "attention" )
            lowerCamelCase     :   List[Any]								=		layer_norm
            lowerCamelCase     :   Optional[int]								=		k.T
            lowerCamelCase     :   Union[str, Any]								=		o.T
            lowerCamelCase     :   Any								=		q.T
            lowerCamelCase     :   Dict								=		v.T

            # Block i, layer 1 (MLP).
            lowerCamelCase     :   Optional[int]								=		tax_layer_norm_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "encoder"   ,   "pre_mlp_layer_norm" )
            lowerCamelCase     :   Optional[int]								=		tax_mlp_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "encoder"   ,   lowerCamelCase__ )
            lowerCamelCase     :   Any								=		layer_norm
            if split_mlp_wi:
                  lowerCamelCase     :   Tuple								=		wi[0].T
                  lowerCamelCase     :   Optional[int]								=		wi[1].T
            else:
                  lowerCamelCase     :   Optional[int]								=		wi.T
            lowerCamelCase     :   int								=		wo.T
            if scalable_attention:
                  # convert the rel_embedding of each layer
                  lowerCamelCase     :   Dict								=		tax_relpos_bias_lookup(
                      lowerCamelCase__   ,   lowerCamelCase__   ,   "encoder" ).T

      lowerCamelCase     :   Tuple								=		old["encoder/encoder_norm/scale"]

      if not scalable_attention:
            lowerCamelCase     :   Tuple								=		tax_relpos_bias_lookup(
                lowerCamelCase__   ,   0   ,   "encoder" ).T
            lowerCamelCase     :   Optional[Any]								=		tax_relpos_bias_lookup(
                lowerCamelCase__   ,   0   ,   "decoder" ).T

      if not is_encoder_only:
            # Decoder.
            for i in range(lowerCamelCase__ ):
                  # Block i, layer 0 (Self Attention).
                  lowerCamelCase     :   int								=		tax_layer_norm_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "decoder"   ,   "pre_self_attention_layer_norm" )
                  lowerCamelCase     :   Optional[Any]								=		tax_attention_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "decoder"   ,   "self_attention" )
                  lowerCamelCase     :   Dict								=		layer_norm
                  lowerCamelCase     :   List[str]								=		k.T
                  lowerCamelCase     :   Dict								=		o.T
                  lowerCamelCase     :   Optional[int]								=		q.T
                  lowerCamelCase     :   Dict								=		v.T

                  # Block i, layer 1 (Cross Attention).
                  lowerCamelCase     :   Union[str, Any]								=		tax_layer_norm_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "decoder"   ,   "pre_cross_attention_layer_norm" )
                  lowerCamelCase     :   List[Any]								=		tax_attention_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "decoder"   ,   "encoder_decoder_attention" )
                  lowerCamelCase     :   Tuple								=		layer_norm
                  lowerCamelCase     :   Union[str, Any]								=		k.T
                  lowerCamelCase     :   Optional[Any]								=		o.T
                  lowerCamelCase     :   Optional[Any]								=		q.T
                  lowerCamelCase     :   str								=		v.T

                  # Block i, layer 2 (MLP).
                  lowerCamelCase     :   List[Any]								=		tax_layer_norm_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "decoder"   ,   "pre_mlp_layer_norm" )
                  lowerCamelCase     :   List[Any]								=		tax_mlp_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "decoder"   ,   lowerCamelCase__ )
                  lowerCamelCase     :   str								=		layer_norm
                  if split_mlp_wi:
                        lowerCamelCase     :   Any								=		wi[0].T
                        lowerCamelCase     :   List[Any]								=		wi[1].T
                  else:
                        lowerCamelCase     :   Optional[Any]								=		wi.T
                  lowerCamelCase     :   int								=		wo.T

                  if scalable_attention:
                        # convert the rel_embedding of each layer
                        lowerCamelCase     :   int								=		tax_relpos_bias_lookup(lowerCamelCase__   ,   lowerCamelCase__   ,   "decoder" ).T

            lowerCamelCase     :   Any								=		old["decoder/decoder_norm/scale"]

            # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
            if "decoder/logits_dense/kernel" in old:
                  lowerCamelCase     :   int								=		old["decoder/logits_dense/kernel"].T

      return new

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__ ):


      """simple docstring"""



      lowerCamelCase     :   Optional[int]								=		collections.OrderedDict([(k, torch.from_numpy(v.copy() )) for (k, v) in converted_params.items()] )

      # Add what is missing.
      if "encoder.embed_tokens.weight" not in state_dict:
            lowerCamelCase     :   Tuple								=		state_dict["shared.weight"]

      if not is_encoder_only:
            if "decoder.embed_tokens.weight" not in state_dict:
                  lowerCamelCase     :   List[Any]								=		state_dict["shared.weight"]

            if "lm_head.weight" not in state_dict:  # For old 1.0 models.
                  print("Using shared word embeddings as lm_head." )
                  lowerCamelCase     :   List[str]								=		state_dict["shared.weight"]

      return state_dict

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__ ):


      """simple docstring"""



      lowerCamelCase     :   Any								=		checkpoints.load_tax_checkpoint(lowerCamelCase__ )
      lowerCamelCase     :   int								=		convert_tax_to_pytorch(
          lowerCamelCase__   ,   num_layers=config.num_layers   ,   is_encoder_only=lowerCamelCase__   ,   scalable_attention=lowerCamelCase__ )
      lowerCamelCase     :   Tuple								=		make_state_dict(lowerCamelCase__   ,   lowerCamelCase__ )
      model.load_state_dict(lowerCamelCase__   ,   strict=lowerCamelCase__ )

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__ = False   ,   UpperCAmelCase__ = False   ,   ):


      """simple docstring"""



      lowerCamelCase     :   Dict								=		MTaConfig.from_json_file(lowerCamelCase__ )
      print(F"""Building PyTorch model from configuration: {config}""" )
      # Non-v1.1 checkpoints could also use T5Model, but this works for all.
      # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
      if is_encoder_only:
            lowerCamelCase     :   Optional[int]								=		UMTaEncoderModel(lowerCamelCase__ )
      else:
            lowerCamelCase     :   Optional[int]								=		UMTaForConditionalGeneration(lowerCamelCase__ )

      # Load weights from tf checkpoint
      load_tax_weights_in_ta(lowerCamelCase__   ,   lowerCamelCase__   ,   lowerCamelCase__   ,   lowerCamelCase__   ,   lowerCamelCase__ )

      # Save pytorch-model
      print(F"""Save PyTorch model to {pytorch_dump_path}""" )
      model.save_pretrained(lowerCamelCase__ )

      # Verify that we can load the checkpoint.
      model.from_pretrained(lowerCamelCase__ )
      print("Done" )


if __name__ == "__main__":
     a_						:     List[str]	     =		argparse.ArgumentParser(description='Converts a native T5X checkpoint into a PyTorch checkpoint.')
     # Required parameters
     parser.add_argument(
         '--t5x_checkpoint_path', default=None, type=str, required=True, help='Path to the T5X checkpoint.'
     )
     parser.add_argument(
         '--config_file',
         default=None,
         type=str,
         required=True,
         help='The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.',
     )
     parser.add_argument(
         '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
     )
     parser.add_argument(
         '--is_encoder_only', action='store_true', help='Check if the model is encoder-decoder model', default=False
     )
     parser.add_argument(
         '--scalable_attention',
         action='store_true',
         help='Whether the model uses scaled attention (umt5 model)',
         default=False,
     )
     a_						:     Optional[Any]	     =		parser.parse_args()
     convert_tax_checkpoint_to_pytorch(
         args.tax_checkpoint_path,
         args.config_file,
         args.pytorch_dump_path,
         args.is_encoder_only,
         args.scalable_attention,
     )




from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import numpy
     import tensorflow as tf

     from transformers import (
         TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
         BertConfig,
         DPRConfig,
         TFDPRContextEncoder,
         TFDPRQuestionEncoder,
         TFDPRReader,
     )





class    lowerCamelCase__  :
     """simple docstring"""



     def __init__(self				,			__a				,			__a=13				,			__a=7				,			__a=True				,			__a=True				,			__a=True				,			__a=True				,			__a=99				,			__a=32				,			__a=2				,			__a=4				,			__a=37				,			__a="gelu"				,			__a=0.1				,			__a=0.1				,			__a=5_12				,			__a=16				,			__a=2				,			__a=0.02				,			__a=3				,			__a=4				,			__a=None				,			__a=0				,			):






           '''simple docstring'''



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







     def 						_a      (self    ):






           '''simple docstring'''



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

           lowerCamelCase								=		None
           if self.use_input_mask:
                 # follow test_modeling_tf_ctrl.py
                 lowerCamelCase								=		random_attention_mask([self.batch_size, self.seq_length]    )

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

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

           lowerCamelCase								=		BertConfig(
               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				,			)
           lowerCamelCase								=		DPRConfig(projection_dim=self.projection_dim				,			**config.to_dict()    )

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







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRContextEncoder(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(result.pooler_output.shape				,			(self.batch_size, self.projection_dim or self.hidden_size)    )







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRQuestionEncoder(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(result.pooler_output.shape				,			(self.batch_size, self.projection_dim or self.hidden_size)    )







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRReader(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__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)    )
           self.parent.assertEqual(result.relevance_logits.shape				,			(self.batch_size,)    )







     def 						_a      (self    ):






           '''simple docstring'''



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





@require_tf
class    lowerCamelCase__  (						UpperCAmelCase_       ,    UpperCAmelCase_       ,    unittest.TestCase):
     """simple docstring"""



     _A       =			(
         (
             TFDPRContextEncoder,
             TFDPRQuestionEncoder,
             TFDPRReader,
         )
         if is_tf_available()
         else ()
     )
     _A       =			{'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

     _A       =			False
     _A       =			False
     _A       =			False
     _A       =			False
     _A       =			False







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRModelTester(self    )
           lowerCamelCase								=		ConfigTester(self				,			config_class=__a				,			hidden_size=37    )







     def 						_a      (self    ):






           '''simple docstring'''



           self.config_tester.run_common_tests()







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_context_encoder(*__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_question_encoder(*__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_reader(*__a    )







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRContextEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRContextEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRQuestionEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRReader.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )





@require_tf
class    lowerCamelCase__  (						unittest.TestCase):
     """simple docstring"""



     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"    )

           lowerCamelCase								=		tf.constant(
               [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]]    )  # [CLS] hello, is my dog cute? [SEP]
           lowerCamelCase								=		model(__a    )[0]  # embedding shape = (1, 768)
           # compare the actual values for a slice.
           lowerCamelCase								=		tf.constant(
               [
                   [
                       0.0323_6253,
                       0.1275_3335,
                       0.1681_8509,
                       0.0027_9786,
                       0.389_6933,
                       0.2426_4945,
                       0.217_8971,
                       -0.0233_5227,
                       -0.0848_1959,
                       -0.1432_4117,
                   ]
               ]    )
           self.assertTrue(numpy.allclose(output[:, :10].numpy()				,			expected_slice.numpy()				,			atol=1E-4    )    )