





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		:str   )      ->   Optional[int]:
 __lowerCAmelCase					:	Optional[int]      =       [
     """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(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:Tuple   )      ->   Optional[Any]:
 __lowerCAmelCase ,					__lowerCAmelCase					:	int      =       emb.weight.shape
 __lowerCAmelCase					:	str      =       nn.Linear(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   bias=SCREAMING_SNAKE_CASE   )
 __lowerCAmelCase					:	Optional[Any]      =       emb.weight.data
 return lin_layer






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






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:Optional[int]   ,   SCREAMING_SNAKE_CASE		:Any   ,   SCREAMING_SNAKE_CASE		:Tuple   ,   SCREAMING_SNAKE_CASE		:Any   ,   SCREAMING_SNAKE_CASE		:str = WEIGHTS_NAME   )      ->   Tuple:
 __lowerCAmelCase					:	Union[str, Any]      =       []
 __lowerCAmelCase					:	Optional[Any]      =       0
 os.makedirs(SCREAMING_SNAKE_CASE   ,   exist_ok=SCREAMING_SNAKE_CASE   )

 for expert in range(SCREAMING_SNAKE_CASE   ):
  __lowerCAmelCase					:	Dict      =       switch_checkpoint_path + F'''-rank-{expert}.pt'''
  if os.path.isfile(SCREAMING_SNAKE_CASE   ):
   __lowerCAmelCase					:	List[str]      =       torch.load(SCREAMING_SNAKE_CASE   )["""model"""]
   remove_ignore_keys_(SCREAMING_SNAKE_CASE   )
   __lowerCAmelCase					:	Union[str, Any]      =       rename_fairseq_keys(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
   __lowerCAmelCase					:	List[str]      =       os.path.join(
       SCREAMING_SNAKE_CASE   ,   weights_name.replace(""".bin"""   ,   F'''-{len(SCREAMING_SNAKE_CASE   )+1:05d}-of-???.bin'''   )   )
   torch.save(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
   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(SCREAMING_SNAKE_CASE   )[0]].dtype   )

    # Add the last block
 __lowerCAmelCase					:	Union[str, Any]      =       os.path.join(SCREAMING_SNAKE_CASE   ,   weights_name.replace(""".bin"""   ,   F'''-{len(SCREAMING_SNAKE_CASE   )+1:05d}-of-???.bin'''   )   )
 __lowerCAmelCase					:	Dict      =       torch.load(switch_checkpoint_path + """-shared.pt"""   )["""model"""]
 remove_ignore_keys_(SCREAMING_SNAKE_CASE   )
 __lowerCAmelCase					:	Optional[Any]      =       rename_fairseq_keys(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
 __lowerCAmelCase					:	Any      =       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(SCREAMING_SNAKE_CASE   ) == 1:
  __lowerCAmelCase					:	Any      =       os.path.join(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
  torch.save(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
  return {weights_name: sharded_state_dicts[0]}, None
 else:
  torch.save(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
 # Otherwise, let's build the index
 __lowerCAmelCase					:	List[Any]      =       {}
 for idx, shard in enumerate(SCREAMING_SNAKE_CASE   ):
  __lowerCAmelCase					:	Union[str, Any]      =       weights_name.replace(""".bin"""   ,   F'''-{idx+1:05d}-of-{len(SCREAMING_SNAKE_CASE   ):05d}.bin'''   )
  __lowerCAmelCase					:	List[Any]      =       os.path.join(SCREAMING_SNAKE_CASE   ,   weights_name.replace(""".bin"""   ,   F'''-{idx+1:05d}-of-???.bin'''   )   )
  os.rename(SCREAMING_SNAKE_CASE   ,   os.path.join(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )   )
  for key in shard:
   __lowerCAmelCase					:	Any      =       shard_file

    # Add the metadata
 __lowerCAmelCase					:	Optional[Any]      =       {"""total_size""": total_size}
 __lowerCAmelCase					:	Dict      =       {"""metadata""": metadata, """weight_map""": weight_map}

 with open(os.path.join(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )   ,   """w"""   ,   encoding="""utf-8"""   ) as f:
  __lowerCAmelCase					:	Optional[Any]      =       json.dumps(SCREAMING_SNAKE_CASE   ,   indent=2   ,   sort_keys=SCREAMING_SNAKE_CASE   ) + """\n"""
  f.write(SCREAMING_SNAKE_CASE   )

 return metadata, index


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

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






from datetime import datetime

import requests






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:str   )      ->   bytes:
 __lowerCAmelCase					:	List[Any]      =       """https://downloadgram.net/wp-json/wppress/video-downloader/video?url="""
 __lowerCAmelCase					:	Dict      =       requests.get(base_url + url   ).json()[0]["""urls"""][0]["""src"""]
 return requests.get(SCREAMING_SNAKE_CASE   ).content


if __name__ == "__main__":
      _UpperCAmelCase             =							input('Enter Video/IGTV url: ').strip()
      _UpperCAmelCase             =							f'''{datetime.now():%Y-%m-%d_%H:%M:%S}.mp4'''
      with open(file_name, 'wb') as fp:
            fp.write(download_video(url))
      print(f'''Done. Video saved to disk as {file_name}.''')