class_docstr: CoCa model.
  
  Args:
      dim (`int`):
          model dimension
      img_encoder (`nn.Module`):
          vision transformer - image encoder, returning image embeddings as (batch, seq, dim).
      image_dim (`int`):
          image embedding dimension, if not the same as model dimensions.
      num_tokens (`int`):
          number of text tokens.
      unimodal_depth (`int`):
          depth of the unimodal transformer.
      multimodal_depth (`int`):
          depth of the multimodal transformer.
      dim_head (`int`):
          dimension per attention head.
      heads (`int`):
          number of attention heads.
      caption_loss_weight (`float`):
          weight on the autoregressive caption loss.
      contrastive_loss_weight (`float`):
          weight on the contrastive loss between image and text CLS embeddings.
class_docstr_tok: ["CoCa","model",".","Args",":","dim","(","`","int","`",")",":","model","dimension","img_encoder","(","`","nn.Module","`",")",":","vision","transformer","-","image","encoder",",","returning","image","embeddings","as","(","batch",",","seq",",","dim",")",".","image_dim","(","`","int","`",")",":","image","embedding","dimension",",","if","not","the","same","as","model","dimensions",".","num_tokens","(","`","int","`",")",":","number","of","text","tokens",".","unimodal_depth","(","`","int","`",")",":","depth","of","the","unimodal","transformer",".","multimodal_depth","(","`","int","`",")",":","depth","of","the","multimodal","transformer",".","dim_head","(","`","int","`",")",":","dimension","per","attention","head",".","heads","(","`","int","`",")",":","number","of","attention","heads",".","caption_loss_weight","(","`","float","`",")",":","weight","on","the","autoregressive","caption","loss",".","contrastive_loss_weight","(","`","float","`",")",":","weight","on","the","contrastive","loss","between","image","and","text","CLS","embeddings","."]
code: def forward(
      self,
      text,
      images=None,
      image_tokens=None,
      labels=None,
      return_loss=False,
      return_embeddings=False,
  ):
      batch, device = text.shape[0], text.device
  
      if return_loss and (labels is None):
          text, labels = text[:, :-1], text[:, 1:]
  
      text_embeds, text_tokens = self.embed_text(text)
  
      image_embeds, image_tokens = self.embed_image(
          images=images, image_tokens=image_tokens
      )
  
      # return embeddings if that is what the researcher wants
  
      if return_embeddings:
          return text_embeds, image_embeds
  
      # go through multimodal layers
  
      for attn_ff, cross_attn in self.multimodal_layers:
          text_tokens = attn_ff(text_tokens)
          text_tokens = cross_attn(text_tokens, image_tokens)
  
      logits = self.to_logits(text_tokens)
  
      if not return_loss:
          return logits
  
      # shorthand
  
      ce = F.cross_entropy
  
      # calculate caption loss (cross entropy loss)
  
      logits = rearrange(logits, "b n c -> b c n")
      caption_loss = ce(logits, labels, ignore_index=self.pad_id)
      caption_loss = caption_loss * self.caption_loss_weight
  
      # calculate contrastive loss
  
      sim = einsum("i d, j d -> i j", text_embeds, image_embeds)
      sim = sim * self.temperature.exp()
      contrastive_labels = torch.arange(batch, device=device)
  
      contrastive_loss = (
          ce(sim, contrastive_labels) + ce(sim.t(), contrastive_labels)
      ) * 0.5
      contrastive_loss = contrastive_loss * self.contrastive_loss_weight
  
      return caption_loss + contrastive_loss
  
file_path: towhee/towhee/models/coca/coca.py
imports: import torch
  from torch import einsum, nn
  import torch.nn.functional
  from einops import rearrange, repeat
  from towhee.models.layers.cross_attention import LayerNorm, Residual, RotaryEmbedding, apply_rotary_pos_emb, CrossAttention
  from towhee.models.layers.activations.swiglu import SwiGLU
label_desc: Use image node_id 4 for calling the CoCa obj's underlying member method code with example usage: obj.forward(text, images, image_tokens, labels, return_loss, return_embeddings) and returns: unknown, text_embeds,  image_embeds, logits