






"""simple docstring"""



import argparse

import torch
from clip import load

from transformers import CLIPConfig, CLIPModel
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    List[str]						,   __snake_case     :    Any       ):





  '''simple docstring'''





  lowercase       ,	lowercase       ,	lowercase				       =			pt_attn_layer.in_proj_weight.chunk(3						,   dim=0       )
  lowercase       ,	lowercase       ,	lowercase				       =			pt_attn_layer.in_proj_bias.chunk(3						,   dim=0       )

  lowercase				       =			pt_attn_layer.out_proj.weight
  lowercase				       =			pt_attn_layer.out_proj.bias

  lowercase				       =			q_proj
  lowercase				       =			q_proj_bias

  lowercase				       =			k_proj
  lowercase				       =			k_proj_bias

  lowercase				       =			v_proj
  lowercase				       =			v_proj_bias

  lowercase				       =			out_proj_weights
  lowercase				       =			out_proj_bias
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Dict						,   __snake_case     :    str       ):





  '''simple docstring'''





  copy_linear(hf_mlp.fca						,   pt_mlp.c_fc       )
  copy_linear(hf_mlp.fca						,   pt_mlp.c_proj       )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Dict						,   __snake_case     :    Tuple       ):





  '''simple docstring'''





  lowercase				       =			pt_linear.weight
  lowercase				       =			pt_linear.bias
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Union[str, Any]						,   __snake_case     :    Optional[int]       ):





  '''simple docstring'''





  copy_linear(hf_layer.layer_norma						,   pt_layer.ln_a       )
  copy_linear(hf_layer.layer_norma						,   pt_layer.ln_a       )

  # copy MLP
  copy_mlp(hf_layer.mlp						,   pt_layer.mlp       )

  # copy attn
  copy_attn_layer(hf_layer.self_attn						,   pt_layer.attn       )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    int						,   __snake_case     :    Union[str, Any]       ):





  '''simple docstring'''





  for hf_layer, pt_layer in zip(__snake_case						,   __snake_case       ):
    copy_layer(__snake_case						,   __snake_case       )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Optional[int]						,   __snake_case     :    Any       ):





  '''simple docstring'''





  lowercase				       =			pt_model.token_embedding.weight
  lowercase				       =			pt_model.positional_embedding

  # copy layer norm
  copy_linear(hf_encoder.final_layer_norm						,   pt_model.ln_final       )

  # copy hidden layers
  copy_layers(hf_encoder.encoder.layers						,   pt_model.transformer.resblocks       )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Optional[Any]						,   __snake_case     :    List[str]       ):





  '''simple docstring'''





  lowercase				       =			pt_model.text_projection.data.T

  # copy text encoder
  copy_encoder(hf_model.text_model						,   __snake_case       )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Dict						,   __snake_case     :    List[Any]       ):





  '''simple docstring'''





  lowercase				       =			pt_model.visual.proj.data.T

  # copy layer norms
  copy_linear(hf_model.vision_model.pre_layrnorm						,   pt_model.visual.ln_pre       )
  copy_linear(hf_model.vision_model.post_layernorm						,   pt_model.visual.ln_post       )

  # copy embeds
  lowercase				       =			pt_model.visual.conva.weight.data
  lowercase				       =			pt_model.visual.class_embedding
  lowercase				       =			pt_model.visual.positional_embedding.data

  # copy encoder
  copy_layers(hf_model.vision_model.encoder.layers						,   pt_model.visual.transformer.resblocks       )
@torch.no_grad()
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Any						,   __snake_case     :    Any						,   __snake_case     :    int=None       ):





  '''simple docstring'''





  if config_path is not None:
    lowercase				       =			CLIPConfig.from_pretrained(__snake_case       )
  else:
    lowercase				       =			CLIPConfig(projection_dim=5_12						,   text_config={}						,   vision_config={}       )

  lowercase				       =			CLIPModel(__snake_case       ).eval()

  lowercase       ,	lowercase				       =			load(__snake_case						,   device='cpu'						,   jit=__snake_case       )
  lowercase				       =			pt_model.eval()

  copy_text_model_and_projection(__snake_case						,   __snake_case       )
  copy_vison_model_and_projection(__snake_case						,   __snake_case       )
  lowercase				       =			pt_model.logit_scale

  lowercase				       =			torch.arange(0						,   77       ).unsqueeze(0       )
  lowercase				       =			torch.randn(1						,   3						,   2_24						,   2_24       )

  lowercase				       =			hf_model(input_ids=__snake_case						,   pixel_values=__snake_case						,   return_dict=__snake_case       )
  lowercase				       =			hf_outputs.logits_per_image
  lowercase				       =			hf_outputs.logits_per_text
  lowercase       ,	lowercase				       =			pt_model(__snake_case						,   __snake_case       )

  assert torch.allclose(__snake_case						,   __snake_case						,   atol=1e-3       )
  assert torch.allclose(__snake_case						,   __snake_case						,   atol=1e-3       )

  hf_model.save_pretrained(__snake_case       )


if __name__ == "__main__":
    _UpperCamelCase						:    int		=						argparse.ArgumentParser()
    parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
    parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
    parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
    _UpperCamelCase						:    Dict		=						parser.parse_args()

    convert_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)








import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig

def       lowerCAmelCase_      (	lowerCamelCase:							str  ,					lowerCamelCase:							List[str]  ,					lowerCamelCase:							List[Any]=0    )       ->      Dict:


      '''simple docstring'''




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

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

def       lowerCAmelCase_      (	lowerCamelCase:							Tuple  ,					lowerCamelCase:							Dict  ,					lowerCamelCase:							List[str]  ,					lowerCamelCase:							Any    )       ->      Optional[int]:


      '''simple docstring'''




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

def       lowerCAmelCase_      (	lowerCamelCase:							List[str]  ,					lowerCamelCase:							int    )       ->      int:


      '''simple docstring'''




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

      lowerCamelCase__   	=					sd_megatron['model.language_model.embedding.word_embeddings.weight'][:vocab_size, :]
      lowerCamelCase__   	=					{
          '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'],
      }

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

            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.input_layernorm.weight"""]
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.input_layernorm.bias"""]

            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.self_attention.query_key_value.weight"""]
            lowerCamelCase__   	=					fix_query_key_value_ordering(lowerCamelCase  ,					3  ,					lowerCamelCase  ,					lowerCamelCase    )
            lowerCamelCase__   	=					vala.transpose(0  ,					1    ).contiguous()

            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.self_attention.query_key_value.bias"""]
            lowerCamelCase__   	=					fix_query_key_value_ordering(lowerCamelCase  ,					3  ,					lowerCamelCase  ,					lowerCamelCase    )
            lowerCamelCase__   	=					vala

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

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

      return sd_hf

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[int]  ,					lowerCamelCase:							List[Any]    )       ->      Tuple:


      '''simple docstring'''




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

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

      return config_hf

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[Any]    )       ->      Dict:


      '''simple docstring'''




      print(lowerCamelCase    )

      lowerCamelCase__   	=					args.checkpoint_path
      lowerCamelCase__   	=					args.save_path
      if isfile(lowerCamelCase    ):
            raise FileNotFoundError(f"""ERROR! could not find file {checkpoint_path}"""    )

      # Load the model.
      lowerCamelCase__   	=					torch.load(lowerCamelCase  ,					map_location='cpu'    )

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

      lowerCamelCase__   	=					checkpoint['state_dict']

      # Convert.
      print('Converting'    )
      lowerCamelCase__   	=					convert_megatron_checkpoint(lowerCamelCase  ,					lowerCamelCase    )

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

      lowerCamelCase__   	=					'GPTSw3Tokenizer'

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

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


if __name__ == "__main__":
      __SCREAMING_SNAKE_CASE     : Optional[int]      =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""")
      __SCREAMING_SNAKE_CASE     : Tuple      =parser.parse_args()
      main(_args)