
import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch


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


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




def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ = 5_0_0_0_0_0_0_0		):
      UpperCamelCase__    :   Any					=				set()
      UpperCamelCase__    :   Any					=				int((limit - 2_4) ** (1 / 2)		)

      UpperCamelCase__    :   Dict					=				set(range(3 ,		prime_square_limit + 1 ,		2		)		)
      primes.add(2		)
      for p in range(3 ,		prime_square_limit + 1 ,		2		):
            if p not in primes:
                  continue
            primes.difference_update(set(range(p * p ,		prime_square_limit + 1 ,		UpperCamelCase__		)		)		)

      for primea in primes:
            UpperCamelCase__    :   Dict					=				primea * primea
            for primea in primes:
                  UpperCamelCase__    :   str					=				primea * primea * primea
                  if square + cube >= limit - 1_6:
                        break
                  for primea in primes:
                        UpperCamelCase__    :   Dict					=				primea * primea * primea * primea
                        UpperCamelCase__    :   Tuple					=				square + cube + tetr
                        if total >= limit:
                              break
                        ret.add(UpperCamelCase__		)

      return len(UpperCamelCase__		)


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


