






"""simple docstring"""


from __future__ import annotations

import unittest

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

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


if is_tf_available():
		import numpy
		import tensorflow as tf

		from transformers.models.esm.modeling_tf_esm import (
		    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFEsmForMaskedLM,
		    TFEsmForSequenceClassification,
		    TFEsmForTokenClassification,
		    TFEsmModel,
		)




class 						__UpperCamelCase     :
					def __init__(							self							:	Union[str, Any]     , UpperCAmelCase							:	Union[str, Any]     , )	->       Dict:
											lowerCAmelCase   :Dict						=			parent
											lowerCAmelCase   :Dict						=			13
											lowerCAmelCase   :Any						=			7
											lowerCAmelCase   :int						=			True
											lowerCAmelCase   :Union[str, Any]						=			True
											lowerCAmelCase   :Dict						=			True
											lowerCAmelCase   :Union[str, Any]						=			99
											lowerCAmelCase   :Optional[int]						=			32
											lowerCAmelCase   :List[Any]						=			2
											lowerCAmelCase   :int						=			4
											lowerCAmelCase   :int						=			37
											lowerCAmelCase   :Union[str, Any]						=			'gelu'
											lowerCAmelCase   :Union[str, Any]						=			0.1
											lowerCAmelCase   :str						=			0.1
											lowerCAmelCase   :int						=			512
											lowerCAmelCase   :int						=			16
											lowerCAmelCase   :Any						=			2
											lowerCAmelCase   :Any						=			0.0_2
											lowerCAmelCase   :Dict						=			3
											lowerCAmelCase   :Any						=			4
											lowerCAmelCase   :Optional[int]						=			None
					def        UpperCAmelCase__ (							self							:	Any   )	->       Any:
											lowerCAmelCase   :Union[str, Any]						=			ids_tensor([self.batch_size, self.seq_length]     , self.vocab_size   )

											lowerCAmelCase   :int						=			None
											if self.use_input_mask:
																	lowerCAmelCase   :Union[str, Any]						=			random_attention_mask([self.batch_size, self.seq_length]   )

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

											lowerCAmelCase   :Optional[int]						=			EsmConfig(
											    vocab_size=self.vocab_size     , hidden_size=self.hidden_size     , num_hidden_layers=self.num_hidden_layers     , pad_token_id=1     , 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     , )

											return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
					def        UpperCAmelCase__ (							self							:	str   )	->       Optional[Any]:
											(
											    (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	,       
											)   :str						=			self.prepare_config_and_inputs()

											lowerCAmelCase   :Dict						=			True
											lowerCAmelCase   :Any						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
											lowerCAmelCase   :Tuple						=			ids_tensor([self.batch_size, self.seq_length]     , vocab_size=2   )

											return (
											    config,
											    input_ids,
											    input_mask,
											    sequence_labels,
											    token_labels,
											    choice_labels,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)
					def        UpperCAmelCase__ (							self							:	List[Any]     , UpperCAmelCase							:	Union[str, Any]     , UpperCAmelCase							:	Optional[int]     , UpperCAmelCase							:	Optional[int]     , UpperCAmelCase							:	int     , UpperCAmelCase							:	List[Any]     , UpperCAmelCase							:	List[Any]   )	->       Dict:
											lowerCAmelCase   :Optional[int]						=			TFEsmModel(config=UpperCamelCase__   )
											lowerCAmelCase   :Dict						=			{'input_ids': input_ids, 'attention_mask': input_mask}
											lowerCAmelCase   :Dict						=			model(UpperCamelCase__   )

											lowerCAmelCase   :Optional[int]						=			[input_ids, input_mask]
											lowerCAmelCase   :str						=			model(UpperCamelCase__   )

											lowerCAmelCase   :Union[str, Any]						=			model(UpperCamelCase__   )

											self.parent.assertEqual(result.last_hidden_state.shape     , (self.batch_size, self.seq_length, self.hidden_size)   )
					def        UpperCAmelCase__ (							self							:	List[Any]     , UpperCAmelCase							:	Optional[int]     , UpperCAmelCase							:	Any     , UpperCAmelCase							:	List[Any]     , UpperCAmelCase							:	Dict     , UpperCAmelCase							:	List[str]     , UpperCAmelCase							:	int     , UpperCAmelCase							:	int     , UpperCAmelCase							:	Union[str, Any]     , )	->       Optional[Any]:
											lowerCAmelCase   :Union[str, Any]						=			True

											lowerCAmelCase   :List[str]						=			TFEsmModel(config=UpperCamelCase__   )
											lowerCAmelCase   :Dict						=			{
											    'input_ids': input_ids,
											    'attention_mask': input_mask,
											    'encoder_hidden_states': encoder_hidden_states,
											    'encoder_attention_mask': encoder_attention_mask,
											}
											lowerCAmelCase   :List[str]						=			model(UpperCamelCase__   )

											lowerCAmelCase   :Optional[Any]						=			[input_ids, input_mask]
											lowerCAmelCase   :int						=			model(UpperCamelCase__     , encoder_hidden_states=UpperCamelCase__   )

											# Also check the case where encoder outputs are not passed
											lowerCAmelCase   :Any						=			model(UpperCamelCase__     , attention_mask=UpperCamelCase__   )

											self.parent.assertEqual(result.last_hidden_state.shape     , (self.batch_size, self.seq_length, self.hidden_size)   )
					def        UpperCAmelCase__ (							self							:	Union[str, Any]     , UpperCAmelCase							:	List[str]     , UpperCAmelCase							:	Optional[Any]     , UpperCAmelCase							:	str     , UpperCAmelCase							:	Tuple     , UpperCAmelCase							:	List[str]     , UpperCAmelCase							:	int   )	->       str:
											lowerCAmelCase   :int						=			TFEsmForMaskedLM(config=UpperCamelCase__   )
											lowerCAmelCase   :Tuple						=			model([input_ids, input_mask]   )
											self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.seq_length, self.vocab_size)   )
					def        UpperCAmelCase__ (							self							:	Any     , UpperCAmelCase							:	Union[str, Any]     , UpperCAmelCase							:	Tuple     , UpperCAmelCase							:	Tuple     , UpperCAmelCase							:	int     , UpperCAmelCase							:	Optional[Any]     , UpperCAmelCase							:	int   )	->       List[Any]:
											lowerCAmelCase   :Union[str, Any]						=			self.num_labels
											lowerCAmelCase   :Optional[Any]						=			TFEsmForTokenClassification(config=UpperCamelCase__   )
											lowerCAmelCase   :Union[str, Any]						=			{'input_ids': input_ids, 'attention_mask': input_mask}
											lowerCAmelCase   :Dict						=			model(UpperCamelCase__   )
											self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.seq_length, self.num_labels)   )
					def        UpperCAmelCase__ (							self							:	str   )	->       Any:
											lowerCAmelCase   :Dict						=			self.prepare_config_and_inputs()
											(
											    (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	, (
											    lowerCAmelCase      
											)	,       
											)   :Any						=			config_and_inputs
											lowerCAmelCase   :Optional[Any]						=			{'input_ids': input_ids, 'attention_mask': input_mask}
											return config, inputs_dict




@require_tf
class 						__UpperCamelCase     (     _snake_case   ,     _snake_case   ,     unittest.TestCase						):
					lowercase_    :				Any      =  (
					    (
					        TFEsmModel,
					        TFEsmForMaskedLM,
					        TFEsmForSequenceClassification,
					        TFEsmForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					lowercase_    :				Optional[Any]      =  (
					    {
					        """feature-extraction""": TFEsmModel,
					        """fill-mask""": TFEsmForMaskedLM,
					        """text-classification""": TFEsmForSequenceClassification,
					        """token-classification""": TFEsmForTokenClassification,
					        """zero-shot""": TFEsmForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					lowercase_    :				Optional[int]      =  False
					lowercase_    :				Dict      =  False
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       Any:
											lowerCAmelCase   :Any						=			TFEsmModelTester(self   )
											lowerCAmelCase   :Dict						=			ConfigTester(self     , config_class=UpperCamelCase__     , hidden_size=37   )
					def        UpperCAmelCase__ (							self							:	Any   )	->       Union[str, Any]:
											self.config_tester.run_common_tests()
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       Any:
											lowerCAmelCase   :Any						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*UpperCamelCase__   )
					def        UpperCAmelCase__ (							self							:	int   )	->       int:
											lowerCAmelCase   :Any						=			self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase__   )
					def        UpperCAmelCase__ (							self							:	Optional[int]   )	->       List[str]:
											lowerCAmelCase   :Union[str, Any]						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__   )
					def        UpperCAmelCase__ (							self							:	int   )	->       Optional[Any]:
											lowerCAmelCase   :List[str]						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__   )
					@slow
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       Union[str, Any]:
											for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowerCAmelCase   :Optional[int]						=			TFEsmModel.from_pretrained(UpperCamelCase__   )
																	self.assertIsNotNone(UpperCamelCase__   )
					@unittest.skip('Protein models do not support embedding resizing.'   )
					def        UpperCAmelCase__ (							self							:	Union[str, Any]   )	->       Dict:
											pass
					@unittest.skip('Protein models do not support embedding resizing.'   )
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       List[str]:
											pass
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       Optional[Any]:
											lowerCAmelCase	, lowerCAmelCase   :Dict						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowerCAmelCase   :List[str]						=			model_class(UpperCamelCase__   )
																	assert isinstance(model.get_input_embeddings()     , tf.keras.layers.Layer   )
																	if model_class is TFEsmForMaskedLM:
																							# Output embedding test differs from the main test because they're a matrix, not a layer
																							lowerCAmelCase   :Dict						=			model.get_bias()
																							assert isinstance(UpperCamelCase__     , UpperCamelCase__   )
																							for k, v in name.items():
																													assert isinstance(UpperCamelCase__     , tf.Variable   )
																	else:
																							lowerCAmelCase   :Optional[int]						=			model.get_output_embeddings()
																							assert x is None
																							lowerCAmelCase   :Any						=			model.get_bias()
																							assert name is None




@require_tf
class 						__UpperCamelCase     (     unittest.TestCase						):
					@slow
					def        UpperCAmelCase__ (							self							:	List[str]   )	->       Dict:
											lowerCAmelCase   :Union[str, Any]						=			TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'   )

											lowerCAmelCase   :Union[str, Any]						=			tf.constant([[0, 1, 2, 3, 4, 5]]   )
											lowerCAmelCase   :List[str]						=			model(UpperCamelCase__   )[0]
											lowerCAmelCase   :Union[str, Any]						=			[1, 6, 33]
											self.assertEqual(list(output.numpy().shape   )     , UpperCamelCase__   )
											# compare the actual values for a slice.
											lowerCAmelCase   :Union[str, Any]						=			tf.constant(
											    [
											        [
											            [8.9_2_1_5_1_8, -1_0.5_8_9_8_1_4, -6.4_6_7_1_3_0_7],
											            [-6.3_9_6_7_1_5_6, -1_3.9_1_1_3_7_7, -1.1_2_1_1_9_1_5],
											            [-7.7_8_1_2_4_7, -1_3.9_5_1_5_5_7, -3.7_4_0_5_9_2],
											        ]
											    ]   )
											self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()     , expected_slice.numpy()     , atol=1e-2   )   )
					@slow
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       List[str]:
											lowerCAmelCase   :Any						=			TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'   )

											lowerCAmelCase   :str						=			tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
											lowerCAmelCase   :List[str]						=			model(UpperCamelCase__   )[0]
											# compare the actual values for a slice.
											lowerCAmelCase   :str						=			tf.constant(
											    [
											        [
											            [0.1_4_4_4_3_0_9_2, 0.5_4_1_2_5_3_2_7, 0.3_2_4_7_7_3_9],
											            [0.3_0_3_4_0_4_8_4, 0.0_0_5_2_6_6_7_6, 0.3_1_0_7_7_7_2_2],
											            [0.3_2_2_7_8_0_4_3, -0.2_4_9_8_7_0_9_6, 0.3_4_1_4_6_2_8],
											        ]
											    ]   )
											self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()     , expected_slice.numpy()     , atol=1e-4   )   )


'''simple docstring'''







import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__="attention"					)      ->      Dict:
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/key/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/out/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/query/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/value/kernel''']
       return k, o, q, v
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__=False					)      ->      Optional[Any]:
       if split_mlp_wi:
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi_0/kernel''']
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi_1/kernel''']
              A_						       =   (wi_a, wi_a)
       else:
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi/kernel''']

       A_						       =   params[F'''{prefix}/layers_{i}/mlp/wo/kernel''']
       return wi, wo
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)      ->      Any:
       return params[F'''{prefix}/layers_{i}/{layer_name}/scale''']
def    UpperCAmelCase__     (   UpperCAmelCase__,							*, UpperCAmelCase__,							UpperCAmelCase__					)      ->      Union[str, Any]:
       A_						       =   traverse_util.flatten_dict(variables["""target"""]					)
       A_						       =   {"""/""".join(UpperCAmelCase__					): v for k, v in old.items()}

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

       A_						       =   collections.OrderedDict()

       # Shared embeddings.
       A_						       =   old["""token_embedder/embedding"""]

       # Encoder.
       for i in range(UpperCAmelCase__					):
              # Block i, layer 0 (Self Attention).
              A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""pre_attention_layer_norm"""					)
              A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""attention"""					)
              A_						       =   layer_norm
              A_						       =   k.T
              A_						       =   o.T
              A_						       =   q.T
              A_						       =   v.T

              # Block i, layer 1 (MLP).
              A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""pre_mlp_layer_norm"""					)
              A_       ,      A_						       =   tax_mlp_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							UpperCAmelCase__					)
              A_						       =   layer_norm
              if split_mlp_wi:
                     A_						       =   wi[0].T
                     A_						       =   wi[1].T
              else:
                     A_						       =   wi.T
              A_						       =   wo.T

       A_						       =   old[
           """encoder/relpos_bias/rel_embedding"""
       ].T
       A_						       =   old["""encoder/encoder_norm/scale"""]

       if not is_encoder_only:
              # Decoder.
              for i in range(UpperCAmelCase__					):
                     # Block i, layer 0 (Self Attention).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_self_attention_layer_norm"""					)
                     A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""self_attention"""					)
                     A_						       =   layer_norm
                     A_						       =   k.T
                     A_						       =   o.T
                     A_						       =   q.T
                     A_						       =   v.T

                     # Block i, layer 1 (Cross Attention).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_cross_attention_layer_norm"""					)
                     A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""encoder_decoder_attention"""					)
                     A_						       =   layer_norm
                     A_						       =   k.T
                     A_						       =   o.T
                     A_						       =   q.T
                     A_						       =   v.T

                     # Block i, layer 2 (MLP).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_mlp_layer_norm"""					)
                     A_       ,      A_						       =   tax_mlp_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							UpperCAmelCase__					)
                     A_						       =   layer_norm
                     if split_mlp_wi:
                            A_						       =   wi[0].T
                            A_						       =   wi[1].T
                     else:
                            A_						       =   wi.T
                     A_						       =   wo.T

              A_						       =   old["""decoder/decoder_norm/scale"""]
              A_						       =   old[
                  """decoder/relpos_bias/rel_embedding"""
              ].T

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

       return new
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__					)      ->      List[Any]:
       A_						       =   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:
              A_						       =   state_dict["""shared.weight"""]

       if not is_encoder_only:
              if "decoder.embed_tokens.weight" not in state_dict:
                     A_						       =   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."""					)
                     A_						       =   state_dict["""shared.weight"""]

       return state_dict
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)      ->      Union[str, Any]:
       A_						       =   checkpoints.load_tax_checkpoint(UpperCAmelCase__					)
       A_						       =   convert_tax_to_pytorch(UpperCAmelCase__,							num_layers=config.num_layers,							is_encoder_only=UpperCAmelCase__					)
       A_						       =   make_state_dict(UpperCAmelCase__,							UpperCAmelCase__					)
       model.load_state_dict(UpperCAmelCase__,							strict=UpperCAmelCase__					)
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__ = False					)      ->      List[Any]:
       A_						       =   TaConfig.from_json_file(UpperCAmelCase__					)
       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:
              A_						       =   TaEncoderModel(UpperCAmelCase__					)
       else:
              A_						       =   TaForConditionalGeneration(UpperCAmelCase__					)

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

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

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


if __name__ == "__main__":
      __lowerCamelCase									=     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
      )
      __lowerCamelCase									=     parser.parse_args()
      convert_tax_checkpoint_to_pytorch(
          args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
      )
