






"""simple docstring"""





import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE     )					->	Union[str, Any]:
    """simple docstring"""

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


if __name__ == "__main__":
  A__					: Dict=	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""")
  A__					: Any=	parser.parse_args()
  convert_tf_gptsan_to_pt(args)












"""simple docstring"""





def   lowerCAmelCase_( SCREAMING_SNAKE_CASE					)					->	bool:
    """simple docstring"""

    UpperCamelCase__         =    (1 + 24 * n) ** 0.5
    return ((1 + root) / 6) % 1 == 0







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE = 50_00					)					->	int:
    """simple docstring"""

    UpperCamelCase__         =    [(i * (3 * i - 1)) // 2 for i in range(1 ,	SCREAMING_SNAKE_CASE					)]
    for i, pentagonal_i in enumerate(SCREAMING_SNAKE_CASE					):
        for j in range(SCREAMING_SNAKE_CASE ,	len(SCREAMING_SNAKE_CASE					)					):
            UpperCamelCase__         =    pentagonal_nums[j]
            UpperCamelCase__         =    pentagonal_i + pentagonal_j
            UpperCamelCase__         =    pentagonal_j - pentagonal_i
            if is_pentagonal(SCREAMING_SNAKE_CASE					) and is_pentagonal(SCREAMING_SNAKE_CASE					):
                return b

    return -1


if __name__ == "__main__":
  print(F"""{solution() = }""")




