

"""simple docstring"""





import math
from typing import Optional

import torch
import torch.utils.checkpoint as checkpoint
from torch import nn

from ....modeling_utils import PreTrainedModel
from ....utils import add_start_docstrings, logging
from ...bert.modeling_bert import BertModel
from .configuration_retribert import RetriBertConfig


_a				   =							logging.get_logger(__name__)

_a				   =							[
    'yjernite/retribert-base-uncased',
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]




class       _lowerCAmelCase						(			lowercase     ):







 """simple docstring"""







 __UpperCAmelCase			:      Union[str, Any]     	=			RetriBertConfig
 __UpperCAmelCase			:      int     	=			None
 __UpperCAmelCase			:      Optional[Any]     	=			"retribert"


 def       _lowercase				(    self							:    Tuple,      UpperCAmelCase__							:    List[str]			):
    if isinstance(UpperCAmelCase__,      nn.Linear			):
       module.weight.data.normal_(mean=0.0,      std=self.config.initializer_range			)
       if module.bias is not None:
          module.bias.data.zero_()
    elif isinstance(UpperCAmelCase__,      nn.Embedding			):
       module.weight.data.normal_(mean=0.0,      std=self.config.initializer_range			)
       if module.padding_idx is not None:
          module.weight.data[module.padding_idx].zero_()
    elif isinstance(UpperCAmelCase__,      nn.LayerNorm			):
       module.bias.data.zero_()
       module.weight.data.fill_(1.0			)







_a				   =							r'\n\n    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the\n    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads\n    etc.)\n\n    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.\n    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage\n    and behavior.\n\n    Parameters:\n        config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'




@add_start_docstrings(
    "Bert Based model to embed queries or document for document retrieval."  ,lowercase  ,)
class       _lowerCAmelCase						(			lowercase     ):







 """simple docstring"""







 def __init__(    self							:    Union[str, Any],      UpperCAmelCase__							:    RetriBertConfig			):
    super().__init__(UpperCAmelCase__			)
    __lowercase			=		config.projection_dim

    __lowercase			=		BertModel(UpperCAmelCase__			)
    __lowercase			=		None if config.share_encoders else BertModel(UpperCAmelCase__			)
    __lowercase			=		nn.Dropout(config.hidden_dropout_prob			)
    __lowercase			=		nn.Linear(config.hidden_size,      config.projection_dim,      bias=UpperCAmelCase__			)
    __lowercase			=		nn.Linear(config.hidden_size,      config.projection_dim,      bias=UpperCAmelCase__			)

    __lowercase			=		nn.CrossEntropyLoss(reduction="mean"			)

    # Initialize weights and apply final processing
    self.post_init()


 def       _lowercase				(    self							:    int,      UpperCAmelCase__							:    str,      UpperCAmelCase__							:    Dict,      UpperCAmelCase__							:    Any,      UpperCAmelCase__							:    int=-1,      ):
    # reproduces BERT forward pass with checkpointing
    if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
       return sent_encoder(UpperCAmelCase__,      attention_mask=UpperCAmelCase__			)[1]
    else:
       # prepare implicit variables
       __lowercase			=		input_ids.device
       __lowercase			=		input_ids.size()
       __lowercase			=		torch.zeros(UpperCAmelCase__,      dtype=torch.long,      device=UpperCAmelCase__			)
       __lowercase			=		[None] * sent_encoder.config.num_hidden_layers
       __lowercase			=		sent_encoder.get_extended_attention_mask(
           UpperCAmelCase__,      UpperCAmelCase__			)

       # define function for checkpointing
       def partial_encode(*UpperCAmelCase__							:    str			):
          __lowercase			=		sent_encoder.encoder(
              inputs[0],      attention_mask=inputs[1],      head_mask=UpperCAmelCase__,      )
          __lowercase			=		encoder_outputs[0]
          __lowercase			=		sent_encoder.pooler(UpperCAmelCase__			)
          return pooled_output

       # run embedding layer on everything at once
       __lowercase			=		sent_encoder.embeddings(
           input_ids=UpperCAmelCase__,      position_ids=UpperCAmelCase__,      token_type_ids=UpperCAmelCase__,      inputs_embeds=UpperCAmelCase__			)
       # run encoding and pooling on one mini-batch at a time
       __lowercase			=		[]
       for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size			)			):
          __lowercase			=		embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
          __lowercase			=		extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
          __lowercase			=		checkpoint.checkpoint(UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__			)
          pooled_output_list.append(UpperCAmelCase__			)
       return torch.cat(UpperCAmelCase__,      dim=0			)


 def       _lowercase				(    self							:    int,      UpperCAmelCase__							:    str,      UpperCAmelCase__							:    int=None,      UpperCAmelCase__							:    int=-1,      ):
    __lowercase			=		self.embed_sentences_checkpointed(
        UpperCAmelCase__,      UpperCAmelCase__,      self.bert_query,      UpperCAmelCase__,      )
    return self.project_query(UpperCAmelCase__			)


 def       _lowercase				(    self							:    str,      UpperCAmelCase__							:    Any,      UpperCAmelCase__							:    int=None,      UpperCAmelCase__							:    Optional[Any]=-1,      ):
    __lowercase			=		self.embed_sentences_checkpointed(
        UpperCAmelCase__,      UpperCAmelCase__,      self.bert_query if self.bert_doc is None else self.bert_doc,      UpperCAmelCase__,      )
    return self.project_doc(UpperCAmelCase__			)



 def       _lowercase				(    self							:    int,      UpperCAmelCase__							:    torch.LongTensor,      UpperCAmelCase__							:    Optional[torch.FloatTensor],      UpperCAmelCase__							:    torch.LongTensor,      UpperCAmelCase__							:    Optional[torch.FloatTensor],      UpperCAmelCase__							:    int = -1,      ):
    __lowercase			=		input_ids_query.device
    __lowercase			=		self.embed_questions(UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__			)
    __lowercase			=		self.embed_answers(UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__			)
    __lowercase			=		torch.mm(UpperCAmelCase__,      a_reps.t()			)
    __lowercase			=		self.ce_loss(UpperCAmelCase__,      torch.arange(compare_scores.shape[1]			).to(UpperCAmelCase__			)			)
    __lowercase			=		self.ce_loss(compare_scores.t(),      torch.arange(compare_scores.shape[0]			).to(UpperCAmelCase__			)			)
    __lowercase			=		(loss_qa + loss_aq) / 2
    return loss







"""simple docstring"""





import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling




def 						_A						(						UpperCamelCase_  : str)			->		Optional[Any]:
   '''simple docstring'''


   __lowercase			=		"https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"

   with open(UpperCamelCase_,    "wb") as f:
      with requests.get(UpperCamelCase_,    stream=UpperCamelCase_) as req:
         for chunk in req.iter_content(chunk_size=2048):
            f.write(UpperCamelCase_)




def 						_A						(						)			->		VivitConfig:
   '''simple docstring'''


   __lowercase			=		VivitConfig()

   __lowercase			=		400
   __lowercase			=		"huggingface/label-files"
   __lowercase			=		"kinetics400-id2label.json"

   __lowercase			=		json.load(open(hf_hub_download(UpperCamelCase_,    UpperCamelCase_,    repo_type="dataset"),    "r"))
   __lowercase			=		{int(UpperCamelCase_): v for k, v in idalabel.items()}
   __lowercase			=		idalabel
   __lowercase			=		{v: k for k, v in idalabel.items()}
   return config




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


   __lowercase			=		hf_hub_download(
       repo_id="hf-internal-testing/spaghetti-video",    filename="eating_spaghetti_32_frames.npy",    repo_type="dataset")
   __lowercase			=		np.load(UpperCamelCase_)
   return list(UpperCamelCase_)




def 						_A						(						UpperCamelCase_  : np.ndarray)			->		Optional[int]:
   '''simple docstring'''


   if np.ndim(UpperCamelCase_) == 2:
      return transform_attention_bias(UpperCamelCase_)

   elif np.ndim(UpperCamelCase_) == 3:
      return transform_attention_kernel(UpperCamelCase_)

   else:
      raise Exception(F"""Invalid number of dimesions: {np.ndim(UpperCamelCase_)}""")




def 						_A						(						UpperCamelCase_  : np.ndarray)			->		Optional[Any]:
   '''simple docstring'''


   return current.flatten()




def 						_A						(						UpperCamelCase_  : np.ndarray)			->		Any:
   '''simple docstring'''


   return np.reshape(UpperCamelCase_,    (current.shape[0], current.shape[1] * current.shape[2])).T




def 						_A						(						UpperCamelCase_  : np.ndarray)			->		Any:
   '''simple docstring'''


   return np.reshape(UpperCamelCase_,    (current.shape[0] * current.shape[1], current.shape[2])).T




def 						_A						(						UpperCamelCase_  : Tuple,    UpperCamelCase_  : Optional[Any])			->		Any:
   '''simple docstring'''


   __lowercase			=		state_dict["optimizer"]["target"]["Transformer"][F"""encoderblock_{i}"""]

   __lowercase			=		F"""encoder.layer.{i}."""
   __lowercase			=		{
       prefix + "intermediate.dense.bias": state["MlpBlock_0"]["Dense_0"]["bias"],
       prefix + "intermediate.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_0"]["kernel"]),
       prefix + "output.dense.bias": state["MlpBlock_0"]["Dense_1"]["bias"],
       prefix + "output.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_1"]["kernel"]),
       prefix + "layernorm_before.bias": state["LayerNorm_0"]["bias"],
       prefix + "layernorm_before.weight": state["LayerNorm_0"]["scale"],
       prefix + "layernorm_after.bias": state["LayerNorm_1"]["bias"],
       prefix + "layernorm_after.weight": state["LayerNorm_1"]["scale"],
       prefix
       + "attention.attention.query.bias": transform_attention(
           state["MultiHeadDotProductAttention_0"]["query"]["bias"]),
       prefix
       + "attention.attention.query.weight": transform_attention(
           state["MultiHeadDotProductAttention_0"]["query"]["kernel"]),
       prefix
       + "attention.attention.key.bias": transform_attention(state["MultiHeadDotProductAttention_0"]["key"]["bias"]),
       prefix
       + "attention.attention.key.weight": transform_attention(
           state["MultiHeadDotProductAttention_0"]["key"]["kernel"]),
       prefix
       + "attention.attention.value.bias": transform_attention(
           state["MultiHeadDotProductAttention_0"]["value"]["bias"]),
       prefix
       + "attention.attention.value.weight": transform_attention(
           state["MultiHeadDotProductAttention_0"]["value"]["kernel"]),
       prefix + "attention.output.dense.bias": state["MultiHeadDotProductAttention_0"]["out"]["bias"],
       prefix
       + "attention.output.dense.weight": transform_attention_output_weight(
           state["MultiHeadDotProductAttention_0"]["out"]["kernel"]),
   }

   return new_state




def 						_A						(						UpperCamelCase_  : str)			->		List[Any]:
   '''simple docstring'''


   return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()])




def 						_A						(						UpperCamelCase_  : Union[str, Any],    UpperCamelCase_  : Any=False)			->		Any:
   '''simple docstring'''


   __lowercase			=		get_n_layers(UpperCamelCase_)

   __lowercase			=		OrderedDict()

   __lowercase			=		state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
   __lowercase			=		state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]

   __lowercase			=		np.transpose(
       state_dict["optimizer"]["target"]["embedding"]["kernel"],    (4, 3, 0, 1, 2))
   __lowercase			=		state_dict["optimizer"]["target"]["embedding"]["bias"]

   __lowercase			=		state_dict["optimizer"]["target"]["cls"]
   __lowercase			=		state_dict["optimizer"]["target"]["Transformer"]["posembed_input"][
       "pos_embedding"
   ]

   for i in range(UpperCamelCase_):
      new_state.update(transform_state_encoder_block(UpperCamelCase_,    UpperCamelCase_))

   if classification_head:
      __lowercase			=		{"vivit." + k: v for k, v in new_state.items()}
      __lowercase			=		np.transpose(state_dict["optimizer"]["target"]["output_projection"]["kernel"])
      __lowercase			=		np.transpose(state_dict["optimizer"]["target"]["output_projection"]["bias"])

   return {k: torch.tensor(UpperCamelCase_) for k, v in new_state.items()}




def 						_A						(						)			->		VivitImageProcessor:
   '''simple docstring'''


   __lowercase			=		VivitImageProcessor()

   assert extractor.do_resize is True
   assert extractor.size == {"shortest_edge": 256}
   assert extractor.do_center_crop is True
   assert extractor.crop_size == {"width": 224, "height": 224}
   assert extractor.resample == PILImageResampling.BILINEAR

   # here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
   # one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
   # which effectively means no normalization (and ViViT does not overwrite those when calling this func)
   assert extractor.do_normalize is False
   assert extractor.do_rescale is True
   assert extractor.rescale_factor == 1 / 255

   # zero-centering = True in original implementation
   assert extractor.do_zero_centering is True

   return extractor




def 						_A						(						UpperCamelCase_  : str)			->		List[str]:
   '''simple docstring'''


   __lowercase			=		"checkpoint"

   if not os.path.exists(UpperCamelCase_):
      download_checkpoint(UpperCamelCase_)

   __lowercase			=		restore_checkpoint(UpperCamelCase_,    UpperCamelCase_)
   __lowercase			=		transform_state(UpperCamelCase_,    classification_head=UpperCamelCase_)

   __lowercase			=		get_vivit_config()

   assert config.image_size == 224
   assert config.num_frames == 32

   __lowercase			=		VivitForVideoClassification(UpperCamelCase_)
   model.load_state_dict(UpperCamelCase_)
   model.eval()

   __lowercase			=		get_processor()

   __lowercase			=		prepare_video()
   __lowercase			=		extractor(UpperCamelCase_,    return_tensors="pt")

   __lowercase			=		model(**UpperCamelCase_)

   __lowercase			=		torch.Size([1, 400])
   __lowercase			=		torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658])

   assert outputs.logits.shape == expected_shape
   assert torch.allclose(outputs.logits[0, :5],    UpperCamelCase_,    atol=1E-4), outputs.logits[0, :5]

   model.save_pretrained(UpperCamelCase_)
   extractor.save_pretrained(UpperCamelCase_)


if __name__ == "__main__":
   _a				   =							argparse.ArgumentParser()

   parser.add_argument('--output_model_name', '-o', type=str, help='Output path for the converted HuggingFace model')

   _a				   =							parser.parse_args()
   convert(args.output_model_name)




