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 __init__(
      self,
      *,
      dim,
      num_tokens,
      unimodal_depth,
      multimodal_depth,
      image_dim=None,
      num_img_queries=256,
      dim_head=64,
      heads=8,
      ff_mult=4,
      img_encoder=None,
      caption_loss_weight=1.0,
      contrastive_loss_weight=1.0,
      pad_id=0,
  ):
      super().__init__()
      self.dim = dim
  
      self.pad_id = pad_id
      self.caption_loss_weight = caption_loss_weight
      self.contrastive_loss_weight = contrastive_loss_weight
  
      # token embeddings
  
      self.token_emb = nn.Embedding(num_tokens, dim)
      self.text_cls_token = nn.Parameter(torch.randn(dim))
  
      # image encoder
  
      self.img_encoder = img_encoder
  
      # attention pooling for image tokens
  
      self.img_queries = nn.Parameter(
          torch.randn(num_img_queries + 1, dim)
      )
      self.img_attn_pool = CrossAttention(
          dim=dim,
          context_dim=image_dim,
          dim_head=dim_head,
          heads=heads,
          norm_context=True,
      )
  
      self.img_attn_pool_norm = LayerNorm(dim)
      self.text_cls_norm = LayerNorm(dim)
  
      # contrastive learning temperature
  
      self.temperature = nn.Parameter(torch.Tensor([1.0]))
  
      # unimodal layers
  
      self.unimodal_layers = nn.ModuleList([])
      # pylint: disable=W0612
      for ind in range(unimodal_depth):
          self.unimodal_layers.append(
              Residual(
                  ParallelTransformerBlock(
                      dim=dim,
                      dim_head=dim_head,
                      heads=heads,
                      ff_mult=ff_mult,
                  )
              ),
          )
  
      # multimodal layers
  
      self.multimodal_layers = nn.ModuleList([])
      for ind in range(multimodal_depth):
          self.multimodal_layers.append(
              nn.ModuleList(
                  [
                      Residual(
                          ParallelTransformerBlock(
                              dim=dim,
                              dim_head=dim_head,
                              heads=heads,
                              ff_mult=ff_mult,
                          )
                      ),
                      Residual(
                          CrossAttention(
                              dim=dim,
                              dim_head=dim_head,
                              heads=heads,
                              parallel_ff=True,
                              ff_mult=ff_mult,
                          )
                      ),
                  ]
              )
          )
  
      # to logits
  
      self.to_logits = nn.Sequential(
          LayerNorm(dim), nn.Linear(dim, num_tokens, bias=False)
      )
  
      # they used embedding weight tied projection out to logits, not common, but works
      self.to_logits[-1].weight = self.token_emb.weight
      nn.init.normal_(self.token_emb.weight, std=0.02)
  
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 1 to create a new CoCa object from inherited base classes: nn with example: obj = CoCa()