

"""simple docstring"""

import concurrent.futures
import json
import os
import shutil
import tempfile
import unittest

from transformers import AutoTokenizer, PreTrainedTokenizerFast
from transformers.testing_utils import require_tokenizers

from ..test_tokenization_common import TokenizerTesterMixin


@require_tokenizers
class       _snake_case  (    UpperCamelCase    ,      unittest.TestCase							):
						__snake_case			      =  PreTrainedTokenizerFast
						__snake_case			      =  False
						__snake_case			      =  True
						__snake_case			      =  'tokenizer_file'







						def 							_UpperCamelCase     (self				)      ->    Tuple:
											__lowercase        =      False  # because we don't have pretrained_vocab_files_map
											super().setUp()
											__lowercase        =      True

											__lowercase        =      ['''robot-test/dummy-tokenizer-fast''', '''robot-test/dummy-tokenizer-wordlevel''']
											__lowercase        =      '''SaulLu/dummy-tokenizer-bytelevel-bpe'''

											# Inclusion of 2 tokenizers to test different types of models (Unigram and WordLevel for the moment)
											__lowercase        =      [(PreTrainedTokenizerFast, model_path, {}) for model_path in model_paths]

											__lowercase        =      PreTrainedTokenizerFast.from_pretrained(model_paths[0]				)
											tokenizer.save_pretrained(self.tmpdirname				)







						def 							_UpperCamelCase     (self				)      ->    Any:
											# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
											# model
											pass







						def 							_UpperCamelCase     (self				)      ->    Optional[int]:
											# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
											# model
											pass







						def 							_UpperCamelCase     (self				)      ->    Optional[int]:
											# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
											# model
											pass







						def 							_UpperCamelCase     (self				)      ->    Dict:
											# PreTrainedTokenizerFast doesn't have tokenizer_file in its signature
											pass







						def 							_UpperCamelCase     (self				)      ->    Dict:
											__lowercase        =      self.tmpdirname
											# Here we want to test the 2 available tokenizers that use 2 different types of models: Unigram and WordLevel.
											for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
																with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''				):
																					try:
																										__lowercase        =      tempfile.mkdtemp()
																										__lowercase        =      self.rust_tokenizer_class.from_pretrained(UpperCAmelCase_  ,				**UpperCAmelCase_				)

																										tokenizer.save_pretrained(self.tmpdirname				)
																										super().test_training_new_tokenizer()
																					finally:
																										# Even if the test fails, we must be sure that the folder is deleted and that the default tokenizer
																										# is restored
																										shutil.rmtree(self.tmpdirname				)
																										__lowercase        =      tmpdirname_orig







						def 							_UpperCamelCase     (self				)      ->    Any:
											__lowercase        =      self.tmpdirname
											# Here we want to test the 2 available tokenizers that use 2 different types of models: Unigram and WordLevel.
											for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
																with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''				):
																					try:
																										__lowercase        =      tempfile.mkdtemp()
																										__lowercase        =      self.rust_tokenizer_class.from_pretrained(UpperCAmelCase_  ,				**UpperCAmelCase_				)

																										tokenizer.save_pretrained(self.tmpdirname				)
																										super().test_training_new_tokenizer_with_special_tokens_change()
																					finally:
																										# Even if the test fails, we must be sure that the folder is deleted and that the default tokenizer
																										# is restored
																										shutil.rmtree(self.tmpdirname				)
																										__lowercase        =      tmpdirname_orig







						def 							_UpperCamelCase     (self				)      ->    List[str]:
											__lowercase        =      self.rust_tokenizer_class.from_pretrained(self.bytelevel_bpe_model_name				)

											__lowercase        =      ('''a''' for _ in range(1_0_0_0				))
											__lowercase        =      tokenizer.train_new_from_iterator(text_iterator=UpperCAmelCase_  ,				length=1_0_0_0  ,				vocab_size=5_0				)

											__lowercase        =      new_tokenizer.encode('''a🤗'''				)
											self.assertEqual(UpperCAmelCase_  ,				[6_4, 1_7_2, 2_5_3, 9_7, 2_4_5]				)



@require_tokenizers
class       _snake_case  (    unittest.TestCase							):







						def 							_UpperCamelCase     (self				)      ->    List[Any]:
											__lowercase        =      AutoTokenizer.from_pretrained('''bert-base-cased'''				)
											__lowercase        =      json.loads(tokenizer._tokenizer.to_str()				)
											__lowercase        =      len(UpperCAmelCase_				)

											with tempfile.TemporaryDirectory() as tmp_dir:
																# Hack to save this in the tokenizer_config.json
																__lowercase        =      ['''tokenizer.4.0.0.json''']
																tokenizer.save_pretrained(UpperCAmelCase_				)
																json.dump(UpperCAmelCase_  ,				open(os.path.join(UpperCAmelCase_  ,				'''tokenizer.4.0.0.json'''				)  ,				'''w'''				)				)

																# This should pick the new tokenizer file as the version of Transformers is > 4.0.0
																__lowercase        =      AutoTokenizer.from_pretrained(UpperCAmelCase_				)
																self.assertEqual(len(UpperCAmelCase_				)  ,				len(UpperCAmelCase_				) + 1				)
																__lowercase        =      json.loads(new_tokenizer._tokenizer.to_str()				)
																self.assertIn('''huggingface'''  ,				json_tokenizer['''model''']['''vocab''']				)

																# Will need to be adjusted if we reach v42 and this test is still here.
																# Should pick the old tokenizer file as the version of Transformers is < 4.0.0
																shutil.move(os.path.join(UpperCAmelCase_  ,				'''tokenizer.4.0.0.json'''				)  ,				os.path.join(UpperCAmelCase_  ,				'''tokenizer.42.0.0.json'''				)				)
																__lowercase        =      ['''tokenizer.42.0.0.json''']
																tokenizer.save_pretrained(UpperCAmelCase_				)
																__lowercase        =      AutoTokenizer.from_pretrained(UpperCAmelCase_				)
																self.assertEqual(len(UpperCAmelCase_				)  ,				len(UpperCAmelCase_				)				)
																__lowercase        =      json.loads(new_tokenizer._tokenizer.to_str()				)
																self.assertNotIn('''huggingface'''  ,				json_tokenizer['''model''']['''vocab''']				)







						def 							_UpperCamelCase     (self				)      ->    Union[str, Any]:
											# This repo has two tokenizer files, one for v4.0.0 and above with an added token, one for versions lower.
											__lowercase        =      '''hf-internal-testing/test-two-tokenizers'''

											# This should pick the new tokenizer file as the version of Transformers is > 4.0.0
											__lowercase        =      AutoTokenizer.from_pretrained(UpperCAmelCase_				)
											self.assertEqual(len(UpperCAmelCase_				)  ,				2_8_9_9_7				)
											__lowercase        =      json.loads(tokenizer._tokenizer.to_str()				)
											self.assertIn('''huggingface'''  ,				json_tokenizer['''model''']['''vocab''']				)

											# Testing an older version by monkey-patching the version in the module it's used.
											import transformers as old_transformers

											__lowercase        =      '''3.0.0'''
											__lowercase        =      old_transformers.models.auto.AutoTokenizer.from_pretrained(UpperCAmelCase_				)
											self.assertEqual(len(UpperCAmelCase_				)  ,				2_8_9_9_6				)
											__lowercase        =      json.loads(old_tokenizer._tokenizer.to_str()				)
											self.assertNotIn('''huggingface'''  ,				json_tokenizer['''model''']['''vocab''']				)



@require_tokenizers
class       _snake_case  (    unittest.TestCase							):







						def 							_UpperCamelCase     (self				)      ->    List[Any]:
											# See https://github.com/huggingface/transformers/pull/12550
											# and https://github.com/huggingface/tokenizers/issues/537
											__lowercase        =      PreTrainedTokenizerFast.from_pretrained('''robot-test/dummy-tokenizer-wordlevel'''				)
											__lowercase        =      '''The Matrix is a 1999 science fiction action film.'''

											with concurrent.futures.ThreadPoolExecutor() as executor:
																__lowercase        =      [executor.submit(self.fetch  ,				UpperCAmelCase_  ,				UpperCAmelCase_				) for i in range(1_0				)]
																__lowercase        =      [future.result() for future in futures]
																self.assertEqual(UpperCAmelCase_  ,				[[1, 1_0, 0, 8, 0, 1_8, 0, 0, 0, 2] for i in range(1_0				)]				)







						def 							_UpperCamelCase     (self  ,				UpperCAmelCase_  ,				UpperCAmelCase_				)      ->    Any:
											return tokenizer.encode(UpperCAmelCase_  ,				truncation='''longest_first'''  ,				padding='''longest'''				)





"""simple docstring"""

import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig





def 					__UpperCamelCase (     a__:			List[Any]			,						a__:			int			,						a__:			Dict=0      )					->      Any:




					"""simple docstring"""



					if name is None:
										__lowercase        =      None
					else:
										__lowercase        =      '''.''' * max(0			,						spaces - 2      ) + '''# {:''' + str(50 - spaces      ) + '''s}'''
										__lowercase        =      fmt.format(a__      )

					# Print and recurse (if needed).
					if isinstance(a__			,						a__      ):
										if msg is not None:
															print(a__      )
										for k in val.keys():
															recursive_print(a__			,						val[k]			,						spaces + 2      )
					elif isinstance(a__			,						torch.Tensor      ):
										print(a__			,						''':'''			,						val.size()      )
					else:
										print(a__			,						''':'''			,						a__      )





def 					__UpperCamelCase (     a__:			Optional[int]			,						a__:			Tuple			,						a__:			List[Any]			,						a__:			List[str]      )					->      str:




					"""simple docstring"""



					__lowercase        =      param.size()
					# other versions store [num_heads * num_splits * hidden_size, :]
					__lowercase        =      (num_heads, num_splits, hidden_size) + input_shape[1:]
					__lowercase        =      param.view(*a__      )
					__lowercase        =      param.transpose(0			,						1      ).contiguous()
					__lowercase        =      param.view(*a__      )
					return param





def 					__UpperCamelCase (     a__:			Tuple			,						a__:			Optional[int]      )					->      List[str]:




					"""simple docstring"""



					__lowercase        =      config.n_positions
					__lowercase        =      config.n_layer
					__lowercase        =      config.vocab_size
					__lowercase        =      config.n_head
					__lowercase        =      config.n_embd // config.n_head

					__lowercase        =      sd_megatron['''model.language_model.embedding.word_embeddings.weight'''][:vocab_size, :]
					__lowercase        =      {
					    '''transformer.wte.weight''': word_embeddings,
					    '''transformer.wpe.weight''': sd_megatron['''model.language_model.embedding.position_embeddings.weight'''],
					    '''transformer.ln_f.weight''': sd_megatron['''model.language_model.encoder.final_layernorm.weight'''],
					    '''transformer.ln_f.bias''': sd_megatron['''model.language_model.encoder.final_layernorm.bias'''],
					}

					__lowercase        =      '''model.language_model.encoder.layers.'''
					for i in range(a__      ):
										__lowercase        =      torch.tril(torch.ones((n_positions, n_positions)			,						dtype=torch.bool      )      )
										__lowercase        =      causal_mask.view(1			,						1			,						a__			,						a__      )
										__lowercase        =      causal_mask
										__lowercase        =      torch.tensor(-1e4			,						dtype=torch.bfloataa      )

										__lowercase        =      sd_megatron[f'''{pf}{i}.input_layernorm.weight''']
										__lowercase        =      sd_megatron[f'''{pf}{i}.input_layernorm.bias''']

										__lowercase        =      sd_megatron[f'''{pf}{i}.self_attention.query_key_value.weight''']
										__lowercase        =      fix_query_key_value_ordering(a__			,						3			,						a__			,						a__      )
										__lowercase        =      vala.transpose(0			,						1      ).contiguous()

										__lowercase        =      sd_megatron[f'''{pf}{i}.self_attention.query_key_value.bias''']
										__lowercase        =      fix_query_key_value_ordering(a__			,						3			,						a__			,						a__      )
										__lowercase        =      vala

										__lowercase        =      sd_megatron[f'''{pf}{i}.self_attention.dense.weight'''].transpose(
										    0			,						1      )
										__lowercase        =      sd_megatron[f'''{pf}{i}.self_attention.dense.bias''']
										__lowercase        =      sd_megatron[f'''{pf}{i}.post_attention_layernorm.weight''']
										__lowercase        =      sd_megatron[f'''{pf}{i}.post_attention_layernorm.bias''']
										__lowercase        =      sd_megatron[f'''{pf}{i}.mlp.dense_h_to_4h.weight'''].transpose(0			,						1      )
										__lowercase        =      sd_megatron[f'''{pf}{i}.mlp.dense_h_to_4h.bias''']
										__lowercase        =      sd_megatron[f'''{pf}{i}.mlp.dense_4h_to_h.weight'''].transpose(
										    0			,						1      )
										__lowercase        =      sd_megatron[f'''{pf}{i}.mlp.dense_4h_to_h.bias''']

					# For LM head, transformers' wants the matrix to weight embeddings.
					__lowercase        =      word_embeddings

					return sd_hf





def 					__UpperCamelCase (     a__:			Any			,						a__:			List[Any]      )					->      List[Any]:




					"""simple docstring"""



					__lowercase        =      6_40_00
					__lowercase        =      config_megatron['''encoder_seq_length''']
					__lowercase        =      config_megatron['''hidden_size''']
					__lowercase        =      config_megatron['''num_layers''']
					__lowercase        =      config_megatron['''num_attention_heads''']
					__lowercase        =      config_megatron['''ffn_hidden_size''']
					__lowercase        =      '''gelu'''
					__lowercase        =      0.1
					__lowercase        =      0.1
					__lowercase        =      0.1
					__lowercase        =      config_megatron['''layernorm_epsilon''']  # 1e-5
					__lowercase        =      config_megatron['''init_method_std''']  # 0.02
					__lowercase        =      config_megatron['''apply_query_key_layer_scaling''']  # True
					__lowercase        =      True
					__lowercase        =      True

					# This identifies the 6.7B (7B) model which uses a different tokenizer
					if config_megatron["hidden_size"] == 40_96:
										__lowercase        =      1  # <|endoftext|>
										__lowercase        =      1  # <|endoftext|>
										__lowercase        =      0  # <unk>
					else:
										__lowercase        =      2  # <s>
										__lowercase        =      3  # <|endoftext|>
										__lowercase        =      0  # <pad>

					return config_hf





def 					__UpperCamelCase (     a__:			Dict      )					->      Union[str, Any]:




					"""simple docstring"""



					print(a__      )

					__lowercase        =      args.checkpoint_path
					__lowercase        =      args.save_path
					if isfile(a__      ):
										raise FileNotFoundError(f'''ERROR! could not find file {checkpoint_path}'''      )

					# Load the model.
					__lowercase        =      torch.load(a__			,						map_location='''cpu'''      )

					# Load the config.
					__lowercase        =      checkpoint['''hyper_parameters''']['''cfg''']
					__lowercase        =      GPTaConfig()
					__lowercase        =      copy_config(config_hf=a__			,						config_megatron=a__      )
					__lowercase        =      ['''GPT2LMHeadModel''']

					__lowercase        =      checkpoint['''state_dict''']

					# Convert.
					print('''Converting'''      )
					__lowercase        =      convert_megatron_checkpoint(a__			,						a__      )

					# Print the structure of converted state dict.
					if args.print_checkpoint_structure:
										recursive_print(a__			,						a__      )

					__lowercase        =      '''GPTSw3Tokenizer'''

					# Store the config to file.
					print('''Saving config'''      )
					config_hf.save_pretrained(a__      )

					# Store the state_dict to file.
					__lowercase        =      os.path.join(a__			,						'''pytorch_model.bin'''      )
					print(f'''Saving checkpoint to "{output_checkpoint_file}"'''      )
					torch.save(a__			,						a__      )


if __name__ == "__main__":
							a__			:						Union[str, Any]							=			argparse.ArgumentParser()
							parser.add_argument(
							    """--checkpoint_path""",
							    type=str,
							    required=True,
							    help="""e.g. megatron_gpt--val_loss=2.42-step=38000-consumed_samples=54720000""",
							)
							parser.add_argument("""--save_path""", type=str, required=True, help="""e.g. /home/user/gpt-sw3/hf""")
							parser.add_argument("""--print-checkpoint-structure""", action="""store_true""")
							a__			:						Dict							=			parser.parse_args()
							main(_args)


