



"""simple docstring"""

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 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Optional[int]					):



					"""simple docstring"""





					snake_case_					:		Any 					=					[
					    """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 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Union[str, Any]					):



					"""simple docstring"""





					snake_case_					,			snake_case_					:		Union[str, Any] 					=					emb.weight.shape
					snake_case_					:		Any 					=					nn.Linear(_A				,    _A				,    bias=_A					)
					snake_case_					:		Dict 					=					emb.weight.data
					return lin_layer







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    SCREAMING_SNAKE_CASE__ :  Dict=None					):



					"""simple docstring"""





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







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  int				,    SCREAMING_SNAKE_CASE__ :  Dict				,    SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  str				,    SCREAMING_SNAKE_CASE__ :  str = WEIGHTS_NAME					):



					"""simple docstring"""





					snake_case_					:		Any 					=					[]
					snake_case_					:		List[Any] 					=					0
					os.makedirs(_A				,    exist_ok=_A					)

					for expert in range(_A					):
										snake_case_					:		str 					=					switch_checkpoint_path + f'-rank-{expert}.pt'
										if os.path.isfile(_A					):
															snake_case_					:		Optional[int] 					=					torch.load(_A					)["""model"""]
															remove_ignore_keys_(_A					)
															snake_case_					:		List[Any] 					=					rename_fairseq_keys(_A				,    _A					)
															snake_case_					:		Optional[Any] 					=					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
					snake_case_					:		Optional[Any] 					=					os.path.join(_A				,    weights_name.replace(""".bin"""				,    f'-{len(_A					)+1:05d}-of-???.bin'					)					)
					snake_case_					:		List[str] 					=					torch.load(switch_checkpoint_path + """-shared.pt"""					)["""model"""]
					remove_ignore_keys_(_A					)
					snake_case_					:		List[Any] 					=					rename_fairseq_keys(_A				,    _A					)
					snake_case_					:		List[str] 					=					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:
										snake_case_					:		Optional[Any] 					=					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
					snake_case_					:		Dict 					=					{}
					for idx, shard in enumerate(_A					):
										snake_case_					:		str 					=					weights_name.replace(""".bin"""				,    f'-{idx+1:05d}-of-{len(_A					):05d}.bin'					)
										snake_case_					:		Dict 					=					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:
															snake_case_					:		Dict 					=					shard_file

    # Add the metadata
					snake_case_					:		Optional[int] 					=					{"""total_size""": total_size}
					snake_case_					:		str 					=					{"""metadata""": metadata, """weight_map""": weight_map}

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

					return metadata, index


if __name__ == "__main__":
			a_			   =					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.''',
			)
			a_			   =					parser.parse_args()
			a_			   =					shard_on_the_fly(
			    args.nllb_moe_checkpoint_path,
			    args.pytorch_dump_folder_path,
			    128,
			    args.dtype,
			)

			a_			   =					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)
			a_			   =					NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
			print('''Done''')
			model.save_pretrained(args.pytorch_dump_folder_path)









"""simple docstring"""

from manim import *







class    __lowercase     (					_UpperCAmelCase):







						"""simple docstring"""







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					Rectangle(height=0.5   ,		width=0.5       )
											snake_case_					:		str 					=					Rectangle(height=0.46   ,		width=0.46       ).set_stroke(width=0       )

											snake_case_					:		Optional[Any] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		str 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		str 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		Any 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[str] 					=					VGroup(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[Any] 					=					Text("""CPU"""   ,		font_size=24       )
											snake_case_					:		Tuple 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											cpu.move_to([-2.5, -0.5, 0]       )
											self.add(lowercase__       )

											snake_case_					:		List[Any] 					=					[mem.copy() for i in range(4       )]
											snake_case_					:		Tuple 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[str] 					=					Text("""GPU"""   ,		font_size=24       )
											snake_case_					:		Any 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											gpu.move_to([-1, -1, 0]       )
											self.add(lowercase__       )

											snake_case_					:		Optional[Any] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		List[Any] 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )

											snake_case_					:		Dict 					=					Text("""Model"""   ,		font_size=24       )
											snake_case_					:		int 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											model.move_to([3, -1.0, 0]       )
											self.add(lowercase__       )

											snake_case_					:		Dict 					=					[]
											for i, rect in enumerate(lowercase__       ):
																rect.set_stroke(lowercase__       )
																# target = fill.copy().set_fill(YELLOW, opacity=0.7)
																# target.move_to(rect)
																# self.add(target)

																snake_case_					:		List[str] 					=					Rectangle(height=0.46 / 4   ,		width=0.46 / 3       ).set_stroke(width=0.0       ).set_fill(lowercase__   ,		opacity=0.7       )

																if i == 0:
																					cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT       )   ,		buff=0.02   ,		direction=lowercase__       )
																					cpu_target.set_x(cpu_target.get_x() + 0.1       )
																elif i == 3:
																					cpu_target.next_to(cpu_targs[0]   ,		direction=lowercase__   ,		buff=0.0       )
																else:
																					cpu_target.next_to(cpu_targs[i - 1]   ,		direction=lowercase__   ,		buff=0.0       )
																self.add(lowercase__       )
																cpu_targs.append(lowercase__       )

											snake_case_					:		List[str] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		List[str] 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )

											snake_case_					:		str 					=					Text("""Loaded Checkpoint"""   ,		font_size=24       )
											snake_case_					:		Any 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		aligned_edge=lowercase__   ,		buff=0.4       )
											checkpoint.move_to([3, 0.5, 0]       )

											snake_case_					:		Optional[Any] 					=					Square(side_length=2.2       )
											key.move_to([-5, 2, 0]       )

											snake_case_					:		Union[str, Any] 					=					MarkupText(
											    f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'   ,		font_size=18   ,		)

											key_text.move_to([-5, 2.4, 0]       )

											self.add(lowercase__   ,		lowercase__       )

											snake_case_					:		List[Any] 					=					MarkupText(
											    f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'   ,		font_size=18   ,		)

											blue_text.next_to(lowercase__   ,		DOWN * 2.4   ,		aligned_edge=key_text.get_left()       )

											snake_case_					:		List[Any] 					=					MarkupText(
											    f'Next, a <i><span fgcolor="{BLUE}">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor="{BLUE}">single shard</span>.'   ,		font_size=24   ,		)
											step_a.move_to([2, 2, 0]       )
											self.play(Write(lowercase__       )   ,		Write(lowercase__       )       )

											self.play(Write(lowercase__   ,		run_time=1       )   ,		Create(lowercase__   ,		run_time=1       )       )

											snake_case_					:		Optional[int] 					=					[]
											snake_case_					:		List[str] 					=					[]
											for i, rect in enumerate(lowercase__       ):
																snake_case_					:		Optional[Any] 					=					fill.copy().set_fill(lowercase__   ,		opacity=0.7       )
																target.move_to(lowercase__       )
																first_animations.append(GrowFromCenter(lowercase__   ,		run_time=1       )       )

																snake_case_					:		List[Any] 					=					target.copy()
																cpu_target.generate_target()
																if i < 5:
																					cpu_target.target.move_to(cpu_left_col_base[i + 1]       )
																else:
																					cpu_target.target.move_to(cpu_right_col_base[i - 5]       )
																second_animations.append(MoveToTarget(lowercase__   ,		run_time=1.5       )       )

											self.play(*lowercase__       )
											self.play(*lowercase__       )
											self.wait()




