





'''simple docstring'''


import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch
def 							UpperCAmelCase_  (						__lowerCamelCase					:					List[Any] ):
   lowercase_			:Dict    			= os.path.join(args.tf_model_dir	,"parameters.json" )
   lowercase_			:str    			= json.loads(open(__snake_case ).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" ):
      lowercase_			:Optional[int]    			= args.output + ".pt"
   lowercase_			:Tuple    			= OrderedDict()
   with tf.device("/CPU:0" ):
      lowercase_			:str    			= tf.train.load_checkpoint(args.tf_model_dir )
      lowercase_			:List[Any]    			= reader.get_variable_to_shape_map()
      for key_name in shapes.keys():
         lowercase_			:Dict    			= reader.get_tensor(__snake_case ).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" ):
               lowercase_			:List[str]    			= int(key_name[9] )
            elif key_name.startswith("pasts/out" ):
               lowercase_			:Optional[int]    			= 8
            lowercase_			:List[str]    			= "model.sqout.%d.weight" % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
            lowercase_			:Any    			= vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
            lowercase_			:int    			= torch.tensor(__snake_case )
         elif key_name.startswith("model/moe" ):
            lowercase_			:Optional[int]    			= int(key_name[9:].split("/" )[0] )
            if key_name.endswith("/switch_gating/kernel" ):
               lowercase_			:Any    			= "model.blocks.%d.feed_forward.mlp.router.classifier.weight" % player
               lowercase_			:Any    			= vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
               lowercase_			:str    			= torch.tensor(__snake_case )
            elif key_name.endswith("/softmlp/kernel" ):
               lowercase_			:Tuple    			= "model.blocks.%d.feed_forward.soft_bypass_mlp.weight" % player
               lowercase_			:Any    			= vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
               lowercase_			:List[str]    			= torch.tensor(__snake_case )
            elif key_name.endswith("/wo/kernel" ) or key_name.endswith("/wi/kernel" ):
               lowercase_			:Tuple    			= key_name[-9:-7]
               for i in range(16 ):
                  lowercase_			:Union[str, Any]    			= "model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight" % (player, i, nlayer)
                  lowercase_			:Optional[Any]    			= (
                      vnp[i].transpose([1, 0] ).copy()
                  )  # In Mesh-Tensorflow, it is one array, so it is divided
                  lowercase_			:List[str]    			= torch.tensor(__snake_case )
         elif key_name.startswith("model/mlp" ):
            lowercase_			:Tuple    			= int(key_name[9:].split("/" )[0] )
            if key_name.endswith("/p1/kernel" ):
               lowercase_			:List[Any]    			= "model.blocks.%d.feed_forward.mlp.wi.weight" % player
               lowercase_			:Optional[Any]    			= vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
               lowercase_			:int    			= torch.tensor(__snake_case )
            elif key_name.endswith("/p1/bias" ):
               lowercase_			:int    			= "model.blocks.%d.feed_forward.mlp.wi.bias" % player
               lowercase_			:Any    			= vnp.copy()  # same because it is one dimensional
               lowercase_			:int    			= torch.tensor(__snake_case )
            elif key_name.endswith("/p2/kernel" ):
               lowercase_			:List[Any]    			= "model.blocks.%d.feed_forward.mlp.wo.weight" % player
               lowercase_			:Dict    			= vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
               lowercase_			:Union[str, Any]    			= torch.tensor(__snake_case )
            elif key_name.endswith("/p2/bias" ):
               lowercase_			:int    			= "model.blocks.%d.feed_forward.mlp.wo.bias" % player
               lowercase_			:Tuple    			= vnp.copy()  # same because it is one dimensional
               lowercase_			:Any    			= torch.tensor(__snake_case )
         elif key_name.startswith("model/ln" ):
            lowercase_			:Any    			= int(key_name[8:].split("/" )[0] )
            if key_name.endswith("/b" ):
               lowercase_			:int    			= "model.blocks.%d.feed_forward.norm.bias" % player
               lowercase_			:Dict    			= vnp.copy()  # same because it is one dimensional
               lowercase_			:Optional[int]    			= torch.tensor(__snake_case )
            elif key_name.endswith("/g" ):
               lowercase_			:Tuple    			= "model.blocks.%d.feed_forward.norm.weight" % player
               lowercase_			:List[str]    			= vnp.copy()  # same because it is one dimensional
               lowercase_			:List[str]    			= torch.tensor(__snake_case )
         elif key_name.startswith("model/att" ):
            lowercase_			:str    			= int(key_name[9:].split("/" )[0] )
            if key_name.endswith("/qkv/kernel" ):
               lowercase_			:Any    			= vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
               lowercase_			:Optional[Any]    			= state[:, 0, :, :]
               lowercase_			:str    			= state[:, 1, :, :]
               lowercase_			:Tuple    			= state[:, 2, :, :]
               lowercase_			:Any    			= (
                   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
               lowercase_			:List[str]    			= (
                   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
               lowercase_			:Tuple    			= (
                   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
               lowercase_			:Optional[Any]    			= "model.blocks.%d.self_attn.self_attn.q_proj.weight" % player
               lowercase_			:List[str]    			= torch.tensor(__snake_case )
               lowercase_			:Optional[Any]    			= "model.blocks.%d.self_attn.self_attn.k_proj.weight" % player
               lowercase_			:Dict    			= torch.tensor(__snake_case )
               lowercase_			:Any    			= "model.blocks.%d.self_attn.self_attn.v_proj.weight" % player
               lowercase_			:Optional[int]    			= torch.tensor(__snake_case )
            elif key_name.endswith("/o/kernel" ):
               lowercase_			:int    			= "model.blocks.%d.self_attn.self_attn.out_proj.weight" % player
               lowercase_			:List[Any]    			= (
                   vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]] ).transpose([1, 0] ).copy()
               )  # Mesh-Tensorflow is a diagonal matrix
               lowercase_			:Dict    			= torch.tensor(__snake_case )
         elif key_name.startswith("model/an" ):
            lowercase_			:Optional[int]    			= int(key_name[8:].split("/" )[0] )
            if key_name.endswith("/b" ):
               lowercase_			:str    			= "model.blocks.%d.self_attn.norm.bias" % player
               lowercase_			:List[Any]    			= vnp.copy()  # same because it is one dimensional
               lowercase_			:Optional[int]    			= torch.tensor(__snake_case )
            elif key_name.endswith("/g" ):
               lowercase_			:Union[str, Any]    			= "model.blocks.%d.self_attn.norm.weight" % player
               lowercase_			:List[Any]    			= vnp.copy()  # same because it is one dimensional
               lowercase_			:Dict    			= torch.tensor(__snake_case )
         elif (
             key_name.startswith("model/wte" )
             or key_name.startswith("model/wpe" )
             or key_name.startswith("model/ete" )
         ):
            lowercase_			:List[str]    			= {"wte": "embed_tokens", "wpe": "position_embeddings", "ete": "extra_position_embeddings"}[
                key_name[-3:]
            ]
            lowercase_			:Optional[Any]    			= "model.%s.weight" % nlayer
            lowercase_			:Optional[int]    			= vnp.copy()  # same in embedded
            lowercase_			:Tuple    			= torch.tensor(__snake_case )
            if key_name.startswith("model/wte" ):
               lowercase_			:List[str]    			= "lm_head.weight"
               lowercase_			:str    			= vnp.copy()  # same in embedded
               lowercase_			:Any    			= torch.tensor(__snake_case )
         elif key_name.startswith("model/wob" ):
            lowercase_			:Optional[Any]    			= "final_logits_bias"
            lowercase_			:List[Any]    			= vnp.copy()  # same in embedded
            lowercase_			:Dict    			= state.reshape((1, -1) )
            lowercase_			:str    			= torch.tensor(__snake_case )
         elif key_name == "model/dense/kernel":
            lowercase_			:Dict    			= "model.last_project.weight"
            lowercase_			:Optional[Any]    			= vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
            lowercase_			:int    			= torch.tensor(__snake_case )
         elif key_name == "model/dense_1/bias":
            lowercase_			:List[Any]    			= "model.last_project.bias"
            lowercase_			:Optional[int]    			= vnp.copy()  # same because it is one dimensional
            lowercase_			:Optional[Any]    			= torch.tensor(__snake_case )
   torch.save(__snake_case	,args.output )


if __name__ == "__main__":
 lowerCAmelCase					:     Optional[Any]												=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''')
 lowerCAmelCase					:     Optional[int]												=parser.parse_args()
 convert_tf_gptsan_to_pt(args)




"""simple docstring"""





from sklearn.metrics import matthews_corrcoef

import datasets


_a														=		"""
Compute the Matthews correlation coefficient (MCC)

The Matthews correlation coefficient is used in machine learning as a
measure of the quality of binary and multiclass classifications. It takes
into account true and false positives and negatives and is generally
regarded as a balanced measure which can be used even if the classes are of
very different sizes. The MCC is in essence a correlation coefficient value
between -1 and +1. A coefficient of +1 represents a perfect prediction, 0
an average random prediction and -1 an inverse prediction.  The statistic
is also known as the phi coefficient. [source: Wikipedia]
"""

_a														=		"""
Args:
    predictions (list of int): Predicted labels, as returned by a model.
    references (list of int): Ground truth labels.
    sample_weight (list of int, float, or bool): Sample weights. Defaults to `None`.
Returns:
    matthews_correlation (dict containing float): Matthews correlation.
Examples:
    Example 1, a basic example with only predictions and references as inputs:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3])
        >>> print(round(results['matthews_correlation'], 2))
        0.54

    Example 2, the same example as above, but also including sample weights:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3],
        ...                                     sample_weight=[0.5, 3, 1, 1, 1, 2])
        >>> print(round(results['matthews_correlation'], 2))
        0.1

    Example 3, the same example as above, but with sample weights that cause a negative correlation:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3],
        ...                                     sample_weight=[0.5, 1, 0, 0, 0, 1])
        >>> print(round(results['matthews_correlation'], 2))
        -0.25
"""

_a														=		"""\
@article{scikit-learn,
  title={Scikit-learn: Machine Learning in {P}ython},
  author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
  journal={Journal of Machine Learning Research},
  volume={12},
  pages={2825--2830},
  year={2011}
}
"""






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION      ,     _KWARGS_DESCRIPTION	)
class     _UpperCAmelCase(    datasets.Metric	):



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






							'''simple docstring'''
							return datasets.MetricInfo(
							    description=_DESCRIPTION				,		citation=_CITATION				,		inputs_description=_KWARGS_DESCRIPTION				,		features=datasets.Features(
							        {
							            '''predictions''': datasets.Value('''int32'''),
							            '''references''': datasets.Value('''int32'''),
							        })				,		reference_urls=[
							        '''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.matthews_corrcoef.html'''
							    ]				,		)




		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a=None)       ->       Dict:






							'''simple docstring'''
							return {
							    "matthews_correlation": float(matthews_corrcoef(__a				,		__a				,		sample_weight=__a)),
							}






