
"""simple docstring"""






from sklearn.metrics import mean_squared_error

import datasets


__lowerCamelCase										=		"\\n@article{scikit-learn,\n title={Scikit-learn: Machine Learning in {P}ython},\n author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n journal={Journal of Machine Learning Research},\n volume={12},\n pages={2825--2830},\n year={2011}\n}\n"

__lowerCamelCase										=		"\\nMean Squared Error(MSE) is the average of the square of difference between the predicted\nand actual values.\n"


__lowerCamelCase										=		"\nArgs:\n    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Estimated target values.\n    references: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Ground truth (correct) target values.\n    sample_weight: array-like of shape (n_samples,), default=None\n        Sample weights.\n    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"\n        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.\n\n        \"raw_values\" : Returns a full set of errors in case of multioutput input.\n\n        \"uniform_average\" : Errors of all outputs are averaged with uniform weight.\n\n    squared : bool, default=True\n        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.\n\nReturns:\n    mse : mean squared error.\nExamples:\n\n    >>> mse_metric = datasets.load_metric(\"mse\")\n    >>> predictions = [2.5, 0.0, 2, 8]\n    >>> references = [3, -0.5, 2, 7]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.375}\n    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)\n    >>> print(rmse_result)\n    {'mse': 0.6123724356957945}\n\n    If you're using multi-dimensional lists, then set the config as follows :\n\n    >>> mse_metric = datasets.load_metric(\"mse\", \"multilist\")\n    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]\n    >>> references = [[0, 2], [-1, 2], [8, -5]]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.7083333333333334}\n    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput='raw_values')\n    >>> print(results) # doctest: +NORMALIZE_WHITESPACE\n    {'mse': array([0.41666667, 1. ])}\n"
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION   , _KWARGS_DESCRIPTION     )
class        _lowercase					(     datasets.Metric     ):







     def 	lowerCAmelCase__    (      self       ):
      return datasets.MetricInfo(
          description=_DESCRIPTION ,						citation=_CITATION ,						inputs_description=_KWARGS_DESCRIPTION ,						features=datasets.Features(self._get_feature_types()       ) ,						reference_urls=[
              '''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html'''
          ] ,						)







     def 	lowerCAmelCase__    (      self       ):




      if self.config_name == "multilist":
       return {
           "predictions": datasets.Sequence(datasets.Value('''float'''       )       ),
           "references": datasets.Sequence(datasets.Value('''float'''       )       ),
       }
      else:
       return {
           "predictions": datasets.Value('''float'''       ),
           "references": datasets.Value('''float'''       ),
       }







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_ ,						UpperCamelCase_ ,						UpperCamelCase_=None ,						UpperCamelCase_="uniform_average" ,						UpperCamelCase_=True       ):
      __magic_name__     =      mean_squared_error(
          lowerCamelCase_ ,						lowerCamelCase_ ,						sample_weight=lowerCamelCase_ ,						multioutput=lowerCamelCase_ ,						squared=lowerCamelCase_       )

      return {"mse": mse}




"""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 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase						) ->      List[str]:
 __magic_name__     =      XCLIPTextConfig()

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

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

  __magic_name__     =      1024
  __magic_name__     =      4096
  __magic_name__     =      16
  __magic_name__     =      24
  __magic_name__     =      768
  __magic_name__     =      3072

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

 __magic_name__     =      XCLIPConfig.from_text_vision_configs(__UpperCamelCase      ,   __UpperCamelCase						)

 if "large" in model_name:
  __magic_name__     =      768

 return config





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

 return name





def 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase						) ->      str:
 for key in orig_state_dict.copy().keys():
  __magic_name__     =      orig_state_dict.pop(__UpperCamelCase						)

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

 return orig_state_dict





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





def 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase=None      ,   __UpperCamelCase=False						) ->      Tuple:
 __magic_name__     =      {
     # 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''',
 }

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

 __magic_name__     =      get_xclip_config(__UpperCamelCase      ,   __UpperCamelCase						)
 __magic_name__     =      XCLIPModel(__UpperCamelCase						)
 model.eval()

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

 __magic_name__     =      convert_state_dict(__UpperCamelCase      ,   __UpperCamelCase						)

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

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

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

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

 with torch.no_grad():
  __magic_name__     =      model(**__UpperCamelCase						)

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

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


if __name__ == "__main__":
 __lowerCamelCase										=		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										=		parser.parse_args()
 convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)

