



import argparse
import json
import os

import torch
from torch import nn

from transformers import NllbMoeConfig, NllbMoeModel
from transformers.modeling_utils import dtype_byte_size
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
def  UpperCamelCase						(       _A			:				Tuple    )->       Dict:





						"""simple docstring"""





						A__							=       [
						    "encoder.version",
						    "decoder.version",
						    "model.encoder.version",
						    "model.decoder.version",
						    "decoder.output_projection.weight",
						    "_float_tensor",
						    "encoder.embed_positions._float_tensor",
						    "decoder.embed_positions._float_tensor",
						]
						for k in ignore_keys:
												state_dict.pop(_A			,							_A    )
def  UpperCamelCase						(       _A			:				int    )->       Optional[Any]:





						"""simple docstring"""





						A__						,       A__							=       emb.weight.shape
						A__							=       nn.Linear(_A			,							_A			,							bias=_A    )
						A__							=       emb.weight.data
						return lin_layer
def  UpperCamelCase						(       _A			:				str			,							_A			:				Optional[Any]=None    )->       str:





						"""simple docstring"""





						A__							=       {}
						for old_key in state_dict.keys():
												A__							=       old_key
												if "moe_layer.experts." in key:
																		if expert_idx is not None:
																								A__							=       key.replace("moe_layer.experts.0"			,							f"""ffn.experts.expert_{expert_idx}"""    )
																		else:
																								A__							=       key.replace("moe_layer.experts."			,							"ffn.experts.expert_"    )
												if "gate" in key:
																		A__							=       key.replace(".moe_layer.gate.wg"			,							".ffn.router.classifier"    )
												if "fc2" and "experts" not in key:
																		A__							=       key.replace(".fc2."			,							".ffn.fc2."    )
												if "fc1" and "experts" not in key:
																		A__							=       key.replace(".fc1."			,							".ffn.fc1."    )
												if ".encoder_attn." in key:
																		A__							=       key.replace(".encoder_attn."			,							".cross_attention."    )
												if "encoder_attn_layer_norm" in key:
																		A__							=       key.replace("encoder_attn_layer_norm"			,							"cross_attention_layer_norm"    )
												if "final_layer_norm" in key:
																		A__							=       key.replace("final_layer_norm"			,							"ff_layer_norm"    )
												A__							=       state_dict[old_key]
						return new_dict
def  UpperCamelCase						(       _A			:				Tuple			,							_A			:				Tuple			,							_A			:				int			,							_A			:				str			,							_A			:				str = WEIGHTS_NAME    )->       List[str]:





						"""simple docstring"""





						A__							=       []
						A__							=       0
						os.makedirs(_A			,							exist_ok=_A    )

						for expert in range(_A    ):
												A__							=       switch_checkpoint_path + f"""-rank-{expert}.pt"""
												if os.path.isfile(_A    ):
																		A__							=       torch.load(_A    )["model"]
																		remove_ignore_keys_(_A    )
																		A__							=       rename_fairseq_keys(_A			,							_A    )
																		A__							=       os.path.join(
																		    _A			,							weights_name.replace(".bin"			,							f"""-{len(_A    )+1:05d}-of-???.bin"""    )    )
																		torch.save(_A			,							_A    )
																		sharded_state_dicts.append(expert_state.keys()    )
																		total_size += sum([value.numel() for key, value in expert_state.items()]    ) * dtype_byte_size(
																		    expert_state[list(_A    )[0]].dtype    )

    # Add the last block
						A__							=       os.path.join(_A			,							weights_name.replace(".bin"			,							f"""-{len(_A    )+1:05d}-of-???.bin"""    )    )
						A__							=       torch.load(switch_checkpoint_path + "-shared.pt"    )["model"]
						remove_ignore_keys_(_A    )
						A__							=       rename_fairseq_keys(_A			,							_A    )
						A__							=       shared_weights["decoder.embed_tokens.weight"]
						sharded_state_dicts.append(shared_weights.keys()    )

						# If we only have the shared weights (dummy model/experts saved on the same file)
						if len(_A    ) == 1:
												A__							=       os.path.join(_A			,							_A    )
												torch.save(_A			,							_A    )
												return {weights_name: sharded_state_dicts[0]}, None
						else:
												torch.save(_A			,							_A    )
						# Otherwise, let's build the index
						A__							=       {}
						for idx, shard in enumerate(_A    ):
												A__							=       weights_name.replace(".bin"			,							f"""-{idx+1:05d}-of-{len(_A    ):05d}.bin"""    )
												A__							=       os.path.join(_A			,							weights_name.replace(".bin"			,							f"""-{idx+1:05d}-of-???.bin"""    )    )
												os.rename(_A			,							os.path.join(_A			,							_A    )    )
												for key in shard:
																		A__							=       shard_file

    # Add the metadata
						A__							=       {"total_size": total_size}
						A__							=       {"metadata": metadata, "weight_map": weight_map}

						with open(os.path.join(_A			,							_A    )			,							"w"			,							encoding="utf-8"    ) as f:
												A__							=       json.dumps(_A			,							indent=2			,							sort_keys=_A    ) + "\n"
												f.write(_A    )

						return metadata, index


if __name__ == "__main__":
						UpperCAmelCase_							:      str      =  argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--nllb_moe_checkpoint_path",
						    default="/home/arthur_huggingface_co/fairseq/weights/checkpoints/model_moe_54b/checkpoint_2_300000",
						    type=str,
						    required=False,
						    help="Path to a directory containing a folder per layer. Follows the original Google format.",
						)
						parser.add_argument("--dtype", default="float32", type=str, required=False, help="dtype of the saved model")
						parser.add_argument(
						    "--pytorch_dump_folder_path",
						    default="/home/arthur_huggingface_co/fairseq/weights/checkpoints/hf-converted-moe-54b",
						    type=str,
						    required=False,
						    help="Path to the output pytorch model.",
						)
						UpperCAmelCase_							:      Union[str, Any]      =  parser.parse_args()
						UpperCAmelCase_   ,	UpperCAmelCase_							:      Optional[int]      =  shard_on_the_fly(
						    args.nllb_moe_checkpoint_path,
						    args.pytorch_dump_folder_path,
						    128,
						    args.dtype,
						)

						UpperCAmelCase_							:      Any      =  NllbMoeConfig.from_pretrained(
						    "facebook/nllb-200-3.3B", encoder_sparse_step=4, decoder_sparse_step=4, num_experts=128
						)
						config.save_pretrained(args.pytorch_dump_folder_path)
						UpperCAmelCase_							:      Tuple      =  NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
						print("Done")
						model.save_pretrained(args.pytorch_dump_folder_path)










import argparse
import struct
import unittest






class      UpperCamelCase    :



	def __init__(       self			,	UpperCAmelCase__		):
							A__							=       data

							# Initialize hash values
							A__							=       [
							    0x6A_09E_667,
							    0xBB_67A_E85,
							    0x3C_6EF_372,
							    0xA5_4FF_53A,
							    0x51_0E5_27F,
							    0x9B_056_88C,
							    0x1F_83D_9AB,
							    0x5B_E0C_D19,
							]

							# Initialize round constants
							A__							=       [
							    0x42_8A2_F98,
							    0x71_374_491,
							    0xB5_C0F_BCF,
							    0xE9_B5D_BA5,
							    0x39_56C_25B,
							    0x59_F11_1F1,
							    0x92_3F8_2A4,
							    0xAB_1C5_ED5,
							    0xD8_07A_A98,
							    0x12_835_B01,
							    0x24_318_5BE,
							    0x55_0C7_DC3,
							    0x72_BE5_D74,
							    0x80_DEB_1FE,
							    0x9B_DC0_6A7,
							    0xC1_9BF_174,
							    0xE4_9B6_9C1,
							    0xEF_BE4_786,
							    0x0F_C19_DC6,
							    0x24_0CA_1CC,
							    0x2D_E92_C6F,
							    0x4A_748_4AA,
							    0x5C_B0A_9DC,
							    0x76_F98_8DA,
							    0x98_3E5_152,
							    0xA8_31C_66D,
							    0xB0_032_7C8,
							    0xBF_597_FC7,
							    0xC6_E00_BF3,
							    0xD5_A79_147,
							    0x06_CA6_351,
							    0x14_292_967,
							    0x27_B70_A85,
							    0x2E_1B2_138,
							    0x4D_2C6_DFC,
							    0x53_380_D13,
							    0x65_0A7_354,
							    0x76_6A0_ABB,
							    0x81_C2C_92E,
							    0x92_722_C85,
							    0xA2_BFE_8A1,
							    0xA8_1A6_64B,
							    0xC2_4B8_B70,
							    0xC7_6C5_1A3,
							    0xD1_92E_819,
							    0xD6_990_624,
							    0xF4_0E3_585,
							    0x10_6AA_070,
							    0x19_A4C_116,
							    0x1E_376_C08,
							    0x27_487_74C,
							    0x34_B0B_CB5,
							    0x39_1C0_CB3,
							    0x4E_D8A_A4A,
							    0x5B_9CC_A4F,
							    0x68_2E6_FF3,
							    0x74_8F8_2EE,
							    0x78_A56_36F,
							    0x84_C87_814,
							    0x8C_C70_208,
							    0x90_BEF_FFA,
							    0xA4_506_CEB,
							    0xBE_F9A_3F7,
							    0xC6_717_8F2,
							]

							A__							=       self.preprocessing(self.data		)
							self.final_hash()



	@staticmethod
	def 			__A							(       UpperCAmelCase__		):
							A__							=       b"\x80" + (b"\x00" * (63 - (len(UpperCAmelCase__		) + 8) % 64))
							A__							=       struct.pack(">Q"			,	(len(UpperCAmelCase__		) * 8)		)
							return data + padding + big_endian_integer



	def 			__A							(       self		):
							# Convert into blocks of 64 bytes
							A__							=       [
							    self.preprocessed_data[x : x + 64]
							    for x in range(0			,	len(self.preprocessed_data		)			,	64		)
							]

							for block in self.blocks:
													# Convert the given block into a list of 4 byte integers
													A__							=       list(struct.unpack(">16L"			,	UpperCAmelCase__		)		)
													# add 48 0-ed integers
													words += [0] * 48

													A__						,       A__						,       A__						,       A__						,       A__						,       A__						,       A__						,       A__							=       self.hashes

													for index in range(0			,	64		):
																			if index > 15:
																									# modify the zero-ed indexes at the end of the array
																									A__							=       (
																									    self.ror(words[index - 15]			,	7		)
																									    ^ self.ror(words[index - 15]			,	18		)
																									    ^ (words[index - 15] >> 3)
																									)
																									A__							=       (
																									    self.ror(words[index - 2]			,	17		)
																									    ^ self.ror(words[index - 2]			,	19		)
																									    ^ (words[index - 2] >> 10)
																									)

																									A__							=       (
																									    words[index - 16] + sa + words[index - 7] + sa
																									) % 0x100_000_000

																			# Compression
																			A__							=       self.ror(UpperCAmelCase__			,	6		) ^ self.ror(UpperCAmelCase__			,	11		) ^ self.ror(UpperCAmelCase__			,	25		)
																			A__							=       (e & f) ^ ((~e & 0xFF_FFF_FFF) & g)
																			A__							=       (
																			    h + sa + ch + self.round_constants[index] + words[index]
																			) % 0x100_000_000
																			A__							=       self.ror(UpperCAmelCase__			,	2		) ^ self.ror(UpperCAmelCase__			,	13		) ^ self.ror(UpperCAmelCase__			,	22		)
																			A__							=       (a & b) ^ (a & c) ^ (b & c)
																			A__							=       (sa + maj) % 0x100_000_000

																			A__						,       A__						,       A__						,       A__						,       A__						,       A__						,       A__						,       A__							=       (
																			    g,
																			    f,
																			    e,
																			    ((d + tempa) % 0x100_000_000),
																			    c,
																			    b,
																			    a,
																			    ((tempa + tempa) % 0x100_000_000),
																			)

													A__							=       [a, b, c, d, e, f, g, h]

													# Modify final values
													A__							=       [
													    ((element + mutated_hash_values[index]) % 0x100_000_000)
													    for index, element in enumerate(self.hashes		)
													]

							A__							=       "".join([hex(UpperCAmelCase__		)[2:].zfill(8		) for value in self.hashes]		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							return 0xFF_FFF_FFF & (value << (32 - rotations)) | (value >> rotations)






class      UpperCamelCase    (						unittest.TestCase     ):



	def 			__A							(       self		):
							import hashlib

							A__							=       bytes("Test String"			,	"utf-8"		)
							self.assertEqual(SHAaaa(UpperCAmelCase__		).hash			,	hashlib.shaaaa(UpperCAmelCase__		).hexdigest()		)





def  UpperCamelCase						(       )->       None:





						"""simple docstring"""





						import doctest

						doctest.testmod()

						A__							=       argparse.ArgumentParser()
						parser.add_argument(
						    "-s"			,							"--string"			,							dest="input_string"			,							default="Hello World!! Welcome to Cryptography"			,							help="Hash the string"			,							)
						parser.add_argument(
						    "-f"			,							"--file"			,							dest="input_file"			,							help="Hash contents of a file"    )

						A__							=       parser.parse_args()

						A__							=       args.input_string

						# hash input should be a bytestring
						if args.input_file:
												with open(args.input_file			,							"rb"    ) as f:
																		A__							=       f.read()
						else:
												A__							=       bytes(_A			,							"utf-8"    )

						print(SHAaaa(_A    ).hash    )


if __name__ == "__main__":
						main()





