





'''simple docstring'''



import argparse

import torch
from clip import load

from transformers import CLIPConfig, CLIPModel




def       _lowerCAmelCase  ( lowerCamelCase_	:     Dict						,				lowerCamelCase_	:     Union[str, Any] ):
			__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       _lowerCAmelCase  ( lowerCamelCase_	:     Optional[int]						,				lowerCamelCase_	:     Dict ):
			copy_linear(hf_mlp.fca						,				pt_mlp.c_fc )
			copy_linear(hf_mlp.fca						,				pt_mlp.c_proj )




def       _lowerCAmelCase  ( lowerCamelCase_	:     Optional[Any]						,				lowerCamelCase_	:     str ):
			__lowercase   =							pt_linear.weight
			__lowercase   =							pt_linear.bias




def       _lowerCAmelCase  ( lowerCamelCase_	:     List[Any]						,				lowerCamelCase_	:     int ):
			# copy layer norms
			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       _lowerCAmelCase  ( lowerCamelCase_	:     Optional[int]						,				lowerCamelCase_	:     List[str] ):
			for hf_layer, pt_layer in zip(lowerCamelCase_						,				lowerCamelCase_ ):
						copy_layer(lowerCamelCase_						,				lowerCamelCase_ )




def       _lowerCAmelCase  ( lowerCamelCase_	:     Tuple						,				lowerCamelCase_	:     Dict ):
			# copy  embeds
			__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       _lowerCAmelCase  ( lowerCamelCase_	:     Tuple						,				lowerCamelCase_	:     Dict ):
			# copy projection
			__lowercase   =							pt_model.text_projection.data.T

			# copy text encoder
			copy_encoder(hf_model.text_model						,				lowerCamelCase_ )




def       _lowerCAmelCase  ( lowerCamelCase_	:     str						,				lowerCamelCase_	:     Dict ):
			# copy projection
			__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       _lowerCAmelCase  ( lowerCamelCase_	:     int						,				lowerCamelCase_	:     Any						,				lowerCamelCase_	:     str=None ):
			if config_path is not None:
						__lowercase   =							CLIPConfig.from_pretrained(lowerCamelCase_ )
			else:
						__lowercase   =							CLIPConfig(projection_dim=5_1_2						,				text_config={}						,				vision_config={} )

			__lowercase   =							CLIPModel(lowerCamelCase_ ).eval()

			__lowercase  ,						__lowercase   =							load(lowerCamelCase_						,				device='''cpu'''						,				jit=lowerCamelCase_ )
			__lowercase   =							pt_model.eval()

			copy_text_model_and_projection(lowerCamelCase_						,				lowerCamelCase_ )
			copy_vison_model_and_projection(lowerCamelCase_						,				lowerCamelCase_ )
			__lowercase   =							pt_model.logit_scale

			__lowercase   =							torch.arange(0						,				7_7 ).unsqueeze(0 )
			__lowercase   =							torch.randn(1						,				3						,				2_2_4						,				2_2_4 )

			__lowercase   =							hf_model(input_ids=lowerCamelCase_						,				pixel_values=lowerCamelCase_						,				return_dict=lowerCamelCase_ )
			__lowercase   =							hf_outputs.logits_per_image
			__lowercase   =							hf_outputs.logits_per_text
			__lowercase  ,						__lowercase   =							pt_model(lowerCamelCase_						,				lowerCamelCase_ )

			assert torch.allclose(lowerCamelCase_						,				lowerCamelCase_						,				atol=1E-3 )
			assert torch.allclose(lowerCamelCase_						,				lowerCamelCase_						,				atol=1E-3 )

			hf_model.save_pretrained(lowerCamelCase_ )


if __name__ == "__main__":
	_SCREAMING_SNAKE_CASE    =			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''')
	_SCREAMING_SNAKE_CASE    =			parser.parse_args()

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







from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


_A  					= logging.get_logger(__name__)

_A  					= {
    "google/bigbird-pegasus-large-arxiv": (
        "https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json"
    ),
    "google/bigbird-pegasus-large-pubmed": (
        "https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json"
    ),
    "google/bigbird-pegasus-large-bigpatent": (
        "https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json"
    ),
    # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
}







class      lowerCamelCase		(  A_  ):
     UpperCAmelCase__				:							List[Any]           =		"bigbird_pegasus"
     UpperCAmelCase__				:							Tuple           =		["past_key_values"]
     UpperCAmelCase__				:							Optional[Any]           =		{
         "num_attention_heads": "encoder_attention_heads",
         "hidden_size": "d_model",
         "attention_probs_dropout_prob": "attention_dropout",
     }





     def __init__(self      :    Any				,    _A      :    List[str]=9_6_1_0_3				,    _A      :    List[Any]=4_0_9_6				,    _A      :    str=1_6				,    _A      :    Tuple=4_0_9_6				,    _A      :    List[str]=1_6				,    _A      :    Optional[int]=1_6				,    _A      :    List[Any]=4_0_9_6				,    _A      :    Dict=1_6				,    _A      :    Tuple=0.0				,    _A      :    List[str]=0.0				,    _A      :    Dict=True				,    _A      :    int=True				,    _A      :    List[Any]="gelu_new"				,    _A      :    int=1_0_2_4				,    _A      :    int=0.1				,    _A      :    List[Any]=0.0				,    _A      :    Optional[int]=0.0				,    _A      :    Any=0.02				,    _A      :    Dict=2				,    _A      :    List[Any]=0.0				,    _A      :    Any=True				,    _A      :    str=0				,    _A      :    Dict=2				,    _A      :    Any=1				,    _A      :    List[str]="block_sparse"				,    _A      :    int=6_4				,    _A      :    Optional[int]=3				,    _A      :    List[Any]=False				,    **_A      :    Tuple				,    )						->						Any:
         snake_case      =					vocab_size
         snake_case      =					max_position_embeddings
         snake_case      =					d_model
         snake_case      =					encoder_ffn_dim
         snake_case      =					encoder_layers
         snake_case      =					encoder_attention_heads
         snake_case      =					decoder_ffn_dim
         snake_case      =					decoder_layers
         snake_case      =					decoder_attention_heads
         snake_case      =					dropout
         snake_case      =					attention_dropout
         snake_case      =					activation_dropout
         snake_case      =					activation_function
         snake_case      =					init_std
         snake_case      =					encoder_layerdrop
         snake_case      =					decoder_layerdrop
         snake_case      =					classifier_dropout
         snake_case      =					use_cache
         snake_case      =					encoder_layers
         snake_case      =					scale_embedding  # scale factor will be sqrt(d_model) if True

         # extra config
         snake_case      =					attention_type
         snake_case      =					block_size
         snake_case      =					num_random_blocks
         snake_case      =					use_bias

         super().__init__(
             pad_token_id=_A				,    bos_token_id=_A				,    eos_token_id=_A				,    is_encoder_decoder=_A				,    decoder_start_token_id=_A				,    **_A				,    )







class      lowerCamelCase		(  A_  ):





     @property
     def 		UpperCAmelCase(self      :    Any   )						->						Mapping[str, Mapping[int, str]]:
         if self.task in ["default", "seq2seq-lm"]:
             snake_case      =					OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                 ]   )

             if self.use_past:
                 snake_case      =					{0: "batch"}
                 snake_case      =					{0: "batch", 1: "past_decoder_sequence + sequence"}
             else:
                 snake_case      =					{0: "batch", 1: "decoder_sequence"}
                 snake_case      =					{0: "batch", 1: "decoder_sequence"}

             if self.use_past:
                 self.fill_with_past_key_values_(_A				,    direction="inputs"   )
         elif self.task == "causal-lm":
             # TODO: figure this case out.
             snake_case      =					OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                 ]   )
             if self.use_past:
                 snake_case     ,							snake_case      =					self.num_layers
                 for i in range(_A   ):
                     snake_case      =					{0: "batch", 2: "past_sequence + sequence"}
                     snake_case      =					{0: "batch", 2: "past_sequence + sequence"}
         else:
             snake_case      =					OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                     ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
                     ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
                 ]   )

         return common_inputs





     @property
     def 		UpperCAmelCase(self      :    str   )						->						Mapping[str, Mapping[int, str]]:
         if self.task in ["default", "seq2seq-lm"]:
             snake_case      =					super().outputs
         else:
             snake_case      =					super(_A				,    self   ).outputs
             if self.use_past:
                 snake_case     ,							snake_case      =					self.num_layers
                 for i in range(_A   ):
                     snake_case      =					{0: "batch", 2: "past_sequence + sequence"}
                     snake_case      =					{0: "batch", 2: "past_sequence + sequence"}
         return common_outputs





     def 		UpperCAmelCase(self      :    Any				,    _A      :    PreTrainedTokenizer				,    _A      :    int = -1				,    _A      :    int = -1				,    _A      :    bool = False				,    _A      :    Optional[TensorType] = None				,    )						->						Mapping[str, Any]:
         snake_case      =					self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             _A				,    _A				,    _A				,    _A				,    _A   )

         # Generate decoder inputs
         snake_case      =					seq_length if not self.use_past else 1
         snake_case      =					self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             _A				,    _A				,    _A				,    _A				,    _A   )
         snake_case      =					{f'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
         snake_case      =					dict(**_A				,    **_A   )

         if self.use_past:
             if not is_torch_available():
                 raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."   )
             else:
                 import torch
             snake_case     ,							snake_case      =					common_inputs["input_ids"].shape
             snake_case      =					common_inputs["decoder_input_ids"].shape[1]
             snake_case     ,							snake_case      =					self.num_attention_heads
             snake_case      =					(
                 batch,
                 num_encoder_attention_heads,
                 encoder_seq_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             snake_case      =					decoder_seq_length + 3
             snake_case      =					(
                 batch,
                 num_decoder_attention_heads,
                 decoder_past_length,
                 self._config.hidden_size // num_decoder_attention_heads,
             )

             snake_case      =					torch.cat(
                 [common_inputs["decoder_attention_mask"], torch.ones(_A				,    _A   )]				,    dim=1   )

             snake_case      =					[]
             # If the number of encoder and decoder layers are present in the model configuration, both are considered
             snake_case     ,							snake_case      =					self.num_layers
             snake_case      =					min(_A				,    _A   )
             snake_case      =					max(_A				,    _A   ) - min_num_layers
             snake_case      =					"encoder" if num_encoder_layers > num_decoder_layers else "decoder"

             for _ in range(_A   ):
                 common_inputs["past_key_values"].append(
                     (
                         torch.zeros(_A   ),
                         torch.zeros(_A   ),
                         torch.zeros(_A   ),
                         torch.zeros(_A   ),
                     )   )
             # TODO: test this.
             snake_case      =					encoder_shape if remaining_side_name == "encoder" else decoder_shape
             for _ in range(_A				,    _A   ):
                 common_inputs["past_key_values"].append((torch.zeros(_A   ), torch.zeros(_A   ))   )
         return common_inputs





     def 		UpperCAmelCase(self      :    Tuple				,    _A      :    PreTrainedTokenizer				,    _A      :    int = -1				,    _A      :    int = -1				,    _A      :    bool = False				,    _A      :    Optional[TensorType] = None				,    )						->						Mapping[str, Any]:
         snake_case      =					self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             _A				,    _A				,    _A				,    _A				,    _A   )

         if self.use_past:
             if not is_torch_available():
                 raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."   )
             else:
                 import torch
             snake_case     ,							snake_case      =					common_inputs["input_ids"].shape
             # Not using the same length for past_key_values
             snake_case      =					seqlen + 2
             snake_case     ,							snake_case      =					self.num_layers
             snake_case     ,							snake_case      =					self.num_attention_heads
             snake_case      =					(
                 batch,
                 num_encoder_attention_heads,
                 past_key_values_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )

             snake_case      =					common_inputs["attention_mask"].dtype
             snake_case      =					torch.cat(
                 [common_inputs["attention_mask"], torch.ones(_A				,    _A				,    dtype=_A   )]				,    dim=1   )
             snake_case      =					[
                 (torch.zeros(_A   ), torch.zeros(_A   )) for _ in range(_A   )
             ]
         return common_inputs





     def 		UpperCAmelCase(self      :    Dict				,    _A      :    PreTrainedTokenizer				,    _A      :    int = -1				,    _A      :    int = -1				,    _A      :    bool = False				,    _A      :    Optional[TensorType] = None				,    )						->						Mapping[str, Any]:
         # Copied from OnnxConfig.generate_dummy_inputs
         # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
         # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
         snake_case      =					compute_effective_axis_dimension(
             _A				,    fixed_dimension=OnnxConfig.default_fixed_batch				,    num_token_to_add=0   )

         # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
         snake_case      =					tokenizer.num_special_tokens_to_add(_A   )
         snake_case      =					compute_effective_axis_dimension(
             _A				,    fixed_dimension=OnnxConfig.default_fixed_sequence				,    num_token_to_add=_A   )

         # Generate dummy inputs according to compute batch and sequence
         snake_case      =					[" ".join([tokenizer.unk_token]   ) * seq_length] * batch_size
         snake_case      =					dict(tokenizer(_A				,    return_tensors=_A   )   )
         return common_inputs





     def 		UpperCAmelCase(self      :    Any				,    _A      :    PreTrainedTokenizer				,    _A      :    int = -1				,    _A      :    int = -1				,    _A      :    bool = False				,    _A      :    Optional[TensorType] = None				,    )						->						Mapping[str, Any]:
         if self.task in ["default", "seq2seq-lm"]:
             snake_case      =					self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                 _A				,    batch_size=_A				,    seq_length=_A				,    is_pair=_A				,    framework=_A   )

         elif self.task == "causal-lm":
             snake_case      =					self._generate_dummy_inputs_for_causal_lm(
                 _A				,    batch_size=_A				,    seq_length=_A				,    is_pair=_A				,    framework=_A   )
         else:
             snake_case      =					self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                 _A				,    batch_size=_A				,    seq_length=_A				,    is_pair=_A				,    framework=_A   )

         return common_inputs





     def 		UpperCAmelCase(self      :    List[Any]				,    _A      :    str				,    _A      :    Union[str, Any]				,    _A      :    Tuple				,    _A      :    List[Any]   )						->						Any:
         if self.task in ["default", "seq2seq-lm"]:
             snake_case      =					super()._flatten_past_key_values_(_A				,    _A				,    _A				,    _A   )
         else:
             snake_case      =					super(_A				,    self   )._flatten_past_key_values_(
                 _A				,    _A				,    _A				,    _A   )

