






import copy
import tempfile
import unittest

from huggingface_hub import HfFolder, delete_repo
from parameterized import parameterized
from requests.exceptions import HTTPError

from transformers import AutoConfig, GenerationConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test





class   UpperCamelCase_ (       unittest.TestCase  ):



  @parameterized.expand([(None,), ("""foo.json""",)]				)
  def 		UpperCamelCase_				(      self ,		snake_case__				)			->			Optional[int]:



         """simple docstring"""





         UpperCAmelCase	    =			GenerationConfig(
             do_sample=snake_case__ ,		temperature=0.7 ,		length_penalty=1.0 ,		bad_words_ids=[[1, 2, 3], [4, 5]] ,		)
         with tempfile.TemporaryDirectory() as tmp_dir:
                config.save_pretrained(snake_case__ ,		config_name=snake_case__				)
                UpperCAmelCase	    =			GenerationConfig.from_pretrained(snake_case__ ,		config_name=snake_case__				)

         # Checks parameters that were specified
         self.assertEqual(loaded_config.do_sample ,		snake_case__				)
         self.assertEqual(loaded_config.temperature ,		0.7				)
         self.assertEqual(loaded_config.length_penalty ,		1.0				)
         self.assertEqual(loaded_config.bad_words_ids ,		[[1, 2, 3], [4, 5]]				)

         # Checks parameters that were not specified (defaults)
         self.assertEqual(loaded_config.top_k ,		50				)
         self.assertEqual(loaded_config.max_length ,		20				)
         self.assertEqual(loaded_config.max_time ,		snake_case__				)



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



         """simple docstring"""





         UpperCAmelCase	    =			AutoConfig.from_pretrained("""gpt2"""				)
         UpperCAmelCase	    =			GenerationConfig.from_model_config(snake_case__				)
         UpperCAmelCase	    =			GenerationConfig()

         # The generation config has loaded a few non-default parameters from the model config
         self.assertNotEqual(snake_case__ ,		snake_case__				)

         # One of those parameters is eos_token_id -- check if it matches
         self.assertNotEqual(generation_config_from_model.eos_token_id ,		default_generation_config.eos_token_id				)
         self.assertEqual(generation_config_from_model.eos_token_id ,		model_config.eos_token_id				)



  def 		UpperCamelCase_				(      self				)			->			Dict:



         """simple docstring"""





         UpperCAmelCase	    =			GenerationConfig()
         UpperCAmelCase	    =			{
             """max_new_tokens""": 10_24,
             """foo""": """bar""",
         }
         UpperCAmelCase	    =			copy.deepcopy(snake_case__				)
         UpperCAmelCase	    =			generation_config.update(**snake_case__				)

         # update_kwargs was not modified (no side effects)
         self.assertEqual(snake_case__ ,		snake_case__				)

         # update_kwargs was used to update the config on valid attributes
         self.assertEqual(generation_config.max_new_tokens ,		10_24				)

         # `.update()` returns a dictionary of unused kwargs
         self.assertEqual(snake_case__ ,		{"""foo""": """bar"""}				)



  def 		UpperCamelCase_				(      self				)			->			str:



         """simple docstring"""





         UpperCAmelCase	    =			GenerationConfig()
         UpperCAmelCase	    =			"""bar"""

         with tempfile.TemporaryDirectory("""test-generation-config"""				) as tmp_dir:
                generation_config.save_pretrained(snake_case__				)

                UpperCAmelCase	    =			GenerationConfig.from_pretrained(snake_case__				)
         # update_kwargs was used to update the config on valid attributes
         self.assertEqual(new_config.foo ,		"""bar"""				)

         UpperCAmelCase	    =			GenerationConfig.from_model_config(snake_case__				)
         assert not hasattr(snake_case__ ,		"""foo"""				)  # no new kwargs should be initialized if from config



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			GenerationConfig()
         self.assertEqual(default_config.temperature ,		1.0				)
         self.assertEqual(default_config.do_sample ,		snake_case__				)
         self.assertEqual(default_config.num_beams ,		1				)

         UpperCAmelCase	    =			GenerationConfig(
             do_sample=snake_case__ ,		temperature=0.7 ,		length_penalty=1.0 ,		bad_words_ids=[[1, 2, 3], [4, 5]] ,		)
         self.assertEqual(config.temperature ,		0.7				)
         self.assertEqual(config.do_sample ,		snake_case__				)
         self.assertEqual(config.num_beams ,		1				)

         with tempfile.TemporaryDirectory() as tmp_dir:
                config.save_pretrained(snake_case__				)
                UpperCAmelCase	    =			GenerationConfig.from_pretrained(snake_case__ ,		temperature=1.0				)

         self.assertEqual(loaded_config.temperature ,		1.0				)
         self.assertEqual(loaded_config.do_sample ,		snake_case__				)
         self.assertEqual(loaded_config.num_beams ,		1				)  # default value





@is_staging_test
class   UpperCamelCase_ (       unittest.TestCase  ):



  @classmethod
  def 		UpperCamelCase_				(      cls				)			->			Union[str, Any]:



         """simple docstring"""





         UpperCAmelCase	    =			TOKEN
         HfFolder.save_token(snake_case__				)



  @classmethod
  def 		UpperCamelCase_				(      cls				)			->			Optional[Any]:



         """simple docstring"""





         try:
                delete_repo(token=cls._token ,		repo_id="""test-generation-config"""				)
         except HTTPError:
                pass

         try:
                delete_repo(token=cls._token ,		repo_id="""valid_org/test-generation-config-org"""				)
         except HTTPError:
                pass



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			GenerationConfig(
             do_sample=snake_case__ ,		temperature=0.7 ,		length_penalty=1.0 ,		)
         config.push_to_hub("""test-generation-config""" ,		use_auth_token=self._token				)

         UpperCAmelCase	    =			GenerationConfig.from_pretrained(f'''{USER}/test-generation-config'''				)
         for k, v in config.to_dict().items():
                if k != "transformers_version":
                       self.assertEqual(snake_case__ ,		getattr(snake_case__ ,		snake_case__				)				)

        # Reset repo
         delete_repo(token=self._token ,		repo_id="""test-generation-config"""				)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
                config.save_pretrained(
                    snake_case__ ,		repo_id="""test-generation-config""" ,		push_to_hub=snake_case__ ,		use_auth_token=self._token				)

         UpperCAmelCase	    =			GenerationConfig.from_pretrained(f'''{USER}/test-generation-config'''				)
         for k, v in config.to_dict().items():
                if k != "transformers_version":
                       self.assertEqual(snake_case__ ,		getattr(snake_case__ ,		snake_case__				)				)



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



         """simple docstring"""





         UpperCAmelCase	    =			GenerationConfig(
             do_sample=snake_case__ ,		temperature=0.7 ,		length_penalty=1.0 ,		)
         config.push_to_hub("""valid_org/test-generation-config-org""" ,		use_auth_token=self._token				)

         UpperCAmelCase	    =			GenerationConfig.from_pretrained("""valid_org/test-generation-config-org"""				)
         for k, v in config.to_dict().items():
                if k != "transformers_version":
                       self.assertEqual(snake_case__ ,		getattr(snake_case__ ,		snake_case__				)				)

        # Reset repo
         delete_repo(token=self._token ,		repo_id="""valid_org/test-generation-config-org"""				)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
                config.save_pretrained(
                    snake_case__ ,		repo_id="""valid_org/test-generation-config-org""" ,		push_to_hub=snake_case__ ,		use_auth_token=self._token				)

         UpperCAmelCase	    =			GenerationConfig.from_pretrained("""valid_org/test-generation-config-org"""				)
         for k, v in config.to_dict().items():
                if k != "transformers_version":
                       self.assertEqual(snake_case__ ,		getattr(snake_case__ ,		snake_case__				)				)








"""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)
