import argparse
import re
from pathlib import Path

import requests
import torch
from PIL import Image
from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor

from transformers import (
    EfficientFormerConfig,
    EfficientFormerForImageClassificationWithTeacher,
    EfficientFormerImageProcessor,
)
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :str    						=	old_name

    if "patch_embed" in old_name:
        __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase  :Optional[int]    						=	old_name.split('''.'''			)

        if layer == "0":
            __UpperCamelCase  :str    						=	old_name.replace('''0'''						,    '''convolution1'''			)
        elif layer == "1":
            __UpperCamelCase  :Optional[int]    						=	old_name.replace('''1'''						,    '''batchnorm_before'''			)
        elif layer == "3":
            __UpperCamelCase  :Union[str, Any]    						=	old_name.replace('''3'''						,    '''convolution2'''			)
        else:
            __UpperCamelCase  :Union[str, Any]    						=	old_name.replace('''4'''						,    '''batchnorm_after'''			)

    if "network" in old_name and re.search(R'''\d\.\d'''						,    SCREAMING_SNAKE_CASE			):
        __UpperCamelCase  :List[Any]    						=	R'''\b\d{2}\b'''
        if bool(re.search(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			)			):
            __UpperCamelCase  :List[str]    						=	re.search(R'''\d\.\d\d.'''						,    SCREAMING_SNAKE_CASE			).group()
        else:
            __UpperCamelCase  :Any    						=	re.search(R'''\d\.\d.'''						,    SCREAMING_SNAKE_CASE			).group()
        if int(match[0]			) < 6:
            __UpperCamelCase  :List[Any]    						=	old_name.replace(SCREAMING_SNAKE_CASE						,    ''''''			)
            __UpperCamelCase  :List[str]    						=	trimmed_name.replace('''network'''						,    match[0] + '''.meta4D_layers.blocks.''' + match[2:-1]			)
            __UpperCamelCase  :Optional[int]    						=	'''intermediate_stages.''' + trimmed_name
        else:
            __UpperCamelCase  :Tuple    						=	old_name.replace(SCREAMING_SNAKE_CASE						,    ''''''			)
            if int(match[2]			) < num_meta4D_last_stage:
                __UpperCamelCase  :str    						=	trimmed_name.replace('''network'''						,    '''meta4D_layers.blocks.''' + match[2]			)
            else:
                __UpperCamelCase  :List[str]    						=	str(int(match[2]			) - num_meta4D_last_stage			)
                __UpperCamelCase  :Optional[Any]    						=	trimmed_name.replace('''network'''						,    '''meta3D_layers.blocks.''' + layer_index			)
                if "norm1" in old_name:
                    __UpperCamelCase  :Optional[Any]    						=	trimmed_name.replace('''norm1'''						,    '''layernorm1'''			)
                elif "norm2" in old_name:
                    __UpperCamelCase  :Optional[int]    						=	trimmed_name.replace('''norm2'''						,    '''layernorm2'''			)
                elif "fc1" in old_name:
                    __UpperCamelCase  :List[str]    						=	trimmed_name.replace('''fc1'''						,    '''linear_in'''			)
                elif "fc2" in old_name:
                    __UpperCamelCase  :Union[str, Any]    						=	trimmed_name.replace('''fc2'''						,    '''linear_out'''			)

            __UpperCamelCase  :int    						=	'''last_stage.''' + trimmed_name

    elif "network" in old_name and re.search(R'''.\d.'''						,    SCREAMING_SNAKE_CASE			):
        __UpperCamelCase  :Dict    						=	old_name.replace('''network'''						,    '''intermediate_stages'''			)

    if "fc" in new_name:
        __UpperCamelCase  :Optional[Any]    						=	new_name.replace('''fc'''						,    '''convolution'''			)
    elif ("norm1" in new_name) and ("layernorm1" not in new_name):
        __UpperCamelCase  :Optional[Any]    						=	new_name.replace('''norm1'''						,    '''batchnorm_before'''			)
    elif ("norm2" in new_name) and ("layernorm2" not in new_name):
        __UpperCamelCase  :Optional[Any]    						=	new_name.replace('''norm2'''						,    '''batchnorm_after'''			)
    if "proj" in new_name:
        __UpperCamelCase  :str    						=	new_name.replace('''proj'''						,    '''projection'''			)
    if "dist_head" in new_name:
        __UpperCamelCase  :Any    						=	new_name.replace('''dist_head'''						,    '''distillation_classifier'''			)
    elif "head" in new_name:
        __UpperCamelCase  :Dict    						=	new_name.replace('''head'''						,    '''classifier'''			)
    elif "patch_embed" in new_name:
        __UpperCamelCase  :int    						=	'''efficientformer.''' + new_name
    elif new_name == "norm.weight" or new_name == "norm.bias":
        __UpperCamelCase  :int    						=	new_name.replace('''norm'''						,    '''layernorm'''			)
        __UpperCamelCase  :int    						=	'''efficientformer.''' + new_name
    else:
        __UpperCamelCase  :Dict    						=	'''efficientformer.encoder.''' + new_name

    return new_name



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    for key in checkpoint.copy().keys():
        __UpperCamelCase  :Tuple    						=	checkpoint.pop(SCREAMING_SNAKE_CASE			)
        __UpperCamelCase  :Optional[Any]    						=	val

    return checkpoint



def lowerCamelCase						(     ):
    '''simple docstring'''







    __UpperCamelCase  :Tuple    						=	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
    __UpperCamelCase  :Any    						=	Image.open(requests.get(SCREAMING_SNAKE_CASE						,    stream=SCREAMING_SNAKE_CASE			).raw			)

    return image



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Union[str, Any]    						=	torch.load(SCREAMING_SNAKE_CASE						,    map_location='''cpu'''			)['''model''']
    __UpperCamelCase  :Optional[int]    						=	EfficientFormerConfig.from_json_file(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :Optional[int]    						=	EfficientFormerForImageClassificationWithTeacher(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :List[Any]    						=	'''_'''.join(checkpoint_path.split('''/'''			)[-1].split('''.'''			)[0].split('''_'''			)[:-1]			)

    __UpperCamelCase  :Union[str, Any]    						=	config.depths[-1] - config.num_metaad_blocks + 1
    __UpperCamelCase  :List[Any]    						=	convert_torch_checkpoint(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			)

    model.load_state_dict(SCREAMING_SNAKE_CASE			)
    model.eval()

    __UpperCamelCase  :Union[str, Any]    						=	{
        '''bilinear''': PILImageResampling.BILINEAR,
        '''bicubic''': PILImageResampling.BICUBIC,
        '''nearest''': PILImageResampling.NEAREST,
    }

    # prepare image
    __UpperCamelCase  :Any    						=	prepare_img()
    __UpperCamelCase  :int    						=	256
    __UpperCamelCase  :Tuple    						=	224
    __UpperCamelCase  :List[Any]    						=	EfficientFormerImageProcessor(
        size={'''shortest_edge''': image_size}						,    crop_size={'''height''': crop_size, '''width''': crop_size}						,    resample=pillow_resamplings['''bicubic''']						,    )
    __UpperCamelCase  :Tuple    						=	processor(images=SCREAMING_SNAKE_CASE						,    return_tensors='''pt'''			).pixel_values

    # original processing pipeline
    __UpperCamelCase  :List[Any]    						=	Compose(
        [
            Resize(SCREAMING_SNAKE_CASE						,    interpolation=pillow_resamplings['''bicubic''']			),
            CenterCrop(SCREAMING_SNAKE_CASE			),
            ToTensor(),
            Normalize(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			),
        ]			)
    __UpperCamelCase  :Dict    						=	image_transforms(SCREAMING_SNAKE_CASE			).unsqueeze(0			)

    assert torch.allclose(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			)

    __UpperCamelCase  :Optional[int]    						=	model(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :Dict    						=	outputs.logits

    __UpperCamelCase  :Optional[Any]    						=	(1, 1_000)

    if "l1" in model_name:
        __UpperCamelCase  :Optional[Any]    						=	torch.Tensor(
            [-0.1_312, 0.4_353, -1.0_499, -0.5_124, 0.4_183, -0.6_793, -1.3_777, -0.0_893, -0.7_358, -2.4_328]			)
        assert torch.allclose(logits[0, :10]						,    SCREAMING_SNAKE_CASE						,    atol=1e-3			)
        assert logits.shape == expected_shape
    elif "l3" in model_name:
        __UpperCamelCase  :List[Any]    						=	torch.Tensor(
            [-1.3_150, -1.5_456, -1.2_556, -0.8_496, -0.7_127, -0.7_897, -0.9_728, -0.3_052, 0.3_751, -0.3_127]			)
        assert torch.allclose(logits[0, :10]						,    SCREAMING_SNAKE_CASE						,    atol=1e-3			)
        assert logits.shape == expected_shape
    elif "l7" in model_name:
        __UpperCamelCase  :List[Any]    						=	torch.Tensor(
            [-1.0_283, -1.4_131, -0.5_644, -1.3_115, -0.5_785, -1.2_049, -0.7_528, 0.1_992, -0.3_822, -0.0_878]			)
        assert logits.shape == expected_shape
    else:
        raise ValueError(
            f"""Unknown model checkpoint: {checkpoint_path}. Supported version of efficientformer are l1, l3 and l7"""			)

    # Save Checkpoints
    Path(SCREAMING_SNAKE_CASE			).mkdir(exist_ok=SCREAMING_SNAKE_CASE			)
    model.save_pretrained(SCREAMING_SNAKE_CASE			)
    print(f"""Checkpoint successfuly converted. Model saved at {pytorch_dump_path}"""			)
    processor.save_pretrained(SCREAMING_SNAKE_CASE			)
    print(f"""Processor successfuly saved at {pytorch_dump_path}"""			)

    if push_to_hub:
        print('''Pushing model to the hub...'''			)

        model.push_to_hub(
            repo_id=f"""Bearnardd/{pytorch_dump_path}"""						,    commit_message='''Add model'''						,    use_temp_dir=SCREAMING_SNAKE_CASE						,    )
        processor.push_to_hub(
            repo_id=f"""Bearnardd/{pytorch_dump_path}"""						,    commit_message='''Add image processor'''						,    use_temp_dir=SCREAMING_SNAKE_CASE						,    )


if __name__ == "__main__":
    __lowercase											=      argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--pytorch_model_path''',
        default=None,
        type=str,
        required=True,
        help='''Path to EfficientFormer pytorch checkpoint.''',
    )
    parser.add_argument(
        '''--config_file''',
        default=None,
        type=str,
        required=True,
        help='''The json file for EfficientFormer model config.''',
    )
    parser.add_argument(
        '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
    )

    parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image processor to the hub''')
    parser.add_argument(
        '''--no-push_to_hub''',
        dest='''push_to_hub''',
        action='''store_false''',
        help='''Do not push model and image processor to the hub''',
    )
    parser.set_defaults(push_to_hub=True)

    __lowercase											=      parser.parse_args()
    convert_efficientformer_checkpoint(
        checkpoint_path=args.pytorch_model_path,
        efficientformer_config_file=args.config_file,
        pytorch_dump_path=args.pytorch_dump_path,
        push_to_hub=args.push_to_hub,
    )






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 lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE="attention"			):
    '''simple docstring'''







    __UpperCamelCase  :int    						=	params[f"""{prefix}/layers_{i}/{layer_name}/key/kernel"""]
    __UpperCamelCase  :Optional[int]    						=	params[f"""{prefix}/layers_{i}/{layer_name}/out/kernel"""]
    __UpperCamelCase  :List[str]    						=	params[f"""{prefix}/layers_{i}/{layer_name}/query/kernel"""]
    __UpperCamelCase  :Union[str, Any]    						=	params[f"""{prefix}/layers_{i}/{layer_name}/value/kernel"""]
    return k, o, q, v



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=False			):
    '''simple docstring'''







    if split_mlp_wi:
        __UpperCamelCase  :List[str]    						=	params[f"""{prefix}/layers_{i}/mlp/wi_0/kernel"""]
        __UpperCamelCase  :Dict    						=	params[f"""{prefix}/layers_{i}/mlp/wi_1/kernel"""]
        __UpperCamelCase  :Optional[Any]    						=	(wi_a, wi_a)
    else:
        __UpperCamelCase  :List[str]    						=	params[f"""{prefix}/layers_{i}/mlp/wi/kernel"""]

    __UpperCamelCase  :int    						=	params[f"""{prefix}/layers_{i}/mlp/wo/kernel"""]
    return wi, wo



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    return params[f"""{prefix}/layers_{i}/{layer_name}/scale"""]



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    *, SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Tuple    						=	traverse_util.flatten_dict(variables['''target''']			)
    __UpperCamelCase  :Optional[int]    						=	{'''/'''.join(SCREAMING_SNAKE_CASE			): v for k, v in old.items()}

    # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
    __UpperCamelCase  :Union[str, Any]    						=	'''encoder/layers_0/mlp/wi_0/kernel''' in old
    print('''Split MLP:'''						,    SCREAMING_SNAKE_CASE			)

    __UpperCamelCase  :Optional[Any]    						=	collections.OrderedDict()

    # Shared embeddings.
    __UpperCamelCase  :Any    						=	old['''token_embedder/embedding''']

    # Encoder.
    for i in range(SCREAMING_SNAKE_CASE			):
        # Block i, layer 0 (Self Attention).
        __UpperCamelCase  :Tuple    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    '''pre_attention_layer_norm'''			)
        __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase  :Dict    						=	tax_attention_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    '''attention'''			)
        __UpperCamelCase  :str    						=	layer_norm
        __UpperCamelCase  :str    						=	k.T
        __UpperCamelCase  :Optional[Any]    						=	o.T
        __UpperCamelCase  :Optional[int]    						=	q.T
        __UpperCamelCase  :Any    						=	v.T

        # Block i, layer 1 (MLP).
        __UpperCamelCase  :Union[str, Any]    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    '''pre_mlp_layer_norm'''			)
        __UpperCamelCase						, __UpperCamelCase  :Optional[int]    						=	tax_mlp_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    SCREAMING_SNAKE_CASE			)
        __UpperCamelCase  :Any    						=	layer_norm
        if split_mlp_wi:
            __UpperCamelCase  :Optional[Any]    						=	wi[0].T
            __UpperCamelCase  :int    						=	wi[1].T
        else:
            __UpperCamelCase  :Dict    						=	wi.T
        __UpperCamelCase  :str    						=	wo.T

    __UpperCamelCase  :Optional[Any]    						=	old[
        '''encoder/relpos_bias/rel_embedding'''
    ].T
    __UpperCamelCase  :Optional[int]    						=	old['''encoder/encoder_norm/scale''']

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

            # Block i, layer 1 (Cross Attention).
            __UpperCamelCase  :Optional[Any]    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''pre_cross_attention_layer_norm'''			)
            __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase  :int    						=	tax_attention_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''encoder_decoder_attention'''			)
            __UpperCamelCase  :List[str]    						=	layer_norm
            __UpperCamelCase  :Any    						=	k.T
            __UpperCamelCase  :Any    						=	o.T
            __UpperCamelCase  :Optional[int]    						=	q.T
            __UpperCamelCase  :Union[str, Any]    						=	v.T

            # Block i, layer 2 (MLP).
            __UpperCamelCase  :Tuple    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''pre_mlp_layer_norm'''			)
            __UpperCamelCase						, __UpperCamelCase  :Optional[Any]    						=	tax_mlp_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    SCREAMING_SNAKE_CASE			)
            __UpperCamelCase  :Optional[Any]    						=	layer_norm
            if split_mlp_wi:
                __UpperCamelCase  :int    						=	wi[0].T
                __UpperCamelCase  :Any    						=	wi[1].T
            else:
                __UpperCamelCase  :Optional[Any]    						=	wi.T
            __UpperCamelCase  :Dict    						=	wo.T

        __UpperCamelCase  :Optional[int]    						=	old['''decoder/decoder_norm/scale''']
        __UpperCamelCase  :Optional[Any]    						=	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:
            __UpperCamelCase  :List[str]    						=	old['''decoder/logits_dense/kernel'''].T

    return new



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Optional[Any]    						=	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:
        __UpperCamelCase  :Union[str, Any]    						=	state_dict['''shared.weight''']

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

    return state_dict



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :int    						=	checkpoints.load_tax_checkpoint(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :str    						=	convert_tax_to_pytorch(SCREAMING_SNAKE_CASE						,    num_layers=config.num_layers						,    is_encoder_only=SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :Optional[int]    						=	make_state_dict(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			)
    model.load_state_dict(SCREAMING_SNAKE_CASE						,    strict=SCREAMING_SNAKE_CASE			)



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE = False			):
    '''simple docstring'''







    __UpperCamelCase  :Union[str, Any]    						=	TaConfig.from_json_file(SCREAMING_SNAKE_CASE			)
    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:
        __UpperCamelCase  :Dict    						=	TaEncoderModel(SCREAMING_SNAKE_CASE			)
    else:
        __UpperCamelCase  :str    						=	TaForConditionalGeneration(SCREAMING_SNAKE_CASE			)

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

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

    # Verify that we can load the checkpoint.
    model.from_pretrained(SCREAMING_SNAKE_CASE			)
    print('''Done'''			)


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





