code: '''
   * Copyright (c) 2021, salesforce.com, inc.
   * All rights reserved.
   * SPDX-License-Identifier: BSD-3-Clause
   * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
  '''
  
  from functools import partial
  from multiprocessing.sharedctypes import Value
  from models.vit import VisionTransformer, interpolate_pos_embed
  from models.med import BertConfig, BertForMaskedLM
  # from models.modality_wise_ln_med import BertConfig, BertForMaskedLM
  from enum import Enum
  from omegaconf import OmegaConf
  
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  import numpy as np
  import random
  
  
  class VisionLanguageLearner(nn.Module):
      def __init__(self,                 
                   text_encoder = None,
                   tokenizer = None,
                   config = None,    
                   temp = 0.07,
                   init_deit = True
                   ):
          super().__init__()
          
          self.tokenizer = tokenizer 
          self.mlm_probability = config['mlm_probability']
          embed_dim = config['embed_dim']
       
          self.visual_encoder = VisionTransformer(
              img_size=config['image_res'], patch_size=16, embed_dim=768, depth=1, num_heads=12, 
              mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))   
          
          if init_deit:
              checkpoint = torch.hub.load_state_dict_from_url(
                  url="https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth",
                  map_location="cpu", check_hash=True)
              state_dict = checkpoint["model"]
              pos_embed_reshaped = interpolate_pos_embed(state_dict['pos_embed'], self.visual_encoder)
              state_dict['pos_embed'] = pos_embed_reshaped
              msg = self.visual_encoder.load_state_dict(state_dict,strict=False)
              print(f'missing_keys={msg.missing_keys}\tunexpected_keys={msg.unexpected_keys}')         
              
          vision_width = config['vision_width']       
          bert_config = BertConfig.from_dict(OmegaConf.to_container(config.bert_config))
          
          self.text_encoder = BertForMaskedLM.from_pretrained(text_encoder, config=bert_config)      
  
          text_width = self.text_encoder.config.hidden_size
          self.vision_proj = nn.Linear(vision_width, embed_dim)
          self.text_proj = nn.Linear(text_width, embed_dim)         
  
          self.temp = nn.Parameter(torch.ones([]) * config['temp'])   
          self.queue_size = config['queue_size']
          self.momentum = config['momentum']  
  
          # Hardcoded from DALL-E's D-VAE.
          vocab_size = 8192
          self.mim_head = nn.Linear(self.visual_encoder.embed_dim, vocab_size)
  
          # create momentum models
          self.text_encoder_m = BertForMaskedLM.from_pretrained(text_encoder, config=bert_config)       
          self.vision_proj_m = nn.Linear(vision_width, embed_dim)
          self.text_proj_m = nn.Linear(text_width, embed_dim)    
          self.visual_encoder_m = VisionTransformer(
              img_size=config['image_res'], patch_size=16, embed_dim=768, depth=1, num_heads=12, 
              mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))   
          
          self.model_pairs = [
              (self.visual_encoder, self.visual_encoder_m),
              (self.text_encoder, self.text_encoder_m),
              (self.text_proj, self.text_proj_m),
              (self.vision_proj, self.vision_proj_m)
          ]
          
          self.copy_params()
  
          # create the queue
          self.register_buffer("image_queue", torch.randn(embed_dim, self.queue_size))
          self.register_buffer("text_queue", torch.randn(embed_dim, self.queue_size))
          self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))  
                               
          self.image_queue = nn.functional.normalize(self.image_queue, dim=0)
          self.text_queue = nn.functional.normalize(self.text_queue, dim=0)
  
  
      def forward(self, image, text, visual_token_ids, masked_visual_token_pos, masked_visual_tok_labels, alpha=0, return_dict=False):
          with torch.no_grad():
              self.temp.clamp_(0.001,0.5)
  
          ## ================ ITA ====================== ##
          image_embeds = self.visual_encoder(image) 
          image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
          image_embeds = self.text_encoder.bert(
                          inputs_embeds=image_embeds, 
                          attention_mask=image_atts,
                          return_dict=True,
                          mode='image'
                      )
          image_embeds = image_embeds.last_hidden_state
          image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)  
  
          text_output = self.text_encoder.bert(text.input_ids, attention_mask = text.attention_mask,                      
                                          return_dict = True, mode = 'text')            
          text_embeds = text_output.last_hidden_state
          text_feat = F.normalize(self.text_proj(text_embeds[:,0,:]),dim=-1)                 
  
          
          with torch.no_grad():
              self._momentum_update()
  
              image_embeds_m = self.visual_encoder_m(image) 
              image_embeds_m = self.text_encoder.bert(
                          inputs_embeds=image_embeds_m, 
                          attention_mask=image_atts,
                          return_dict=True,
                          mode='image'
                      )
              image_embeds_m = image_embeds_m.last_hidden_state
              image_feat_m = F.normalize(self.vision_proj_m(image_embeds_m[:,0,:]),dim=-1)  
              image_feat_all = torch.cat([image_feat_m.t(),self.image_queue.clone().detach()],dim=1)                                         
  
              text_output_m = self.text_encoder_m.bert(text.input_ids, attention_mask = text.attention_mask,                      
                                                  return_dict = True, mode = 'text')    
              text_feat_m = F.normalize(self.text_proj_m(text_output_m.last_hidden_state[:,0,:]),dim=-1) 
              text_feat_all = torch.cat([text_feat_m.t(),self.text_queue.clone().detach()],dim=1)
  
              sim_i2t_m = image_feat_m @ text_feat_all / self.temp
              sim_t2i_m = text_feat_m @ image_feat_all / self.temp
  
              sim_targets = torch.zeros(sim_i2t_m.size()).to(image.device)
              sim_targets.fill_diagonal_(1)
  
              sim_i2t_targets = alpha * F.softmax(sim_i2t_m, dim=1) + (1 - alpha) * sim_targets
              sim_t2i_targets = alpha * F.softmax(sim_t2i_m, dim=1) + (1 - alpha) * sim_targets        
  
          sim_i2t = image_feat @ text_feat_all / self.temp 
          sim_t2i = text_feat @ image_feat_all / self.temp 
  
          loss_i2t = -torch.sum(F.log_softmax(sim_i2t, dim=1)*sim_i2t_targets,dim=1).mean()
          loss_t2i = -torch.sum(F.log_softmax(sim_t2i, dim=1)*sim_t2i_targets,dim=1).mean() 
  
          loss_ita = (loss_i2t+loss_t2i)/2
          self._dequeue_and_enqueue(image_feat_m, text_feat_m)
  
          
          ##================= MLM ========================##                
          input_ids = text.input_ids.clone()
          labels = input_ids.clone()
  
          probability_matrix = torch.full(labels.shape, self.mlm_probability)                    
          input_ids, labels = self.mask(input_ids, self.text_encoder.config.vocab_size, image.device, targets=labels,
                                        probability_matrix = probability_matrix) 
          
          with torch.no_grad():
              logits_m = self.text_encoder_m(input_ids, 
                                             attention_mask = text.attention_mask,
                                             return_dict = True,
                                             return_logits = True,
                                             mode='text'
                                            )    
          mlm_output = self.text_encoder(input_ids, 
                                         attention_mask = text.attention_mask,
                                         return_dict = True,
                                         labels = labels,   
                                         soft_labels = F.softmax(logits_m,dim=-1),
                                         alpha = alpha,
                                         mode='text'
                                        )                           
          loss_mlm = mlm_output.loss        
  
  
          ##================= MIM ========================##
          post_mask_image_embeds = self.visual_encoder(image, masked_visual_token_pos)
          image_atts = torch.ones(post_mask_image_embeds.size()[:-1],dtype=torch.long).to(image.device)
          post_mask_cross_embeds = self.text_encoder.bert(
                          inputs_embeds=post_mask_image_embeds, 
                          attention_mask=image_atts,
                          return_dict=True,
                          mode='image'
                      )
          # Drop the CLS token, because we don't mask it.
          post_mask_cross_embeds = post_mask_cross_embeds.last_hidden_state[:, 1:]
          predicted_visual_tokens = self.mim_head(post_mask_cross_embeds)
          loss_mim = F.cross_entropy(
              input=predicted_visual_tokens[masked_visual_token_pos], 
              target=masked_visual_tok_labels
          )
  
          if return_dict:
              return {
                  'losses': {
                      'loss_ita': loss_ita,
                      'loss_mlm': loss_mlm,
                      'loss_mim': loss_mim
                  }
              }
  
          return loss_mlm, loss_mim, loss_ita
  
          
  
      @torch.no_grad()    
      def copy_params(self):
          for model_pair in self.model_pairs:           
              for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
                  param_m.data.copy_(param.data)  # initialize
                  param_m.requires_grad = False  # not update by gradient    
  
              
      @torch.no_grad()        
      def _momentum_update(self):
          for model_pair in self.model_pairs:           
              for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
                  param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum)
                  
              
              
      @torch.no_grad()
      def _dequeue_and_enqueue(self, image_feat, text_feat):
          # gather keys before updating queue
          image_feats = concat_all_gather(image_feat)
          text_feats = concat_all_gather(text_feat)
  
          batch_size = image_feats.shape[0]
  
          ptr = int(self.queue_ptr)
          assert self.queue_size % batch_size == 0  # for simplicity
  
          # replace the keys at ptr (dequeue and enqueue)
          self.image_queue[:, ptr:ptr + batch_size] = image_feats.T
          self.text_queue[:, ptr:ptr + batch_size] = text_feats.T
          ptr = (ptr + batch_size) % self.queue_size  # move pointer
  
          self.queue_ptr[0] = ptr 
          
          
      def mask(self, input_ids, vocab_size, device, targets=None, masked_indices=None, probability_matrix=None):
          if masked_indices is None:                                       
              masked_indices = torch.bernoulli(probability_matrix).bool()
                                                 
          masked_indices[input_ids == self.tokenizer.pad_token_id] = False
          masked_indices[input_ids == self.tokenizer.cls_token_id] = False
          
          if targets is not None:
              targets[~masked_indices] = -100 # We only compute loss on masked tokens            
  
          # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
          indices_replaced = torch.bernoulli(torch.full(input_ids.shape, 0.8)).bool() & masked_indices
          input_ids[indices_replaced] = self.tokenizer.mask_token_id
  
          # 10% of the time, we replace masked input tokens with random word
          indices_random = torch.bernoulli(torch.full(input_ids.shape, 0.5)).bool() & masked_indices & ~indices_replaced
          random_words = torch.randint(vocab_size, input_ids.shape, dtype=torch.long).to(device)
          input_ids[indices_random] = random_words[indices_random]                     
          # The rest of the time (10% of the time) we keep the masked input tokens unchanged   
          
          if targets is not None:
              return input_ids, targets
          else:
              return input_ids
          
  
  @torch.no_grad()
  def concat_all_gather(tensor):
      """
      Performs all_gather operation on the provided tensors.
      *** Warning ***: torch.distributed.all_gather has no gradient.
      """
      tensors_gather = [torch.ones_like(tensor)
          for _ in range(torch.distributed.get_world_size())]
      torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
  
      output = torch.cat(tensors_gather, dim=0)
      return output
  
  
apis: torch.ones_like torch.distributed.all_gather torch.ones torch.nn.functional.log_softmax torch.nn.Linear torch.hub.load_state_dict_from_url torch.randint torch.distributed.get_world_size torch.randn torch.nn.functional.normalize torch.nn.functional.softmax torch.no_grad torch.full torch.nn.functional.cross_entropy torch.zeros torch.cat torch.bernoulli