






from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml





class   UpperCamelCase_ (       yaml.SafeLoader  ):



  def 		UpperCamelCase_				(      self ,		snake_case__				)			->			Union[str, Any]:



         """simple docstring"""





         UpperCAmelCase	    =			[self.constructed_objects[key_node] for key_node, _ in node.value]
         UpperCAmelCase	    =			[tuple(snake_case__				) if isinstance(snake_case__ ,		snake_case__				) else key for key in keys]
         UpperCAmelCase	    =			Counter(snake_case__				)
         UpperCAmelCase	    =			[key for key in counter if counter[key] > 1]
         if duplicate_keys:
                raise TypeError(f'''Got duplicate yaml keys: {duplicate_keys}'''				)



  def 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__=False				)			->			Dict:



         """simple docstring"""





         UpperCAmelCase	    =			super().construct_mapping(snake_case__ ,		deep=snake_case__				)
         self._check_no_duplicates_on_constructed_node(snake_case__				)
         return mapping






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			list(readme_content.splitlines()		)
       if full_content and full_content[0] == "---" and "---" in full_content[1:]:
              UpperCAmelCase	    =			full_content[1:].index("""---"""		) + 1
              UpperCAmelCase	    =			'''\n'''.join(full_content[1:sep_idx]		)
              return yamlblock, "\n".join(full_content[sep_idx + 1 :]		)

       return None, "\n".join(__A		)





class   UpperCamelCase_ (       __A  ):
  _A       :					Tuple				=	{'train_eval_index'}  # train-eval-index in the YAML metadata



  @classmethod
  def 		UpperCamelCase_				(      cls ,		snake_case__				)			->			"DatasetMetadata":



         """simple docstring"""





         with open(snake_case__ ,		encoding="""utf-8"""				) as readme_file:
                UpperCAmelCase	    =			_split_yaml_from_readme(readme_file.read()				)
         if yaml_string is not None:
                return cls.from_yaml_string(snake_case__				)
         else:
                return cls()



  def 		UpperCamelCase_				(      self ,		snake_case__				)			->			int:



         """simple docstring"""





         if path.exists():
                with open(snake_case__ ,		encoding="""utf-8"""				) as readme_file:
                       UpperCAmelCase	    =			readme_file.read()
         else:
                UpperCAmelCase	    =			None
         UpperCAmelCase	    =			self._to_readme(snake_case__				)
         with open(snake_case__ ,		"""w""" ,		encoding="""utf-8"""				) as readme_file:
                readme_file.write(snake_case__				)



  def 		UpperCamelCase_				(      self ,		snake_case__ = None				)			->			str:



         """simple docstring"""





         if readme_content is not None:
                UpperCAmelCase	    =			_split_yaml_from_readme(snake_case__				)
                UpperCAmelCase	    =			'''---\n''' + self.to_yaml_string() + '''---\n''' + content
         else:
                UpperCAmelCase	    =			'''---\n''' + self.to_yaml_string() + '''---\n'''
         return full_content



  @classmethod
  def 		UpperCamelCase_				(      cls ,		snake_case__				)			->			"DatasetMetadata":



         """simple docstring"""





         UpperCAmelCase	    =			yaml.load(snake_case__ ,		Loader=_NoDuplicateSafeLoader				) or {}

         # Convert the YAML keys to DatasetMetadata fields
         UpperCAmelCase	    =			{
             (key.replace("""-""" ,		"""_"""				) if key.replace("""-""" ,		"""_"""				) in cls._FIELDS_WITH_DASHES else key): value
             for key, value in metadata_dict.items()
         }
         return cls(**snake_case__				)



  def 		UpperCamelCase_				(      self				)			->			str:



         """simple docstring"""





         return yaml.safe_dump(
             {
                 (key.replace("""_""" ,		"""-"""				) if key in self._FIELDS_WITH_DASHES else key): value
                 for key, value in self.items()
             } ,		sort_keys=snake_case__ ,		allow_unicode=snake_case__ ,		encoding="""utf-8""" ,		).decode("""utf-8"""				)





lowerCAmelCase_	:    Optional[Any]            =      {
    '''image-classification''': [],
    '''translation''': [],
    '''image-segmentation''': [],
    '''fill-mask''': [],
    '''automatic-speech-recognition''': [],
    '''token-classification''': [],
    '''sentence-similarity''': [],
    '''audio-classification''': [],
    '''question-answering''': [],
    '''summarization''': [],
    '''zero-shot-classification''': [],
    '''table-to-text''': [],
    '''feature-extraction''': [],
    '''other''': [],
    '''multiple-choice''': [],
    '''text-classification''': [],
    '''text-to-image''': [],
    '''text2text-generation''': [],
    '''zero-shot-image-classification''': [],
    '''tabular-classification''': [],
    '''tabular-regression''': [],
    '''image-to-image''': [],
    '''tabular-to-text''': [],
    '''unconditional-image-generation''': [],
    '''text-retrieval''': [],
    '''text-to-speech''': [],
    '''object-detection''': [],
    '''audio-to-audio''': [],
    '''text-generation''': [],
    '''conversational''': [],
    '''table-question-answering''': [],
    '''visual-question-answering''': [],
    '''image-to-text''': [],
    '''reinforcement-learning''': [],
    '''voice-activity-detection''': [],
    '''time-series-forecasting''': [],
    '''document-question-answering''': [],
}


if __name__ == "__main__":
 from argparse import ArgumentParser

 lowerCAmelCase_	:    Optional[Any]            =      ArgumentParser(usage='''Validate the yaml metadata block of a README.md file.''')
 ap.add_argument('''readme_filepath''')
 lowerCAmelCase_	:    Optional[int]            =      ap.parse_args()

 lowerCAmelCase_	:    Dict            =      Path(args.readme_filepath)
 lowerCAmelCase_	:    Any            =      DatasetMetadata.from_readme(readme_filepath)
 print(dataset_metadata)
 dataset_metadata.to_readme(readme_filepath)








"""simple docstring"""







import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			os.path.join(args.tf_model_dir	,      """parameters.json"""		)
       UpperCAmelCase	    =			json.loads(open(lowerCAmelCase		).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(lowerCAmelCase		).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(lowerCAmelCase		)
                     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(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                     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(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                     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(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                     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(lowerCAmelCase		)
                                   UpperCAmelCase	    =			"""model.blocks.%d.self_attn.self_attn.k_proj.weight""" % player
                                   UpperCAmelCase	    =			torch.tensor(lowerCAmelCase		)
                                   UpperCAmelCase	    =			"""model.blocks.%d.self_attn.self_attn.v_proj.weight""" % player
                                   UpperCAmelCase	    =			torch.tensor(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                     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(lowerCAmelCase		)
                            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(lowerCAmelCase		)
                     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(lowerCAmelCase		)
                            if key_name.startswith("""model/wte"""		):
                                   UpperCAmelCase	    =			"""lm_head.weight"""
                                   UpperCAmelCase	    =			vnp.copy()  # same in embedded
                                   UpperCAmelCase	    =			torch.tensor(lowerCAmelCase		)
                     elif key_name.startswith("""model/wob"""		):
                            UpperCAmelCase	    =			"""final_logits_bias"""
                            UpperCAmelCase	    =			vnp.copy()  # same in embedded
                            UpperCAmelCase	    =			state.reshape((1, -1)		)
                            UpperCAmelCase	    =			torch.tensor(lowerCAmelCase		)
                     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(lowerCAmelCase		)
                     elif key_name == "model/dense_1/bias":
                            UpperCAmelCase	    =			"""model.last_project.bias"""
                            UpperCAmelCase	    =			vnp.copy()  # same because it is one dimensional
                            UpperCAmelCase	    =			torch.tensor(lowerCAmelCase		)
       torch.save(lowerCAmelCase	,      args.output		)


if __name__ == "__main__":
 lowerCAmelCase_	:    Union[str, 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_	:    Any            =      parser.parse_args()
 convert_tf_gptsan_to_pt(args)
