"""simple docstring"""






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 			UpperCAmelCase__       (				lowerCAmelCase__     :List[str]					,		lowerCAmelCase__     :str		) ->    Optional[Any]:
      '''simple docstring'''


      lowercase														=	XCLIPTextConfig()

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

      if "large" in model_name:
            lowercase														=	7_6_8
            lowercase														=	3_0_7_2
            lowercase														=	1_2

            lowercase														=	1_0_2_4
            lowercase														=	4_0_9_6
            lowercase														=	1_6
            lowercase														=	2_4
            lowercase														=	7_6_8
            lowercase														=	3_0_7_2

      if model_name == "xclip-large-patch14-16-frames":
            lowercase														=	3_3_6

      lowercase														=	XCLIPConfig.from_text_vision_configs(lowerCAmelCase__					,		lowerCAmelCase__		)

      if "large" in model_name:
            lowercase														=	7_6_8

      return config


def 			UpperCAmelCase__       (				lowerCAmelCase__     :List[Any]		) ->    Any:
      '''simple docstring'''


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

      return name


def 			UpperCAmelCase__       (				lowerCAmelCase__     :Dict					,		lowerCAmelCase__     :Dict		) ->    List[str]:
      '''simple docstring'''


      for key in orig_state_dict.copy().keys():
            lowercase														=	orig_state_dict.pop(lowerCAmelCase__		)

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

      return orig_state_dict


def 			UpperCAmelCase__       (				lowerCAmelCase__     :List[str]		) ->    Tuple:
      '''simple docstring'''


      if num_frames == 8:
            lowercase														=	'''eating_spaghetti_8_frames.npy'''
      elif num_frames == 1_6:
            lowercase														=	'''eating_spaghetti.npy'''
      elif num_frames == 3_2:
            lowercase														=	'''eating_spaghetti_32_frames.npy'''
      lowercase														=	hf_hub_download(
          repo_id="""hf-internal-testing/spaghetti-video"""					,		filename=lowerCAmelCase__					,		repo_type="""dataset"""					,		)
      lowercase														=	np.load(lowerCAmelCase__		)
      return list(lowerCAmelCase__		)


def 			UpperCAmelCase__       (				lowerCAmelCase__     :Optional[int]					,		lowerCAmelCase__     :Tuple=None					,		lowerCAmelCase__     :int=False		) ->    Tuple:
      '''simple docstring'''


      lowercase														=	{
          # 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''',
      }

      lowercase														=	model_to_url[model_name]
      lowercase														=	8
      if "16-frames" in model_name:
            lowercase														=	1_6
      elif "shot" in model_name:
            lowercase														=	3_2

      lowercase														=	get_xclip_config(lowerCAmelCase__					,		lowerCAmelCase__		)
      lowercase														=	XCLIPModel(lowerCAmelCase__		)
      model.eval()

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

      lowercase														=	convert_state_dict(lowerCAmelCase__					,		lowerCAmelCase__		)

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

      lowercase														=	3_3_6 if model_name == '''xclip-large-patch14-16-frames''' else 2_2_4
      lowercase														=	VideoMAEImageProcessor(size=lowerCAmelCase__		)
      lowercase														=	CLIPTokenizer.from_pretrained("""openai/clip-vit-base-patch32"""		)
      lowercase														=	CLIPTokenizerFast.from_pretrained("""openai/clip-vit-base-patch32"""		)
      lowercase														=	XCLIPProcessor(image_processor=lowerCAmelCase__					,		tokenizer=lowerCAmelCase__		)

      lowercase														=	prepare_video(lowerCAmelCase__		)
      lowercase														=	processor(
          text=["""playing sports""", """eating spaghetti""", """go shopping"""]					,		videos=lowerCAmelCase__					,		return_tensors="""pt"""					,		padding=lowerCAmelCase__		)

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

      with torch.no_grad():
            lowercase														=	model(**lowerCAmelCase__		)

      # Verify outputs
      lowercase														=	outputs.logits_per_video
      lowercase														=	logits_per_video.softmax(dim=1		)
      print("""Probs:"""					,		lowerCAmelCase__		)
      # kinetics-400
      if model_name == "xclip-base-patch32":
            lowercase														=	torch.tensor([[0.0_019, 0.9_951, 0.0_030]]		)
      elif model_name == "xclip-base-patch32-16-frames":
            lowercase														=	torch.tensor([[7.09_99e-04, 9.98_83e-01, 4.55_80e-04]]		)
      elif model_name == "xclip-base-patch16":
            lowercase														=	torch.tensor([[0.0_083, 0.9_681, 0.0_236]]		)
      elif model_name == "xclip-base-patch16-16-frames":
            lowercase														=	torch.tensor([[7.69_37e-04, 9.97_28e-01, 1.94_73e-03]]		)
      elif model_name == "xclip-large-patch14":
            lowercase														=	torch.tensor([[0.0_062, 0.9_864, 0.0_075]]		)
      elif model_name == "xclip-large-patch14-16-frames":
            lowercase														=	torch.tensor([[3.38_77e-04, 9.99_37e-01, 2.88_88e-04]]		)
      # kinetics-600
      elif model_name == "xclip-base-patch16-kinetics-600":
            lowercase														=	torch.tensor([[0.0_555, 0.8_914, 0.0_531]]		)
      elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
            lowercase														=	torch.tensor([[3.85_54e-04, 9.99_29e-01, 3.27_54e-04]]		)
      elif model_name == "xclip-large-patch14-kinetics-600":
            lowercase														=	torch.tensor([[0.0_036, 0.9_920, 0.0_045]]		)
      # few shot
      elif model_name == "xclip-base-patch16-hmdb-2-shot":
            lowercase														=	torch.tensor([[7.18_90e-06, 9.99_94e-01, 5.65_59e-05]]		)
      elif model_name == "xclip-base-patch16-hmdb-4-shot":
            lowercase														=	torch.tensor([[1.03_20e-05, 9.99_93e-01, 6.24_35e-05]]		)
      elif model_name == "xclip-base-patch16-hmdb-8-shot":
            lowercase														=	torch.tensor([[4.13_77e-06, 9.99_90e-01, 9.83_86e-05]]		)
      elif model_name == "xclip-base-patch16-hmdb-16-shot":
            lowercase														=	torch.tensor([[4.13_47e-05, 9.99_62e-01, 3.34_11e-04]]		)
      elif model_name == "xclip-base-patch16-ucf-2-shot":
            lowercase														=	torch.tensor([[8.58_57e-05, 9.99_28e-01, 6.32_91e-04]]		)
      elif model_name == "xclip-base-patch16-ucf-4-shot":
            lowercase														=	torch.tensor([[8.58_57e-05, 9.99_28e-01, 6.32_91e-04]]		)
      elif model_name == "xclip-base-patch16-ucf-8-shot":
            lowercase														=	torch.tensor([[0.0_027, 0.9_904, 0.0_070]]		)
      elif model_name == "xclip-base-patch16-ucf-16-shot":
            lowercase														=	torch.tensor([[9.82_19e-04, 9.95_93e-01, 3.08_63e-03]]		)
      # zero shot
      elif model_name == "xclip-base-patch16-zero-shot":
            lowercase														=	torch.tensor([[3.50_82e-04, 9.97_85e-01, 1.79_66e-03]]		)
      else:
            raise ValueError(f'Model name {model_name} not supported'		)
      assert torch.allclose(lowerCAmelCase__					,		lowerCAmelCase__					,		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(lowerCAmelCase__		)

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


if __name__ == "__main__":
       __lowerCAmelCase    :       Union[str, Any]					     =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."""
       )

       __lowerCAmelCase    :       Dict					     =parser.parse_args()
       convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







'''simple docstring'''
import os
from collections import namedtuple

import pytest

from datasets import ClassLabel, Features, Sequence, Value
from datasets.commands.test import TestCommand
from datasets.info import DatasetInfo, DatasetInfosDict


UpperCamelCase__    					=      namedtuple(
    '''_TestCommandArgs''',
    [
        '''dataset''',
        '''name''',
        '''cache_dir''',
        '''data_dir''',
        '''all_configs''',
        '''save_infos''',
        '''ignore_verifications''',
        '''force_redownload''',
        '''clear_cache''',
    ],
    defaults=[None, None, None, False, False, False, False, False],
)





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Optional[int]:
    return (abs(source - target						) / target) < 0.0_1







@pytest.mark.integration
def 					a__       (					lowerCAmelCase__						)    ->	Optional[int]:
    UpperCAmelCase__			:		str													=				_TestCommandArgs(dataset=lowerCAmelCase__  ,					all_configs=lowerCAmelCase__  ,					save_infos=lowerCAmelCase__						)
    UpperCAmelCase__			:		Tuple													=				TestCommand(*lowerCAmelCase__						)
    test_command.run()
    UpperCAmelCase__			:		List[Any]													=				os.path.join(lowerCAmelCase__  ,					'''README.md'''						)
    assert os.path.exists(lowerCAmelCase__						)
    UpperCAmelCase__			:		List[str]													=				DatasetInfosDict.from_directory(lowerCAmelCase__						)
    UpperCAmelCase__			:		Any													=				DatasetInfosDict(
        {
            '''default''': DatasetInfo(
                features=Features(
                    {
                        '''tokens''': Sequence(Value('''string'''						)						),
                        '''ner_tags''': Sequence(
                            ClassLabel(names=['''O''', '''B-PER''', '''I-PER''', '''B-ORG''', '''I-ORG''', '''B-LOC''', '''I-LOC''']						)						),
                        '''langs''': Sequence(Value('''string'''						)						),
                        '''spans''': Sequence(Value('''string'''						)						),
                    }						)  ,					splits=[
                    {
                        '''name''': '''train''',
                        '''num_bytes''': 2_35_15_63,
                        '''num_examples''': 1_00_00,
                    },
                    {
                        '''name''': '''validation''',
                        '''num_bytes''': 23_84_18,
                        '''num_examples''': 10_00,
                    },
                ]  ,					download_size=3_94_06_80  ,					dataset_size=2_58_99_81  ,					)
        }						)
    assert dataset_infos.keys() == expected_dataset_infos.keys()
    for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
        UpperCAmelCase__						,						UpperCAmelCase__			:		str													=				getattr(dataset_infos['''default''']  ,					lowerCAmelCase__						), getattr(expected_dataset_infos['''default''']  ,					lowerCAmelCase__						)
        if key == "num_bytes":
            assert is_apercent_close(lowerCAmelCase__  ,					lowerCAmelCase__						)
        elif key == "splits":
            assert list(lowerCAmelCase__						) == list(lowerCAmelCase__						)
            for split in result:
                assert result[split].name == expected[split].name
                assert result[split].num_examples == expected[split].num_examples
                assert is_apercent_close(result[split].num_bytes  ,					expected[split].num_bytes						)
        else:
            result == expected
