
'''simple docstring'''
from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed







def        A_     (						__SCREAMING_SNAKE_CASE						:  Union[str, Any]			,     __SCREAMING_SNAKE_CASE						:  Tuple			,     __SCREAMING_SNAKE_CASE						:  Optional[Any]			,     __SCREAMING_SNAKE_CASE						:  Any ):





  """simple docstring"""

  for param, grad_param in zip(model_a.parameters()			,     model_b.parameters() ):
    if not param.requires_grad:
      continue
    if not did_step:
      # Grads should not be in sync
      assert (
          torch.allclose(param.grad			,     grad_param.grad ) is False
      ), F"Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"
    else:
      # Grads should be in sync
      assert (
          torch.allclose(param.grad			,     grad_param.grad ) is True
      ), F"Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"







def        A_     (						__SCREAMING_SNAKE_CASE						:  Tuple			,     __SCREAMING_SNAKE_CASE						:  Optional[Any]			,     __SCREAMING_SNAKE_CASE						:  Union[str, Any]			,     __SCREAMING_SNAKE_CASE						:  List[Any]			,     __SCREAMING_SNAKE_CASE						:  Any=True ):





  """simple docstring"""

  model.train()
  __A		:     Any												=		model(__SCREAMING_SNAKE_CASE )
  __A		:     Any												=		F.mse_loss(__SCREAMING_SNAKE_CASE			,     target.to(output.device ) )
  if not do_backward:
    loss /= accelerator.gradient_accumulation_steps
    loss.backward()
  else:
    accelerator.backward(__SCREAMING_SNAKE_CASE )







def        A_     (						__SCREAMING_SNAKE_CASE						:  Optional[int]			,     __SCREAMING_SNAKE_CASE						:  Optional[int]=False ):





  """simple docstring"""

  set_seed(42 )
  __A		:     Union[str, Any]												=		RegressionModel()
  __A		:     Union[str, Any]												=		deepcopy(__SCREAMING_SNAKE_CASE )
  __A		:     List[str]												=		RegressionDataset(length=80 )
  __A		:     Tuple												=		DataLoader(__SCREAMING_SNAKE_CASE			,     batch_size=16 )
  model.to(accelerator.device )
  if sched:
    __A		:     Optional[int]												=		AdamW(params=model.parameters()			,     lr=1e-3 )
    __A		:     Optional[int]												=		AdamW(params=ddp_model.parameters()			,     lr=1e-3 )
    __A		:     Tuple												=		LambdaLR(__SCREAMING_SNAKE_CASE			,     lr_lambda=lambda __SCREAMING_SNAKE_CASE : epoch**0.65 )
    __A		:     List[Any]												=		LambdaLR(__SCREAMING_SNAKE_CASE			,     lr_lambda=lambda __SCREAMING_SNAKE_CASE : epoch**0.65 )
  # Make a copy of `model`
  if sched:
    __A		:     str												=		accelerator.prepare(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
  else:
    __A		:     List[str]												=		accelerator.prepare(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
  if sched:
    return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
  return model, ddp_model, dataloader







def        A_     (						__SCREAMING_SNAKE_CASE						:  Tuple ):





  """simple docstring"""

  __A		:     str												=		get_training_setup(__SCREAMING_SNAKE_CASE )
  # Use a single batch
  __A		:     int												=		next(iter(__SCREAMING_SNAKE_CASE ) ).values()
  for iteration in range(3 ):
    # Gather the distributed inputs and targs for the base model
    __A		:     Tuple												=		accelerator.gather((ddp_input, ddp_target) )
    __A		:     Any												=		input.to(accelerator.device ), target.to(accelerator.device )
    # Perform our initial ground truth step in non "DDP"
    step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    # Do "gradient accumulation" (noop)
    if iteration % 2 == 0:
      # Accumulate grads locally
      with accelerator.no_sync(__SCREAMING_SNAKE_CASE ):
        step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    else:
      # Sync grads
      step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )

    # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
    check_model_parameters(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    for param, ddp_param in zip(model.parameters()			,     ddp_model.parameters() ):
      if not param.requires_grad:
        continue
      assert torch.allclose(
          param.grad			,     ddp_param.grad ), F"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"

    # Shuffle ddp_input on each iteration
    torch.manual_seed(1337 + iteration )
    __A		:     int												=		ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE ) )]







def        A_     (						__SCREAMING_SNAKE_CASE						:  Tuple ):





  """simple docstring"""

  __A		:     Union[str, Any]												=		get_training_setup(__SCREAMING_SNAKE_CASE )
  # Use a single batch
  __A		:     Optional[int]												=		next(iter(__SCREAMING_SNAKE_CASE ) ).values()
  for iteration in range(3 ):
    # Gather the distributed inputs and targs for the base model
    __A		:     Tuple												=		accelerator.gather((ddp_input, ddp_target) )
    __A		:     str												=		input.to(accelerator.device ), target.to(accelerator.device )
    # Perform our initial ground truth step in non "DDP"
    step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    # Do "gradient accumulation" (noop)
    if iteration % 2 == 0:
      # Accumulate grads locally
      with accelerator.no_sync(__SCREAMING_SNAKE_CASE ):
        step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    else:
      # Sync grads
      step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )

    # DDP model and model should only be in sync when not (iteration % 2 == 0)
    for param, ddp_param in zip(model.parameters()			,     ddp_model.parameters() ):
      if not param.requires_grad:
        continue
      if iteration % 2 == 0:
        # Grads should not be in sync
        assert (
            torch.allclose(param.grad			,     ddp_param.grad ) is False
        ), F"Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"
      else:
        # Grads should be in sync
        assert (
            torch.allclose(param.grad			,     ddp_param.grad ) is True
        ), F"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"

        # Shuffle ddp_input on each iteration
    torch.manual_seed(1337 + iteration )
    __A		:     Optional[Any]												=		ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE ) )]







def        A_     (						__SCREAMING_SNAKE_CASE						:  List[Any]=False			,     __SCREAMING_SNAKE_CASE						:  str=False ):





  """simple docstring"""

  __A		:     Any												=		Accelerator(
      split_batches=__SCREAMING_SNAKE_CASE			,     dispatch_batches=__SCREAMING_SNAKE_CASE			,     gradient_accumulation_steps=2 )
  # Test that context manager behaves properly
  __A		:     Tuple												=		get_training_setup(__SCREAMING_SNAKE_CASE )
  for iteration, batch in enumerate(__SCREAMING_SNAKE_CASE ):
    __A		:     Dict												=		batch.values()
    # Gather the distributed inputs and targs for the base model
    __A		:     Dict												=		accelerator.gather((ddp_input, ddp_target) )
    __A		:     Any												=		input.to(accelerator.device ), target.to(accelerator.device )
    # Perform our initial ground truth step in non "DDP"
    step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    # Do "gradient accumulation" (noop)
    with accelerator.accumulate(__SCREAMING_SNAKE_CASE ):
      step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )

    # DDP model and model should only be in sync when not (iteration % 2 == 0)
    for param, ddp_param in zip(model.parameters()			,     ddp_model.parameters() ):
      if not param.requires_grad:
        continue
      if ((iteration + 1) % 2 == 0) or (iteration == len(__SCREAMING_SNAKE_CASE ) - 1):
        # Grads should be in sync
        assert (
            torch.allclose(param.grad			,     ddp_param.grad ) is True
        ), F"Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"
      else:
        # Grads should not be in sync
        assert (
            torch.allclose(param.grad			,     ddp_param.grad ) is False
        ), F"Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"

        # Shuffle ddp_input on each iteration
    torch.manual_seed(1337 + iteration )
    __A		:     Any												=		ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE ) )]
  GradientState._reset_state()







def        A_     (						__SCREAMING_SNAKE_CASE						:  List[str]=False			,     __SCREAMING_SNAKE_CASE						:  Tuple=False ):





  """simple docstring"""

  __A		:     Tuple												=		Accelerator(
      split_batches=__SCREAMING_SNAKE_CASE			,     dispatch_batches=__SCREAMING_SNAKE_CASE			,     gradient_accumulation_steps=2 )
  # Test that context manager behaves properly
  __A		:     int												=		get_training_setup(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
  for iteration, batch in enumerate(__SCREAMING_SNAKE_CASE ):
    __A		:     Optional[int]												=		batch.values()
    # Gather the distributed inputs and targs for the base model
    __A		:     str												=		accelerator.gather((ddp_input, ddp_target) )
    __A		:     Optional[int]												=		input.to(accelerator.device ), target.to(accelerator.device )
    # Perform our initial ground truth step in non "DDP"
    model.train()
    ddp_model.train()
    step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    opt.step()

    if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(__SCREAMING_SNAKE_CASE )):
      if split_batches:
        sched.step()
      else:
        for _ in range(accelerator.num_processes ):
          sched.step()
    opt.zero_grad()
    # Perform gradient accumulation under wrapper
    with accelerator.accumulate(__SCREAMING_SNAKE_CASE ):
      step_model(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
      ddp_opt.step()
      ddp_sched.step()
      ddp_opt.zero_grad()

    # Learning rates should be the same
    assert (
        opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
    ), F"Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n"
    __A		:     Any												=		(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(__SCREAMING_SNAKE_CASE ))
    if accelerator.num_processes > 1:
      check_model_parameters(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
    # Shuffle ddp_input on each iteration
    torch.manual_seed(1337 + iteration )
  GradientState._reset_state()







def        A_     (						):





  """simple docstring"""

  __A		:     Tuple												=		Accelerator()

  __A		:     Optional[Any]												=		RegressionDataset(length=80 )
  __A		:     int												=		DataLoader(__SCREAMING_SNAKE_CASE			,     batch_size=16 )
  __A		:     int												=		RegressionDataset(length=96 )
  __A		:     str												=		DataLoader(__SCREAMING_SNAKE_CASE			,     batch_size=16 )
  __A		:     Any												=		accelerator.prepare(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
  assert accelerator.gradient_state.active_dataloader is None
  for iteration, _ in enumerate(__SCREAMING_SNAKE_CASE ):
    assert id(accelerator.gradient_state.active_dataloader ) == id(__SCREAMING_SNAKE_CASE )
    if iteration < len(__SCREAMING_SNAKE_CASE ) - 1:
      assert not accelerator.gradient_state.end_of_dataloader
      if iteration == 1:
        for batch_num, _ in enumerate(__SCREAMING_SNAKE_CASE ):
          assert id(accelerator.gradient_state.active_dataloader ) == id(__SCREAMING_SNAKE_CASE )
          if batch_num < len(__SCREAMING_SNAKE_CASE ) - 1:
            assert not accelerator.gradient_state.end_of_dataloader
          else:
            assert accelerator.gradient_state.end_of_dataloader
    else:
      assert accelerator.gradient_state.end_of_dataloader
  assert accelerator.gradient_state.active_dataloader is None







def        A_     (						):





  """simple docstring"""

  __A		:     Union[str, Any]												=		Accelerator()
  __A		:     Union[str, Any]												=		accelerator.state
  if state.local_process_index == 0:
    print("""**Test `accumulate` gradient accumulation with dataloader break**""" )
  test_dataloader_break()
  if state.distributed_type == DistributedType.NO:
    if state.local_process_index == 0:
      print("""**Test NOOP `no_sync` context manager**""" )
    test_noop_sync(__SCREAMING_SNAKE_CASE )
  if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
    if state.local_process_index == 0:
      print("""**Test Distributed `no_sync` context manager**""" )
    test_distributed_sync(__SCREAMING_SNAKE_CASE )
  if state.distributed_type == DistributedType.MULTI_GPU:
    for split_batch in [True, False]:
      for dispatch_batches in [True, False]:
        if state.local_process_index == 0:
          print(
              """**Test `accumulate` gradient accumulation, """			,     F"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"			,     )
        test_gradient_accumulation(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )

    # Currently will break on torch 2.0 +, need to investigate why
  if is_torch_version("""<"""			,     """2.0""" ) or state.distributed_type == DistributedType.NO:
    if state.local_process_index == 0:
      print(
          """**Test `accumulate` gradient accumulation with optimizer and scheduler, """			,     """`split_batches=False`, `dispatch_batches=False`**"""			,     )
    test_gradient_accumulation_with_opt_and_scheduler()
    if state.distributed_type == DistributedType.MULTI_GPU:
      for split_batch in [True, False]:
        for dispatch_batches in [True, False]:
          if not split_batch and not dispatch_batches:
            continue
          if state.local_process_index == 0:
            print(
                """**Test `accumulate` gradient accumulation with optimizer and scheduler, """			,     F"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"			,     )
          test_gradient_accumulation_with_opt_and_scheduler(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )







def        A_     (						__SCREAMING_SNAKE_CASE						:  int ):





  """simple docstring"""

  main()


if __name__ == "__main__":
     main()





'''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        A_     (						__SCREAMING_SNAKE_CASE						:  List[Any]			,     __SCREAMING_SNAKE_CASE						:  Tuple )				->       Any:





  """simple docstring"""

  __A		:     str												=		XCLIPTextConfig()

  # derive patch size from model name
  __A		:     List[str]												=		model_name.find("""patch""" )
  __A		:     int												=		int(model_name[start_idx + len("""patch""" ) : start_idx + len("""patch""" ) + 2] )
  __A		:     Union[str, Any]												=		XCLIPVisionConfig(patch_size=__SCREAMING_SNAKE_CASE			,     num_frames=__SCREAMING_SNAKE_CASE )

  if "large" in model_name:
    __A		:     Optional[int]												=		768
    __A		:     Tuple												=		3072
    __A		:     Tuple												=		12

    __A		:     int												=		1024
    __A		:     Optional[int]												=		4096
    __A		:     Union[str, Any]												=		16
    __A		:     Union[str, Any]												=		24
    __A		:     Dict												=		768
    __A		:     str												=		3072

  if model_name == "xclip-large-patch14-16-frames":
    __A		:     Optional[Any]												=		336

  __A		:     int												=		XCLIPConfig.from_text_vision_configs(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )

  if "large" in model_name:
    __A		:     int												=		768

  return config







def        A_     (						__SCREAMING_SNAKE_CASE						:  Dict )				->       Any:





  """simple docstring"""

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

  return name







def        A_     (						__SCREAMING_SNAKE_CASE						:  List[Any]			,     __SCREAMING_SNAKE_CASE						:  Tuple )				->       Optional[Any]:





  """simple docstring"""

  for key in orig_state_dict.copy().keys():
    __A		:     int												=		orig_state_dict.pop(__SCREAMING_SNAKE_CASE )

    if "attn.in_proj" in key:
      __A		:     Dict												=		key.split(""".""" )
      if key.startswith("""visual""" ):
        __A		:     List[str]												=		key_split[3]
        __A		:     List[str]												=		config.vision_config.hidden_size
        if "message_attn" in key:
          if "weight" in key:
            __A		:     str												=		val[
                :dim, :
            ]
            __A		:     str												=		val[
                dim : dim * 2, :
            ]
            __A		:     int												=		val[
                -dim:, :
            ]
          else:
            __A		:     int												=		val[
                :dim
            ]
            __A		:     Optional[Any]												=		val[
                dim : dim * 2
            ]
            __A		:     Any												=		val[
                -dim:
            ]
        else:
          if "weight" in key:
            __A		:     Union[str, Any]												=		val[
                :dim, :
            ]
            __A		:     Optional[Any]												=		val[
                dim : dim * 2, :
            ]
            __A		:     str												=		val[
                -dim:, :
            ]
          else:
            __A		:     Dict												=		val[:dim]
            __A		:     Union[str, Any]												=		val[
                dim : dim * 2
            ]
            __A		:     List[str]												=		val[-dim:]
      elif key.startswith("""mit""" ):
        __A		:     Union[str, Any]												=		key_split[2]
        __A		:     List[Any]												=		config.vision_config.mit_hidden_size
        if "weight" in key:
          __A		:     Dict												=		val[:dim, :]
          __A		:     Any												=		val[dim : dim * 2, :]
          __A		:     int												=		val[-dim:, :]
        else:
          __A		:     Any												=		val[:dim]
          __A		:     Union[str, Any]												=		val[dim : dim * 2]
          __A		:     Union[str, Any]												=		val[-dim:]
      else:
        __A		:     int												=		key_split[2]
        __A		:     Dict												=		config.text_config.hidden_size
        if "weight" in key:
          __A		:     int												=		val[:dim, :]
          __A		:     List[Any]												=		val[
              dim : dim * 2, :
          ]
          __A		:     Tuple												=		val[-dim:, :]
        else:
          __A		:     str												=		val[:dim]
          __A		:     Dict												=		val[
              dim : dim * 2
          ]
          __A		:     Optional[Any]												=		val[-dim:]
    else:
      __A		:     Union[str, Any]												=		rename_key(__SCREAMING_SNAKE_CASE )
      if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
        __A		:     Union[str, Any]												=		val.T
      __A		:     List[Any]												=		val

  return orig_state_dict







def        A_     (						__SCREAMING_SNAKE_CASE						:  Any )				->       List[str]:





  """simple docstring"""

  if num_frames == 8:
    __A		:     str												=		"""eating_spaghetti_8_frames.npy"""
  elif num_frames == 16:
    __A		:     Optional[int]												=		"""eating_spaghetti.npy"""
  elif num_frames == 32:
    __A		:     Dict												=		"""eating_spaghetti_32_frames.npy"""
  __A		:     int												=		hf_hub_download(
      repo_id="""hf-internal-testing/spaghetti-video"""			,     filename=__SCREAMING_SNAKE_CASE			,     repo_type="""dataset"""			,     )
  __A		:     List[Any]												=		np.load(__SCREAMING_SNAKE_CASE )
  return list(__SCREAMING_SNAKE_CASE )







def        A_     (						__SCREAMING_SNAKE_CASE						:  str			,     __SCREAMING_SNAKE_CASE						:  Dict=None			,     __SCREAMING_SNAKE_CASE						:  List[Any]=False )				->       Optional[Any]:





  """simple docstring"""

  __A		:     Optional[int]												=		{
      # 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		:     Tuple												=		model_to_url[model_name]
  __A		:     Optional[int]												=		8
  if "16-frames" in model_name:
    __A		:     Optional[Any]												=		16
  elif "shot" in model_name:
    __A		:     int												=		32

  __A		:     Any												=		get_xclip_config(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
  __A		:     Any												=		XCLIPModel(__SCREAMING_SNAKE_CASE )
  model.eval()

  if "drive" in checkpoint_url:
    __A		:     Optional[int]												=		"""pytorch_model.bin"""
    gdown.cached_download(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     quiet=__SCREAMING_SNAKE_CASE )
    __A		:     int												=		torch.load(__SCREAMING_SNAKE_CASE			,     map_location="""cpu""" )["""model"""]
  else:
    __A		:     str												=		torch.hub.load_state_dict_from_url(__SCREAMING_SNAKE_CASE )["""model"""]

  __A		:     int												=		convert_state_dict(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )

  __A		:     Optional[int]												=		XCLIPModel(__SCREAMING_SNAKE_CASE )
  __A							,					__A		:     Any												=		model.load_state_dict(__SCREAMING_SNAKE_CASE			,     strict=__SCREAMING_SNAKE_CASE )
  assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
  model.eval()

  __A		:     List[Any]												=		336 if model_name == """xclip-large-patch14-16-frames""" else 224
  __A		:     Any												=		VideoMAEImageProcessor(size=__SCREAMING_SNAKE_CASE )
  __A		:     Any												=		CLIPTokenizer.from_pretrained("""openai/clip-vit-base-patch32""" )
  __A		:     List[Any]												=		CLIPTokenizerFast.from_pretrained("""openai/clip-vit-base-patch32""" )
  __A		:     int												=		XCLIPProcessor(image_processor=__SCREAMING_SNAKE_CASE			,     tokenizer=__SCREAMING_SNAKE_CASE )

  __A		:     Optional[int]												=		prepare_video(__SCREAMING_SNAKE_CASE )
  __A		:     Union[str, Any]												=		processor(
      text=["""playing sports""", """eating spaghetti""", """go shopping"""]			,     videos=__SCREAMING_SNAKE_CASE			,     return_tensors="""pt"""			,     padding=__SCREAMING_SNAKE_CASE )

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

  with torch.no_grad():
    __A		:     str												=		model(**__SCREAMING_SNAKE_CASE )

  # Verify outputs
  __A		:     Union[str, Any]												=		outputs.logits_per_video
  __A		:     Optional[int]												=		logits_per_video.softmax(dim=1 )
  print("""Probs:"""			,     __SCREAMING_SNAKE_CASE )
  # kinetics-400
  if model_name == "xclip-base-patch32":
    __A		:     Optional[Any]												=		torch.tensor([[0.0_019, 0.9_951, 0.0_030]] )
  elif model_name == "xclip-base-patch32-16-frames":
    __A		:     Dict												=		torch.tensor([[7.0_999e-04, 9.9_883e-01, 4.5_580e-04]] )
  elif model_name == "xclip-base-patch16":
    __A		:     Optional[int]												=		torch.tensor([[0.0_083, 0.9_681, 0.0_236]] )
  elif model_name == "xclip-base-patch16-16-frames":
    __A		:     Optional[Any]												=		torch.tensor([[7.6_937e-04, 9.9_728e-01, 1.9_473e-03]] )
  elif model_name == "xclip-large-patch14":
    __A		:     List[Any]												=		torch.tensor([[0.0_062, 0.9_864, 0.0_075]] )
  elif model_name == "xclip-large-patch14-16-frames":
    __A		:     Dict												=		torch.tensor([[3.3_877e-04, 9.9_937e-01, 2.8_888e-04]] )
  # kinetics-600
  elif model_name == "xclip-base-patch16-kinetics-600":
    __A		:     List[Any]												=		torch.tensor([[0.0_555, 0.8_914, 0.0_531]] )
  elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
    __A		:     Optional[int]												=		torch.tensor([[3.8_554e-04, 9.9_929e-01, 3.2_754e-04]] )
  elif model_name == "xclip-large-patch14-kinetics-600":
    __A		:     Optional[Any]												=		torch.tensor([[0.0_036, 0.9_920, 0.0_045]] )
  # few shot
  elif model_name == "xclip-base-patch16-hmdb-2-shot":
    __A		:     Optional[Any]												=		torch.tensor([[7.1_890e-06, 9.9_994e-01, 5.6_559e-05]] )
  elif model_name == "xclip-base-patch16-hmdb-4-shot":
    __A		:     Optional[int]												=		torch.tensor([[1.0_320e-05, 9.9_993e-01, 6.2_435e-05]] )
  elif model_name == "xclip-base-patch16-hmdb-8-shot":
    __A		:     List[str]												=		torch.tensor([[4.1_377e-06, 9.9_990e-01, 9.8_386e-05]] )
  elif model_name == "xclip-base-patch16-hmdb-16-shot":
    __A		:     List[str]												=		torch.tensor([[4.1_347e-05, 9.9_962e-01, 3.3_411e-04]] )
  elif model_name == "xclip-base-patch16-ucf-2-shot":
    __A		:     Any												=		torch.tensor([[8.5_857e-05, 9.9_928e-01, 6.3_291e-04]] )
  elif model_name == "xclip-base-patch16-ucf-4-shot":
    __A		:     Dict												=		torch.tensor([[8.5_857e-05, 9.9_928e-01, 6.3_291e-04]] )
  elif model_name == "xclip-base-patch16-ucf-8-shot":
    __A		:     int												=		torch.tensor([[0.0_027, 0.9_904, 0.0_070]] )
  elif model_name == "xclip-base-patch16-ucf-16-shot":
    __A		:     List[Any]												=		torch.tensor([[9.8_219e-04, 9.9_593e-01, 3.0_863e-03]] )
  # zero shot
  elif model_name == "xclip-base-patch16-zero-shot":
    __A		:     Any												=		torch.tensor([[3.5_082e-04, 9.9_785e-01, 1.7_966e-03]] )
  else:
    raise ValueError(F"Model name {model_name} not supported" )
  assert torch.allclose(__SCREAMING_SNAKE_CASE			,     __SCREAMING_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(__SCREAMING_SNAKE_CASE )

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


if __name__ == "__main__":
     A__   :		Optional[int] 		=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.'
     )

     A__   :		List[str] 		=parser.parse_args()
     convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



