





'''simple docstring'''


import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


lowerCAmelCase_   :   Optional[Any]       =					get_logger(__name__)


lowerCAmelCase_   :   List[Any]       =					R"\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n"






class       UpperCamelCase__  :







	@add_start_docstrings(lowerCamelCase		)
	def __call__(				self			:   List[str] , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray		):




				'''simple docstring'''
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






class       UpperCamelCase__  :







	@add_start_docstrings(lowerCamelCase		)
	def __call__(				self			:   Dict , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray		):




				'''simple docstring'''
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	@add_start_docstrings(lowerCamelCase		)
	def __call__(				self			:   Any , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int , **lowerCamelCase			:   Any		):




				'''simple docstring'''
				for processor in self:
							a__     = inspect.signature(processor.__call__		).parameters
							if len(lowerCamelCase		) > 3:
										if not all(arg in kwargs for arg in list(function_args.keys()		)[2:]		):
													raise ValueError(
													    F'''Make sure that all the required parameters: {list(function_args.keys()		)} for '''
													    F'''{processor.__class__} are passed to the logits processor.'''		)
										a__     = processor(lowerCamelCase , lowerCamelCase , lowerCamelCase , **lowerCamelCase		)
							else:
										a__     = processor(lowerCamelCase , lowerCamelCase , lowerCamelCase		)
				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   Optional[Any] , lowerCamelCase			:   float		):




				'''simple docstring'''
				if not isinstance(lowerCamelCase , lowerCamelCase		) or not (temperature > 0):
							raise ValueError(F'''`temperature` has to be a strictly positive float, but is {temperature}'''		)

				a__     = temperature







	def __call__(				self			:   List[Any] , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = scores / self.temperature
				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   Tuple , lowerCamelCase			:   float , lowerCamelCase			:   float = -float("Inf"		) , lowerCamelCase			:   int = 1		):




				'''simple docstring'''
				if not isinstance(lowerCamelCase , lowerCamelCase		) or (top_p < 0 or top_p > 1.0):
							raise ValueError(F'''`top_p` has to be a float > 0 and < 1, but is {top_p}'''		)
				if not isinstance(lowerCamelCase , lowerCamelCase		) or (min_tokens_to_keep < 1):
							raise ValueError(F'''`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}'''		)

				a__     = top_p
				a__     = filter_value
				a__     = min_tokens_to_keep







	def __call__(				self			:   Union[str, Any] , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__   ,				a__     = lax.top_k(lowerCamelCase , scores.shape[-1]		)

				a__     = jnp.full_like(lowerCamelCase , self.filter_value		)
				a__     = jax.nn.softmax(lowerCamelCase , axis=-1		).cumsum(axis=-1		)
				a__     = cumulative_probs < self.top_p

				# include the token that is higher than top_p as well
				a__     = jnp.roll(lowerCamelCase , 1		)
				score_mask |= score_mask.at[:, 0].set(lowerCamelCase		)

				# min tokens to keep
				a__     = score_mask.at[:, : self.min_tokens_to_keep].set(lowerCamelCase		)

				a__     = jnp.where(lowerCamelCase , lowerCamelCase , lowerCamelCase		)
				a__     = jax.lax.sort_key_val(lowerCamelCase , lowerCamelCase		)[-1]

				return next_scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   str , lowerCamelCase			:   int , lowerCamelCase			:   float = -float("Inf"		) , lowerCamelCase			:   int = 1		):




				'''simple docstring'''
				if not isinstance(lowerCamelCase , lowerCamelCase		) or top_k <= 0:
							raise ValueError(F'''`top_k` has to be a strictly positive integer, but is {top_k}'''		)

				a__     = max(lowerCamelCase , lowerCamelCase		)
				a__     = filter_value







	def __call__(				self			:   Any , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__   ,				a__     = scores.shape
				a__     = jnp.full(batch_size * vocab_size , self.filter_value		)

				a__     = min(self.top_k , scores.shape[-1]		)  # Safety check
				a__   ,				a__     = lax.top_k(lowerCamelCase , lowerCamelCase		)
				a__     = jnp.broadcast_to((jnp.arange(lowerCamelCase		) * vocab_size)[:, None] , (batch_size, topk)		).flatten()
				a__     = topk_scores.flatten()
				a__     = topk_indices.flatten() + shift

				a__     = next_scores_flat.at[topk_indices_flat].set(lowerCamelCase		)
				a__     = next_scores_flat.reshape(lowerCamelCase , lowerCamelCase		)
				return next_scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   str , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = bos_token_id







	def __call__(				self			:   Union[str, Any] , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = jnp.full(scores.shape , -float("inf"		)		)

				a__     = 1 - jnp.bool_(cur_len - 1		)

				a__     = jnp.where(lowerCamelCase , new_scores.at[:, self.bos_token_id].set(0		) , lowerCamelCase		)

				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   List[Any] , lowerCamelCase			:   int , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = max_length
				a__     = eos_token_id







	def __call__(				self			:   Dict , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = jnp.full(scores.shape , -float("inf"		)		)

				a__     = 1 - jnp.bool_(cur_len - self.max_length + 1		)

				a__     = jnp.where(lowerCamelCase , new_scores.at[:, self.eos_token_id].set(0		) , lowerCamelCase		)

				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   List[Any] , lowerCamelCase			:   int , lowerCamelCase			:   int		):




				'''simple docstring'''
				if not isinstance(lowerCamelCase , lowerCamelCase		) or min_length < 0:
							raise ValueError(F'''`min_length` has to be a positive integer, but is {min_length}'''		)

				if not isinstance(lowerCamelCase , lowerCamelCase		) or eos_token_id < 0:
							raise ValueError(F'''`eos_token_id` has to be a positive integer, but is {eos_token_id}'''		)

				a__     = min_length
				a__     = eos_token_id







	def __call__(				self			:   Optional[int] , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				# create boolean flag to decide if min length penalty should be applied
				a__     = 1 - jnp.clip(cur_len - self.min_length , 0 , 1		)

				a__     = jnp.where(lowerCamelCase , scores.at[:, self.eos_token_id].set(-float("inf"		)		) , lowerCamelCase		)

				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   Any , lowerCamelCase			:   List[Any] , lowerCamelCase			:   List[str]		):




				'''simple docstring'''
				a__     = list(lowerCamelCase		)
				a__     = begin_index







	def __call__(				self			:   Optional[int] , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = 1 - jnp.bool_(cur_len - self.begin_index		)

				a__     = jnp.where(lowerCamelCase , scores.at[:, self.begin_suppress_tokens].set(-float("inf"		)		) , lowerCamelCase		)

				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   List[Any] , lowerCamelCase			:   list		):




				'''simple docstring'''
				a__     = list(lowerCamelCase		)







	def __call__(				self			:   List[Any] , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = scores.at[..., self.suppress_tokens].set(-float("inf"		)		)

				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   List[Any] , lowerCamelCase			:   Union[str, Any]		):




				'''simple docstring'''
				a__     = dict(lowerCamelCase		)
				# Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
				# index of the array corresponds to the index of the token to be forced, for XLA compatibility.
				# Indexes without forced tokens will have a negative value.
				a__     = jnp.ones((max(force_token_map.keys()		) + 1) , dtype=jnp.intaa		) * -1
				for index, token in force_token_map.items():
							if token is not None:
										a__     = force_token_array.at[index].set(lowerCamelCase		)
				a__     = jnp.intaa(lowerCamelCase		)







	def __call__(				self			:   int , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   jnp.ndarray , lowerCamelCase			:   int		):




				'''simple docstring'''
				def _force_token(lowerCamelCase			:   Tuple		):
							a__     = scores.shape[0]
							a__     = self.force_token_array[generation_idx]

							a__     = jnp.ones_like(lowerCamelCase , dtype=scores.dtype		) * -float("inf"		)
							a__     = jnp.zeros((batch_size, 1) , dtype=scores.dtype		)
							a__     = lax.dynamic_update_slice(lowerCamelCase , lowerCamelCase , (0, current_token)		)
							return new_scores

				a__     = lax.cond(
				    cur_len >= self.force_token_array.shape[0] , lambda: scores , lambda: lax.cond(
				        self.force_token_array[cur_len] >= 0 , lambda: _force_token(lowerCamelCase		) , lambda: scores , ) , )
				return scores






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   Tuple , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   str , lowerCamelCase			:   Tuple		):




				'''simple docstring'''
				a__     = generate_config.eos_token_id
				a__     = generate_config.no_timestamps_token_id
				a__     = generate_config.no_timestamps_token_id + 1

				a__     = decoder_input_length + 1

				if generate_config.is_multilingual:
							# room for language token and task token
							self.begin_index += 2
				if hasattr(lowerCamelCase , "max_initial_timestamp_index"		):
							a__     = generate_config.max_initial_timestamp_index
				else:
							a__     = model_config.vocab_size
				if self.max_initial_timestamp_index is None:
							a__     = model_config.vocab_size







	def __call__(				self			:   List[str] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   List[Any] , lowerCamelCase			:   str		):




				'''simple docstring'''
				# suppress <|notimestamps|> which is handled by without_timestamps
				a__     = scores.at[:, self.no_timestamps_token_id].set(-float("inf"		)		)

				def handle_pairs(lowerCamelCase			:   List[Any] , lowerCamelCase			:   List[str]		):
							a__     = jnp.where((cur_len - self.begin_index) >= 1 , lowerCamelCase , lowerCamelCase		)
							a__     = jnp.where(
							    input_ids_k[cur_len - 1] >= self.timestamp_begin , True and last_was_timestamp , lowerCamelCase , )

							a__     = jnp.where((cur_len - self.begin_index) < 2 , lowerCamelCase , lowerCamelCase		)
							a__     = jnp.where(
							    input_ids_k[cur_len - 2] >= self.timestamp_begin , lowerCamelCase , lowerCamelCase , )

							return jnp.where(
							    lowerCamelCase , jnp.where(
							        penultimate_was_timestamp > 0 , scores_k.at[self.timestamp_begin :].set(-float("inf"		)		) , scores_k.at[: self.eos_token_id].set(-float("inf"		)		) , ) , lowerCamelCase , )

				a__     = jax.vmap(lowerCamelCase		)(lowerCamelCase , lowerCamelCase		)

				a__     = jnp.where(cur_len == self.begin_index , lowerCamelCase , lowerCamelCase		)
				a__     = jnp.where(
				    self.max_initial_timestamp_index is not None , True and apply_max_initial_timestamp , lowerCamelCase , )

				a__     = self.timestamp_begin + self.max_initial_timestamp_index

				a__     = jnp.where(
				    lowerCamelCase , scores.at[:, last_allowed + 1 :].set(-float("inf"		)		) , lowerCamelCase , )

				# if sum of probability over timestamps is above any other token, sample timestamp
				a__     = jax.nn.log_softmax(lowerCamelCase , axis=-1		)

				def handle_cumulative_probs(lowerCamelCase			:   Tuple , lowerCamelCase			:   Optional[Any]		):
							a__     = jax.nn.logsumexp(logprobs_k[self.timestamp_begin :] , axis=-1		)
							a__     = jnp.max(logprobs_k[: self.timestamp_begin]		)
							return jnp.where(
							    timestamp_logprob > max_text_token_logprob , scores_k.at[: self.timestamp_begin].set(-float("inf"		)		) , lowerCamelCase , )

				a__     = jax.vmap(lowerCamelCase		)(lowerCamelCase , lowerCamelCase		)

				return scores










'''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    _lowerCamelCase  (__lowerCamelCase  :     Tuple					,  __lowerCamelCase  :     Dict     )   ->    Tuple:
			a__     = XCLIPTextConfig()

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

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

						a__     = 1024
						a__     = 4096
						a__     = 16
						a__     = 24
						a__     = 768
						a__     = 3072

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

			a__     = XCLIPConfig.from_text_vision_configs(__lowerCamelCase					,  __lowerCamelCase     )

			if "large" in model_name:
						a__     = 768

			return config







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

			return name







def    _lowerCamelCase  (__lowerCamelCase  :     Any					,  __lowerCamelCase  :     str     )   ->    Optional[Any]:
			for key in orig_state_dict.copy().keys():
						a__     = orig_state_dict.pop(__lowerCamelCase     )

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

			return orig_state_dict







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







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple					,  __lowerCamelCase  :     Any=None					,  __lowerCamelCase  :     Optional[Any]=False     )   ->    Optional[Any]:
			a__     = {
			    # 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",
			}

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

			a__     = get_xclip_config(__lowerCamelCase					,  __lowerCamelCase     )
			a__     = XCLIPModel(__lowerCamelCase     )
			model.eval()

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

			a__     = convert_state_dict(__lowerCamelCase					,  __lowerCamelCase     )

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

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

			a__     = prepare_video(__lowerCamelCase     )
			a__     = 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():
						a__     = model(**__lowerCamelCase     )

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



