text: # Copyright (c) Facebook, Inc. and its affiliates.
  import collections.abc
  import logging
  from dataclasses import dataclass
  from typing import Any, Dict, List, Optional, Tuple
  
  import torch
  from mmf.common.registry import registry
  from mmf.models.base_model import BaseModel
  from mmf.models.transformers.heads.utils import build_heads_dict
  from mmf.modules.encoders import TransformerEncoder, ViTEncoder
  from mmf.modules.losses import MMFLoss
  from mmf.utils.build import build_encoder
  from mmf.utils.modeling import get_bert_configured_parameters
  from omegaconf import MISSING, OmegaConf
  from torch import nn, Tensor
  
  
  logger = logging.getLogger()
  
  
  class ViLTImageEmbedding(nn.Module):
      """
      Patch embedding used for ViLT.
      https://arxiv.org/pdf/2102.03334.pdf
      Implementation based off
      https://github.com/dandelin/ViLT/blob/master/vilt/modules/vilt_module.py
      Using huggingface ViT modules.
      Can be built with random init or the embeddings weights from an exisiting
      ViT model from huggingface. Model list: availible at
      https://huggingface.co/models?other=vit&sort=downloads
      """
  
      def __init__(
          self,
          random_init: bool = True,
          pretrained_model_name: str = "google/vit-base-patch16-224",
          image_size: Optional[List] = None,
          hidden_dropout_prob: Optional[float] = None,
          hidden_size: Optional[int] = None,
          patch_size: Optional[int] = None,
          num_channels: Optional[int] = None,
          *args,
          **kwargs,
      ):
          super().__init__()
          config = OmegaConf.create(
              {"random_init": random_init, "pretrained_model_name": pretrained_model_name}
          )
          if image_size is not None:
              config.image_size = image_size
          if hidden_dropout_prob is not None:
              config.hidden_dropout_prob = hidden_dropout_prob
          if hidden_size is not None:
              config.hidden_size = hidden_size
          if patch_size is not None:
              config.patch_size = patch_size
          if num_channels is not None:
              config.num_channels = num_channels
  
          encoder = ViTEncoder(config)
          self.embedding = encoder.embeddings
          hidden_size = encoder.hf_config.hidden_size
          self.token_type_embeddings = nn.Embedding(2, hidden_size)
  
      def forward(self, image: Tensor) -> Tensor:
          if image.dim() == 5:
              image = image.permute(1, 0, 2, 3, 4).flatten(start_dim=0, end_dim=1)
  
          img_embeddings = self.embedding(image)
  
          img_segment_ids = torch.ones(
              img_embeddings.size()[:-1],
              dtype=img_embeddings.dtype,
              device=img_embeddings.device,
          ).long()
          img_type_embed = self.token_type_embeddings(img_segment_ids)
          img_embeddings = img_embeddings + img_type_embed
          return img_embeddings
  
  
  class ViLTTextEmbedding(nn.Module):
      def __init__(
          self,
          random_init: bool = True,
          bert_model_name: str = "bert-base-uncased",
          hidden_size: Optional[int] = None,
          max_position_embeddings: Optional[int] = None,
          *args,
          **kwargs,
      ):
  
          super().__init__()
          config = OmegaConf.create(
              {"bert_model_name": bert_model_name, "random_init": random_init}
          )
          if hidden_size is not None:
              config.hidden_size = hidden_size
          if max_position_embeddings is not None:
              config.max_position_embeddings = max_position_embeddings
  
          text_encoder = TransformerEncoder(config)
          self.text_embeddings = text_encoder.embeddings
          # the hidden_size param enables hidden_size overrides
          # if hidden_size is None, hidden_size is loaded
          # from the default hf config for the model
          # the actual size of the embeddings will always be in the encoder configs
          hidden_size = text_encoder.config.hidden_size
          self.token_type_embeddings = nn.Embedding(2, hidden_size)
  
      def forward(self, input_ids: Tensor, segment_ids: Tensor) -> Tensor:
          text_embedding = self.text_embeddings(input_ids, token_type_ids=segment_ids)
          # official vilt repo adds type embeddings twice, once in the bert embeddings
          # and a separate time directly
          text_type_embed = self.token_type_embeddings(segment_ids)
          return text_embedding + text_type_embed
  
  
  @registry.register_model("vilt")
  class ViLT(BaseModel):
      @dataclass
      class Config(BaseModel.Config):
          name: str = "ViLT"
          text_embeddings: Any = MISSING
          image_encoder: Any = MISSING
  
      @classmethod
      def config_path(cls):
          return "configs/models/vilt/defaults.yaml"
  
      def build(self):
          self.text_embeddings = ViLTTextEmbedding(**self.config.text_embeddings)
          self.image_embeddings = ViLTImageEmbedding(**self.config.image_encoder.params)
          self.encoder = build_encoder(self.config.image_encoder)
  
          head_configs = self.config.get("heads", {})
          self.tasks = self.config.get("tasks", head_configs.keys())
          if isinstance(self.tasks, str):
              self.tasks = self.tasks.split(",")
  
          self.losses = nn.ModuleDict()
          self.heads_dict = build_heads_dict(head_configs, self.tasks, self.losses)
          self.modality_keys = self.modality_type = ["text", "image"]
  
      def init_losses(self):
          loss_configs = self.config.get("losses", {})
          for loss_name, loss_config in loss_configs.items():
              self.losses[loss_name] = MMFLoss(loss_config)
  
      def forward(self, sample_list: Dict[str, Tensor]) -> Dict[str, Tensor]:
          text_embedding = self.text_embeddings(
              sample_list["input_ids"], sample_list["segment_ids"]
          )
          image_embedding = self.image_embeddings(sample_list["image"])
          self.preprocess_sample(sample_list, image_embedding)
  
          # Feed through encoder
          embeddings = torch.cat([text_embedding, image_embedding], dim=1)
          attention_mask = self.get_attention_mask(
              sample_list, text_embedding, image_embedding
          )
          sequence, _ = self.encoder(embeddings, attention_mask=attention_mask)
          if sequence.dim() != 3:
              sequence = sequence.unsqueeze(1)
  
          outputs = self.heads_dict(sample_list["dataset_name"], sequence, sample_list)
          return outputs
  
      def preprocess_sample(
          self, sample_list: Dict[str, Tensor], image_embedding: Tensor
      ):
          head_names = self.heads_dict.head_names
          if isinstance(head_names, collections.abc.Mapping):
              head_names = head_names[sample_list["dataset_name"]]
  
          head_string = " ".join(head_names)
          prepare_itm = "itm" in head_string
          prepare_mlm = "mlm" in head_string
  
          if prepare_itm:
              sample_list["itm_labels"] = self._infer_itm_labels(sample_list)
          if prepare_mlm:
              sample_list["mlm_labels"] = self._infer_mlm_labels(
                  sample_list, image_embedding.size()[:-1]
              )
              self._encode_mlm(sample_list, image_embedding)
  
      def get_optimizer_parameters(self, config):
          if hasattr(self.encoder, "get_optimizer_parameters"):
              params = self.encoder.get_optimizer_parameters(config)
          else:
              params = [{"params": self.encoder.parameters()}]
          params += get_bert_configured_parameters(self.text_embeddings)
          params += get_bert_configured_parameters(self.heads_dict)
          params += [{"params": self.image_embeddings.parameters()}]
          return params
  
      def get_attention_mask(
          self,
          sample_list: Dict[str, Tensor],
          text_embedding: Tensor,
          image_embedding: Tensor,
      ) -> Tensor:
          text_mask = getattr(sample_list, "input_mask", None)
          image_mask = getattr(sample_list, "image_mask", None)
  
          if text_mask is None and image_mask is None:
              return None
  
          if text_mask is None:
              text_mask = torch.ones(
                  text_embedding.size()[:-1],
                  dtype=text_embedding.dtype,
                  device=text_embedding.device,
              )
  
          if image_mask is None:
              image_mask = torch.ones(
                  image_embedding.size()[:-1],
                  dtype=image_embedding.dtype,
                  device=image_embedding.device,
              )
  
          attention_mask = torch.cat((text_mask, image_mask), dim=-1)
          return attention_mask
  
      def _infer_itm_labels(self, sample_list: Dict[str, Tensor]) -> Dict[str, Tensor]:
          input_ids = sample_list["input_ids"]
          itm_labels = {}
          if "is_correct" in sample_list:
              itm_labels["is_correct"] = sample_list["is_correct"]
          else:
              itm_labels["is_correct"] = torch.tensor(
                  True, dtype=torch.long, device=input_ids.device
              )
  
          return itm_labels
  
      def _infer_mlm_labels(
          self, sample_list: Dict[str, Tensor], image_embeddings_size: Tuple[int, int]
      ):
          input_ids = sample_list["input_ids"]
          mlm_labels = {}
          current_text_idx = 0
          if "lm_label_ids" in sample_list:
              if sample_list["lm_label_ids"].dim() > 2:
                  mlm_labels["text"] = sample_list["lm_label_ids"][:, current_text_idx]
                  current_text_idx += 1
              else:
                  mlm_labels["text"] = sample_list["lm_label_ids"]
          else:
              mlm_labels["text"] = torch.full(
                  input_ids.size(),
                  fill_value=-1,
                  dtype=torch.long,
                  device=input_ids.device,
              )
          mlm_labels["image"] = torch.full(
              image_embeddings_size,
              fill_value=-1,
              dtype=torch.long,
              device=input_ids.device,
          )
          mlm_labels["combined_labels"] = torch.cat(
              [mlm_labels["text"], mlm_labels["image"]], dim=-1
          )
          return mlm_labels
  
      def _encode_mlm(self, sample_list: Dict[str, Tensor], image_embedding: Tensor):
          assert "lm_label_ids" in sample_list
  
          input_ids = sample_list.get("input_ids_masked", sample_list["input_ids"])
          segment_ids = sample_list["segment_ids"]
          text_embedding = self.text_embeddings(input_ids, segment_ids)
  
          embeddings = torch.cat([image_embedding, text_embedding], dim=1)
          attention_mask = self.get_attention_mask(
              sample_list, text_embedding, image_embedding
          )
          sequence, _ = self.encoder(embeddings, attention_mask=attention_mask)
          if sequence.dim() != 3:
              sequence = sequence.unsqueeze(1)
  
          sample_list["hs_masked_for_mlm"] = sequence
  