


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    __lowerCamelCase   (      A__						: int						,						A__						: str						,						A__						: Any					)     ->    Dict:
   return params[f'''{prefix}/{prefix}/relpos_bias/rel_embedding'''][:, i, :]


def    __lowerCamelCase   (      A__						: List[Any]						,						A__						: Optional[int]						,						A__						: Any						,						A__						: str="attention"					)     ->    Dict:
   lowerCamelCase_						:   Optional[int]		= np.ascontiguousarray(params[f'''{prefix}/{prefix}/{layer_name}/key/kernel'''][:, i, :, :]					)
   lowerCamelCase_						:   Union[str, Any]		= k_tmp.reshape(k_tmp.shape[0]						,						k_tmp.shape[1] * k_tmp.shape[2]					)
   lowerCamelCase_						:   List[Any]		= np.ascontiguousarray(params[f'''{prefix}/{prefix}/{layer_name}/out/kernel'''][:, i, :, :]					)
   lowerCamelCase_						:   str		= o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]						,						o_tmp.shape[2]					)
   lowerCamelCase_						:   str		= 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_						:   Tuple		= v_tmp.reshape(v_tmp.shape[0]						,						v_tmp.shape[1] * v_tmp.shape[2]					)
   return k, o, q, v


def    __lowerCamelCase   (      A__						: Dict						,						A__						: Optional[int]						,						A__						: str						,						A__						: List[str]=False					)     ->    Optional[Any]:
   if split_mlp_wi:
      lowerCamelCase_						:   str		= params[f'''{prefix}/{prefix}/mlp/wi_0/kernel'''][:, i, :]
      lowerCamelCase_						:   Dict		= params[f'''{prefix}/{prefix}/mlp/wi_1/kernel'''][:, i, :]
      lowerCamelCase_						:   Tuple		= (wi_a, wi_a)
   else:
      lowerCamelCase_						:   Union[str, Any]		= params[f'''{prefix}/{prefix}/mlp/wi/kernel'''][:, i, :]

   lowerCamelCase_						:   Tuple		= params[f'''{prefix}/{prefix}/mlp/wo/kernel'''][:, i, :]
   return wi, wo


def    __lowerCamelCase   (      A__						: Optional[int]						,						A__						: List[str]						,						A__						: List[Any]						,						A__						: Optional[Any]					)     ->    str:
   return params[f'''{prefix}/{prefix}/{layer_name}/scale'''][:, i]


def    __lowerCamelCase   (      A__						: dict						,						*, A__						: int						,						A__						: bool						,						A__						: bool = False					)     ->    Optional[Any]:
   lowerCamelCase_						:   Union[str, Any]		= traverse_util.flatten_dict(variables["""target"""]					)
   lowerCamelCase_						:   Union[str, Any]		= {"""/""".join(A__					): v for k, v in old.items()}

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

   lowerCamelCase_						:   Union[str, Any]		= collections.OrderedDict()

   # Shared embeddings.
   lowerCamelCase_						:   List[str]		= old["""token_embedder/embedding"""]

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

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

   lowerCamelCase_						:   Union[str, Any]		= old["""encoder/encoder_norm/scale"""]

   if not scalable_attention:
      lowerCamelCase_						:   List[str]		= tax_relpos_bias_lookup(
          A__						,						0						,						"""encoder"""					).T
      lowerCamelCase_						:   Any		= tax_relpos_bias_lookup(
          A__						,						0						,						"""decoder"""					).T

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

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

         # Block i, layer 2 (MLP).
         lowerCamelCase_						:   Union[str, Any]		= tax_layer_norm_lookup(A__						,						A__						,						"""decoder"""						,						"""pre_mlp_layer_norm"""					)
         lowerCamelCase_,	lowerCamelCase_						:   int		= tax_mlp_lookup(A__						,						A__						,						"""decoder"""						,						A__					)
         lowerCamelCase_						:   int		= layer_norm
         if split_mlp_wi:
            lowerCamelCase_						:   str		= wi[0].T
            lowerCamelCase_						:   Dict		= wi[1].T
         else:
            lowerCamelCase_						:   Tuple		= wi.T
         lowerCamelCase_						:   Tuple		= wo.T

         if scalable_attention:
            # convert the rel_embedding of each layer
            lowerCamelCase_						:   List[str]		= tax_relpos_bias_lookup(A__						,						A__						,						"""decoder"""					).T

      lowerCamelCase_						:   Union[str, 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_						:   Optional[Any]		= old["""decoder/logits_dense/kernel"""].T

   return new


def    __lowerCamelCase   (      A__						: Tuple						,						A__						: bool					)     ->    List[Any]:
   lowerCamelCase_						:   List[str]		= 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_						:   Optional[Any]		= 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_						:   int		= state_dict["""shared.weight"""]

   return state_dict


def    __lowerCamelCase   (      A__						: Union[str, Any]						,						A__						: Optional[int]						,						A__						: Optional[int]						,						A__						: Any						,						A__						: Union[str, Any]					)     ->    Optional[int]:
   lowerCamelCase_						:   Optional[Any]		= checkpoints.load_tax_checkpoint(A__					)
   lowerCamelCase_						:   str		= convert_tax_to_pytorch(
       A__						,						num_layers=config.num_layers						,						is_encoder_only=A__						,						scalable_attention=A__					)
   lowerCamelCase_						:   Dict		= make_state_dict(A__						,						A__					)
   model.load_state_dict(A__						,						strict=A__					)


def    __lowerCamelCase   (      A__						: List[Any]						,						A__						: List[str]						,						A__						: List[Any]						,						A__						: bool = False						,						A__						: bool = False						,						)     ->    str:
   lowerCamelCase_						:   Union[str, Any]		= MTaConfig.from_json_file(A__					)
   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_						:   Tuple		= UMTaEncoderModel(A__					)
   else:
      lowerCamelCase_						:   List[str]		= UMTaForConditionalGeneration(A__					)

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

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

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


if __name__ == "__main__":
      snake_case__				:	List[Any]					       =					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,
      )
      snake_case__				:	str					       =					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,
      )









# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import importlib.metadata
import json
import os
from dataclasses import dataclass
from typing import Any, Dict, Union

from packaging import version

from ..utils import is_torch_available, logging


if is_torch_available():
      import torch


snake_case__				:	int					       =					logging.get_logger(__name__)





@dataclass
class   SCREAMING_SNAKE_CASE_			:




       '''simple docstring'''



       def __init__(						self				:				int		,				__a				:				Optional[Any]=False		,				__a				:				Tuple=False		,				__a				:				Tuple=6.0		,				__a				:				str=None		,				__a				:				List[str]=False		,				__a				:				str=False		,				__a				:				Optional[Any]=None		,				__a				:				Dict="fp4"		,				__a				:				Optional[int]=False		,				**__a				:				Dict		,				)					->str:
          lowerCamelCase_						:   str		= load_in_abit
          lowerCamelCase_						:   Union[str, Any]		= load_in_abit
          lowerCamelCase_						:   Tuple		= llm_inta_threshold
          lowerCamelCase_						:   int		= llm_inta_skip_modules
          lowerCamelCase_						:   List[str]		= llm_inta_enable_fpaa_cpu_offload
          lowerCamelCase_						:   Optional[Any]		= llm_inta_has_fpaa_weight
          lowerCamelCase_						:   Dict		= bnb_abit_quant_type
          lowerCamelCase_						:   int		= bnb_abit_use_double_quant

          if bnb_abit_compute_dtype is None:
             lowerCamelCase_						:   int		= torch.floataa
          elif isinstance(__a		,				__a							):
             lowerCamelCase_						:   Tuple		= getattr(__a		,				__a							)
          elif isinstance(__a		,				torch.dtype							):
             lowerCamelCase_						:   List[str]		= bnb_abit_compute_dtype
          else:
             raise ValueError("""bnb_4bit_compute_dtype must be a string or a torch.dtype"""							)

          self.post_init()



       def 				_lowerCAmelCase							(						self				:				Any							)					->Union[str, Any]:
          if not isinstance(self.llm_inta_threshold		,				__a							):
             raise ValueError("""llm_int8_threshold must be a float"""							)

          if self.llm_inta_skip_modules is not None and not isinstance(self.llm_inta_skip_modules		,				__a							):
             raise ValueError("""llm_int8_skip_modules must be a list of strings"""							)
          if not isinstance(self.llm_inta_enable_fpaa_cpu_offload		,				__a							):
             raise ValueError("""llm_int8_enable_fp32_cpu_offload must be a boolean"""							)

          if not isinstance(self.llm_inta_has_fpaa_weight		,				__a							):
             raise ValueError("""llm_int8_has_fp16_weight must be a boolean"""							)

          if self.bnb_abit_compute_dtype is not None and not isinstance(self.bnb_abit_compute_dtype		,				torch.dtype							):
             raise ValueError("""bnb_4bit_compute_dtype must be torch.dtype"""							)

          if not isinstance(self.bnb_abit_quant_type		,				__a							):
             raise ValueError("""bnb_4bit_quant_type must be a string"""							)

          if not isinstance(self.bnb_abit_use_double_quant		,				__a							):
             raise ValueError("""bnb_4bit_use_double_quant must be a boolean"""							)

          if self.load_in_abit and not version.parse(importlib.metadata.version("""bitsandbytes"""							)							) >= version.parse(
              """0.39.0"""							):
             raise ValueError(
                 """4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version"""							)



       def 				_lowerCAmelCase							(						self				:				Optional[int]							)					->int:
          return self.load_in_abit or self.load_in_abit



       def 				_lowerCAmelCase							(						self				:				List[str]							)					->List[Any]:
          if self.load_in_abit:
             return "llm_int8"
          elif self.load_in_abit and self.bnb_abit_quant_type == "fp4":
             return "fp4"
          elif self.load_in_abit and self.bnb_abit_quant_type == "nf4":
             return "nf4"
          else:
             return None



       @classmethod
       def 				_lowerCAmelCase							(						cls				:				int		,				__a				:				Dict		,				__a				:				str		,				**__a				:				Any							)					->Optional[Any]:
          lowerCamelCase_						:   Optional[Any]		= cls(**__a							)

          lowerCamelCase_						:   List[Any]		= []
          for key, value in kwargs.items():
             if hasattr(__a		,				__a							):
                setattr(__a		,				__a		,				__a							)
                to_remove.append(__a							)
          for key in to_remove:
             kwargs.pop(__a		,				__a							)

          if return_unused_kwargs:
             return config, kwargs
          else:
             return config



       def 				_lowerCAmelCase							(						self				:				Optional[Any]		,				__a				:				Union[str, os.PathLike]							)					->Any:
          with open(__a		,				"""w"""		,				encoding="""utf-8"""							) as writer:
             lowerCamelCase_						:   List[Any]		= self.to_dict()
             lowerCamelCase_						:   Dict		= json.dumps(__a		,				indent=2		,				sort_keys=__a							) + """\n"""

             writer.write(__a							)



       def 				_lowerCAmelCase							(						self				:				Tuple							)					->Dict[str, Any]:
          lowerCamelCase_						:   Dict		= copy.deepcopy(self.__dict__							)
          lowerCamelCase_						:   Tuple		= str(output["""bnb_4bit_compute_dtype"""]							).split("""."""							)[1]

          return output



       def __repr__(						self				:				List[str]							)					->Any:
          return F'''{self.__class__.__name__} {self.to_json_string()}'''



       def 				_lowerCAmelCase							(						self				:				Optional[int]		,				__a				:				bool = True							)					->str:
          if use_diff is True:
             lowerCamelCase_						:   Dict		= self.to_diff_dict()
          else:
             lowerCamelCase_						:   Any		= self.to_dict()
          return json.dumps(__a		,				indent=2		,				sort_keys=__a							) + "\n"




       def 				_lowerCAmelCase							(						self				:				Optional[Any]							)					->Dict[str, Any]:
          lowerCamelCase_						:   Any		= self.to_dict()

          # get the default config dict
          lowerCamelCase_						:   List[str]		= BitsAndBytesConfig().to_dict()

          lowerCamelCase_						:   str		= {}

          # only serialize values that differ from the default config
          for key, value in config_dict.items():
             if value != default_config_dict[key]:
                lowerCamelCase_						:   List[Any]		= value

          return serializable_config_dict





