





import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)






def lowerCamelCase__   (							snake_case_     :  Dict	,				snake_case_     :  List[str]	)				->    List[Any]:
   __snake_case								=       XCLIPTextConfig()

   # derive patch size from model name
   __snake_case								=       model_name.find('''patch'''	)
   __snake_case								=       int(model_name[start_idx + len('''patch'''	) : start_idx + len('''patch'''	) + 2]	)
   __snake_case								=       XCLIPVisionConfig(patch_size=snake_case_	,				num_frames=snake_case_	)

   if "large" in model_name:
      __snake_case								=       768
      __snake_case								=       3072
      __snake_case								=       12

      __snake_case								=       1024
      __snake_case								=       4096
      __snake_case								=       16
      __snake_case								=       24
      __snake_case								=       768
      __snake_case								=       3072

   if model_name == "xclip-large-patch14-16-frames":
      __snake_case								=       336

   __snake_case								=       XCLIPConfig.from_text_vision_configs(snake_case_	,				snake_case_	)

   if "large" in model_name:
      __snake_case								=       768

   return config






def lowerCamelCase__   (							snake_case_     :  List[str]	)				->    Tuple:
   # text encoder
   if name == "token_embedding.weight":
      __snake_case								=       name.replace('''token_embedding.weight'''	,				'''text_model.embeddings.token_embedding.weight'''	)
   if name == "positional_embedding":
      __snake_case								=       name.replace('''positional_embedding'''	,				'''text_model.embeddings.position_embedding.weight'''	)
   if "ln_1" in name:
      __snake_case								=       name.replace('''ln_1'''	,				'''layer_norm1'''	)
   if "ln_2" in name:
      __snake_case								=       name.replace('''ln_2'''	,				'''layer_norm2'''	)
   if "c_fc" in name:
      __snake_case								=       name.replace('''c_fc'''	,				'''fc1'''	)
   if "c_proj" in name:
      __snake_case								=       name.replace('''c_proj'''	,				'''fc2'''	)
   if name.startswith('''transformer.resblocks'''	):
      __snake_case								=       name.replace('''transformer.resblocks'''	,				'''text_model.encoder.layers'''	)
   if "attn.out_proj" in name and "message" not in name:
      __snake_case								=       name.replace('''attn.out_proj'''	,				'''self_attn.out_proj'''	)
   if "ln_final" in name:
      __snake_case								=       name.replace('''ln_final'''	,				'''text_model.final_layer_norm'''	)
   # visual encoder
   if name == "visual.class_embedding":
      __snake_case								=       name.replace('''visual.class_embedding'''	,				'''vision_model.embeddings.class_embedding'''	)
   if name == "visual.positional_embedding":
      __snake_case								=       name.replace('''visual.positional_embedding'''	,				'''vision_model.embeddings.position_embedding.weight'''	)
   if name.startswith('''visual.transformer.resblocks'''	):
      __snake_case								=       name.replace('''visual.transformer.resblocks'''	,				'''vision_model.encoder.layers'''	)
   if "visual.conv1" in name:
      __snake_case								=       name.replace('''visual.conv1'''	,				'''vision_model.embeddings.patch_embedding'''	)
   if "visual.ln_pre" in name:
      __snake_case								=       name.replace('''visual.ln_pre'''	,				'''vision_model.pre_layernorm'''	)
   if "visual.ln_post" in name:
      __snake_case								=       name.replace('''visual.ln_post'''	,				'''vision_model.post_layernorm'''	)
   if "visual.proj" in name:
      __snake_case								=       name.replace('''visual.proj'''	,				'''visual_projection.weight'''	)
   if "text_projection" in name:
      __snake_case								=       name.replace('''text_projection'''	,				'''text_projection.weight'''	)
   # things on top
   if "prompts_visual_proj" in name:
      __snake_case								=       name.replace('''prompts_visual_proj'''	,				'''prompts_visual_projection'''	)
   if "prompts_visual_ln" in name:
      __snake_case								=       name.replace('''prompts_visual_ln'''	,				'''prompts_visual_layernorm'''	)
   # mit
   if name == "mit.positional_embedding":
      __snake_case								=       name.replace('''positional'''	,				'''position'''	)
   if name.startswith('''mit.resblocks'''	):
      __snake_case								=       name.replace('''mit.resblocks'''	,				'''mit.encoder.layers'''	)
   # prompts generator
   if name.startswith('''prompts_generator.norm'''	):
      __snake_case								=       name.replace('''prompts_generator.norm'''	,				'''prompts_generator.layernorm'''	)

   return name






def lowerCamelCase__   (							snake_case_     :  str	,				snake_case_     :  List[str]	)				->    Optional[int]:
   for key in orig_state_dict.copy().keys():
      __snake_case								=       orig_state_dict.pop(snake_case_	)

      if "attn.in_proj" in key:
         __snake_case								=       key.split('''.'''	)
         if key.startswith('''visual'''	):
            __snake_case								=       key_split[3]
            __snake_case								=       config.vision_config.hidden_size
            if "message_attn" in key:
               if "weight" in key:
                  __snake_case								=       val[
                      :dim, :
                  ]
                  __snake_case								=       val[
                      dim : dim * 2, :
                  ]
                  __snake_case								=       val[
                      -dim:, :
                  ]
               else:
                  __snake_case								=       val[
                      :dim
                  ]
                  __snake_case								=       val[
                      dim : dim * 2
                  ]
                  __snake_case								=       val[
                      -dim:
                  ]
            else:
               if "weight" in key:
                  __snake_case								=       val[
                      :dim, :
                  ]
                  __snake_case								=       val[
                      dim : dim * 2, :
                  ]
                  __snake_case								=       val[
                      -dim:, :
                  ]
               else:
                  __snake_case								=       val[:dim]
                  __snake_case								=       val[
                      dim : dim * 2
                  ]
                  __snake_case								=       val[-dim:]
         elif key.startswith('''mit'''	):
            __snake_case								=       key_split[2]
            __snake_case								=       config.vision_config.mit_hidden_size
            if "weight" in key:
               __snake_case								=       val[:dim, :]
               __snake_case								=       val[dim : dim * 2, :]
               __snake_case								=       val[-dim:, :]
            else:
               __snake_case								=       val[:dim]
               __snake_case								=       val[dim : dim * 2]
               __snake_case								=       val[-dim:]
         else:
            __snake_case								=       key_split[2]
            __snake_case								=       config.text_config.hidden_size
            if "weight" in key:
               __snake_case								=       val[:dim, :]
               __snake_case								=       val[
                   dim : dim * 2, :
               ]
               __snake_case								=       val[-dim:, :]
            else:
               __snake_case								=       val[:dim]
               __snake_case								=       val[
                   dim : dim * 2
               ]
               __snake_case								=       val[-dim:]
      else:
         __snake_case								=       rename_key(snake_case_	)
         if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
            __snake_case								=       val.T
         __snake_case								=       val

   return orig_state_dict






def lowerCamelCase__   (							snake_case_     :  Any	)				->    List[Any]:
   if num_frames == 8:
      __snake_case								=       '''eating_spaghetti_8_frames.npy'''
   elif num_frames == 16:
      __snake_case								=       '''eating_spaghetti.npy'''
   elif num_frames == 32:
      __snake_case								=       '''eating_spaghetti_32_frames.npy'''
   __snake_case								=       hf_hub_download(
       repo_id='''hf-internal-testing/spaghetti-video'''	,				filename=snake_case_	,				repo_type='''dataset'''	,				)
   __snake_case								=       np.load(snake_case_	)
   return list(snake_case_	)






def lowerCamelCase__   (							snake_case_     :  int	,				snake_case_     :  str=None	,				snake_case_     :  Optional[int]=False	)				->    str:
   __snake_case								=       {
       # fully supervised kinetics-400 checkpoints
       '''xclip-base-patch32''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth''',
       '''xclip-base-patch32-16-frames''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth'''
       ),
       '''xclip-base-patch16''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth''',
       '''xclip-base-patch16-16-frames''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth'''
       ),
       '''xclip-large-patch14''': '''https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb''',
       '''xclip-large-patch14-16-frames''': '''https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f''',
       # fully supervised kinetics-600 checkpoints
       '''xclip-base-patch16-kinetics-600''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth'''
       ),
       '''xclip-base-patch16-kinetics-600-16-frames''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth'''
       ),
       '''xclip-large-patch14-kinetics-600''': '''https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be''',
       # few shot
       '''xclip-base-patch16-hmdb-2-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth'''
       ),
       '''xclip-base-patch16-hmdb-4-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth'''
       ),
       '''xclip-base-patch16-hmdb-8-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth'''
       ),
       '''xclip-base-patch16-hmdb-16-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth'''
       ),
       '''xclip-base-patch16-ucf-2-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth'''
       ),
       '''xclip-base-patch16-ucf-4-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth'''
       ),
       '''xclip-base-patch16-ucf-8-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth'''
       ),
       '''xclip-base-patch16-ucf-16-shot''': (
           '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth'''
       ),
       # zero shot
       '''xclip-base-patch16-zero-shot''': '''https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth''',
   }

   __snake_case								=       model_to_url[model_name]
   __snake_case								=       8
   if "16-frames" in model_name:
      __snake_case								=       16
   elif "shot" in model_name:
      __snake_case								=       32

   __snake_case								=       get_xclip_config(snake_case_	,				snake_case_	)
   __snake_case								=       XCLIPModel(snake_case_	)
   model.eval()

   if "drive" in checkpoint_url:
      __snake_case								=       '''pytorch_model.bin'''
      gdown.cached_download(snake_case_	,				snake_case_	,				quiet=snake_case_	)
      __snake_case								=       torch.load(snake_case_	,				map_location='''cpu'''	)['''model''']
   else:
      __snake_case								=       torch.hub.load_state_dict_from_url(snake_case_	)['''model''']

   __snake_case								=       convert_state_dict(snake_case_	,				snake_case_	)

   __snake_case								=       XCLIPModel(snake_case_	)
   __snake_case	,     __snake_case								=       model.load_state_dict(snake_case_	,				strict=snake_case_	)
   assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
   model.eval()

   __snake_case								=       336 if model_name == '''xclip-large-patch14-16-frames''' else 224
   __snake_case								=       VideoMAEImageProcessor(size=snake_case_	)
   __snake_case								=       CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''	)
   __snake_case								=       CLIPTokenizerFast.from_pretrained('''openai/clip-vit-base-patch32'''	)
   __snake_case								=       XCLIPProcessor(image_processor=snake_case_	,				tokenizer=snake_case_	)

   __snake_case								=       prepare_video(snake_case_	)
   __snake_case								=       processor(
       text=['''playing sports''', '''eating spaghetti''', '''go shopping''']	,				videos=snake_case_	,				return_tensors='''pt'''	,				padding=snake_case_	)

   print('''Shape of pixel values:'''	,				inputs.pixel_values.shape	)

   with torch.no_grad():
      __snake_case								=       model(**snake_case_	)

   # Verify outputs
   __snake_case								=       outputs.logits_per_video
   __snake_case								=       logits_per_video.softmax(dim=1	)
   print('''Probs:'''	,				snake_case_	)
   # kinetics-400
   if model_name == "xclip-base-patch32":
      __snake_case								=       torch.tensor([[0.0_019, 0.9_951, 0.0_030]]	)
   elif model_name == "xclip-base-patch32-16-frames":
      __snake_case								=       torch.tensor([[7.0999e-04, 9.9883e-01, 4.5580e-04]]	)
   elif model_name == "xclip-base-patch16":
      __snake_case								=       torch.tensor([[0.0_083, 0.9_681, 0.0_236]]	)
   elif model_name == "xclip-base-patch16-16-frames":
      __snake_case								=       torch.tensor([[7.6937e-04, 9.9728e-01, 1.9473e-03]]	)
   elif model_name == "xclip-large-patch14":
      __snake_case								=       torch.tensor([[0.0_062, 0.9_864, 0.0_075]]	)
   elif model_name == "xclip-large-patch14-16-frames":
      __snake_case								=       torch.tensor([[3.3877e-04, 9.9937e-01, 2.8888e-04]]	)
   # kinetics-600
   elif model_name == "xclip-base-patch16-kinetics-600":
      __snake_case								=       torch.tensor([[0.0_555, 0.8_914, 0.0_531]]	)
   elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
      __snake_case								=       torch.tensor([[3.8554e-04, 9.9929e-01, 3.2754e-04]]	)
   elif model_name == "xclip-large-patch14-kinetics-600":
      __snake_case								=       torch.tensor([[0.0_036, 0.9_920, 0.0_045]]	)
   # few shot
   elif model_name == "xclip-base-patch16-hmdb-2-shot":
      __snake_case								=       torch.tensor([[7.1890e-06, 9.9994e-01, 5.6559e-05]]	)
   elif model_name == "xclip-base-patch16-hmdb-4-shot":
      __snake_case								=       torch.tensor([[1.0320e-05, 9.9993e-01, 6.2435e-05]]	)
   elif model_name == "xclip-base-patch16-hmdb-8-shot":
      __snake_case								=       torch.tensor([[4.1377e-06, 9.9990e-01, 9.8386e-05]]	)
   elif model_name == "xclip-base-patch16-hmdb-16-shot":
      __snake_case								=       torch.tensor([[4.1347e-05, 9.9962e-01, 3.3411e-04]]	)
   elif model_name == "xclip-base-patch16-ucf-2-shot":
      __snake_case								=       torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]	)
   elif model_name == "xclip-base-patch16-ucf-4-shot":
      __snake_case								=       torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]	)
   elif model_name == "xclip-base-patch16-ucf-8-shot":
      __snake_case								=       torch.tensor([[0.0_027, 0.9_904, 0.0_070]]	)
   elif model_name == "xclip-base-patch16-ucf-16-shot":
      __snake_case								=       torch.tensor([[9.8219e-04, 9.9593e-01, 3.0863e-03]]	)
   # zero shot
   elif model_name == "xclip-base-patch16-zero-shot":
      __snake_case								=       torch.tensor([[3.5082e-04, 9.9785e-01, 1.7966e-03]]	)
   else:
      raise ValueError(f"""Model name {model_name} not supported"""	)
   assert torch.allclose(snake_case_	,				snake_case_	,				atol=1e-3	)
   print('''Looks ok!'''	)

   if pytorch_dump_folder_path is not None:
      print(f"""Saving model {model_name} to {pytorch_dump_folder_path}"""	)
      model.save_pretrained(snake_case_	)

   if push_to_hub:
      print('''Pushing model, processor and slow tokenizer files to the hub...'''	)
      model.push_to_hub(snake_case_	,				organization='''nielsr'''	)
      processor.push_to_hub(snake_case_	,				organization='''nielsr'''	)
      slow_tokenizer.push_to_hub(snake_case_	,				organization='''nielsr'''	)


if __name__ == "__main__":
    snake_case_           =      argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '--model_name',
        default='xclip-base-patch32',
        type=str,
        help='Name of the model.',
    )
    parser.add_argument(
        '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
    )
    parser.add_argument(
        '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
    )

    snake_case_           =      parser.parse_args()
    convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
    import torch

class 		SCREAMING_SNAKE_CASE__    (	TensorFormatter[Mapping, 'torch.Tensor', Mapping]  ):


      def __init__(self	:				Optional[int]    ,      a__	:				Tuple=None    ,      **a__	:				Optional[int]				):





         """simple docstring"""



         super().__init__(features=a__				)
         __snake_case								=       torch_tensor_kwargs
         import torch  # noqa import torch at initialization


      def        a			(self	:				Union[str, Any]    ,      a__	:				Union[str, Any]				):





         """simple docstring"""



         import torch

         if isinstance(a__    ,      a__				) and column:
            if all(
                isinstance(a__    ,      torch.Tensor				) and x.shape == column[0].shape and x.dtype == column[0].dtype
                for x in column				):
               return torch.stack(a__				)
         return column


      def        a			(self	:				Optional[Any]    ,      a__	:				str				):





         """simple docstring"""



         import torch

         if isinstance(a__    ,      (str, bytes, type(a__				))				):
            return value
         elif isinstance(a__    ,      (np.character, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.character				):
            return value.tolist()

         __snake_case								=       {}

         if isinstance(a__    ,      (np.number, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.integer				):
            __snake_case								=       {'''dtype''': torch.intaa}
         elif isinstance(a__    ,      (np.number, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.floating				):
            __snake_case								=       {'''dtype''': torch.floataa}
         elif config.PIL_AVAILABLE and "PIL" in sys.modules:
            import PIL.Image

            if isinstance(a__    ,      PIL.Image.Image				):
               __snake_case								=       np.asarray(a__				)
         return torch.tensor(a__    ,      **{**default_dtype, **self.torch_tensor_kwargs}				)


      def        a			(self	:				Optional[int]    ,      a__	:				str				):





         """simple docstring"""



         import torch

         # support for torch, tf, jax etc.
         if hasattr(a__    ,      '''__array__'''				) and not isinstance(a__    ,      torch.Tensor				):
            __snake_case								=       data_struct.__array__()
         # support for nested types like struct of list of struct
         if isinstance(a__    ,      np.ndarray				):
            if data_struct.dtype == object:  # torch tensors cannot be instantied from an array of objects
               return self._consolidate([self.recursive_tensorize(a__				) for substruct in data_struct]				)
         elif isinstance(a__    ,      (list, tuple)				):
            return self._consolidate([self.recursive_tensorize(a__				) for substruct in data_struct]				)
         return self._tensorize(a__				)


      def        a			(self	:				Optional[Any]    ,      a__	:				dict				):





         """simple docstring"""



         return map_nested(self._recursive_tensorize    ,      a__    ,      map_list=a__				)


      def        a			(self	:				Optional[int]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_row(a__				)
         __snake_case								=       self.python_features_decoder.decode_row(a__				)
         return self.recursive_tensorize(a__				)


      def        a			(self	:				List[Any]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_column(a__				)
         __snake_case								=       self.python_features_decoder.decode_column(a__    ,      pa_table.column_names[0]				)
         __snake_case								=       self.recursive_tensorize(a__				)
         __snake_case								=       self._consolidate(a__				)
         return column


      def        a			(self	:				Optional[int]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_batch(a__				)
         __snake_case								=       self.python_features_decoder.decode_batch(a__				)
         __snake_case								=       self.recursive_tensorize(a__				)
         for column_name in batch:
            __snake_case								=       self._consolidate(batch[column_name]				)
         return batch

