class_docstr: This is a PyTorch implementation of the paper Disentangled Representation Learning for Text-Video Retrieval.
  Args:
      base_encoder (`str`):
          CLIP encoder backbone. default: `clip_vit_b32`
      agg_module (`str`):
          Feature aggregation module for video. default: `seqTransf`, choices=[`ndone`, `seqLSTM`, `seqTransf`]
      interaction (`str`):
          Interaction type for retrieval. default: `wti`.
      wti_arch (`int)`:
          Select a architecture for weight branch. default: 2.
      cdcr (`int`):
          Channel decorrelation regularization. default: 3.
      cdcr_alpha1 (`float`):
          Coefficient 1 for channel decorrelation regularization. default: 1.0.
      cdcr_alpha2 (`float`):
          Coefficient 2 for channel decorrelation regularization. default: 0.06.
      cdcr_lambda (`float`):
          Coefficient for channel decorrelation regularization. default: 0.001.
      cross_num_hidden_layers (`int`):
          Number of hidden layers for cross transformer interaction.
class_docstr_tok: ["This","is","a","PyTorch","implementation","of","the","paper","Disentangled","Representation","Learning","for","Text-Video","Retrieval",".","Args",":","base_encoder","(","`","str","`",")",":","CLIP","encoder","backbone",".","default",":","`","clip_vit_b32","`","agg_module","(","`","str","`",")",":","Feature","aggregation","module","for","video",".","default",":","`","seqTransf","`",",","choices=","[","`","ndone","`",",","`","seqLSTM","`",",","`","seqTransf","`","]","interaction","(","`","str","`",")",":","Interaction","type","for","retrieval",".","default",":","`","wti","`",".","wti_arch","(","`","int",")","`",":","Select","a","architecture","for","weight","branch",".","default",":","2.","cdcr","(","`","int","`",")",":","Channel","decorrelation","regularization",".","default",":","3.","cdcr_alpha1","(","`","float","`",")",":","Coefficient","1","for","channel","decorrelation","regularization",".","default",":","1.0.","cdcr_alpha2","(","`","float","`",")",":","Coefficient","2","for","channel","decorrelation","regularization",".","default",":","0.06.","cdcr_lambda","(","`","float","`",")",":","Coefficient","for","channel","decorrelation","regularization",".","default",":","0.001.","cross_num_hidden_layers","(","`","int","`",")",":","Number","of","hidden","layers","for","cross","transformer","interaction","."]
code: def wti_interaction(
      self, text_feat, video_feat, text_mask, video_mask
  ):
      if (
          self.training and torch.cuda.is_available()
      ):  # batch merge here
          text_feat = allgather(text_feat, self.config)
          video_feat = allgather(video_feat, self.config)
          text_mask = allgather(text_mask, self.config)
          video_mask = allgather(video_mask, self.config)
          torch.distributed.barrier()  # force sync
  
      if self.config.interaction == "wti":
          text_weight = self.text_weight_fc(text_feat).squeeze(
              2
          )  # bd x nd_t x dd -> bd x nd_t
          text_weight.masked_fill_(
              torch.tensor((1 - text_mask), dtype=torch.bool),
              float("-inf"),
          )  # pylint: disable=not-callable
          text_weight = torch.softmax(text_weight, dim=-1)  # bd x nd_t
  
          video_weight = self.video_weight_fc(video_feat).squeeze(
              2
          )  # bd x nd_v x dd -> bd x nd_v
          video_weight.masked_fill_(
              torch.tensor((1 - video_mask), dtype=torch.bool),
              float("-inf"),
          )  # pylint: disable=not-callable
          video_weight = torch.softmax(
              video_weight, dim=-1
          )  # bd x nd_v
  
      text_feat = text_feat / text_feat.norm(dim=-1, keepdim=True)
      video_feat = video_feat / video_feat.norm(dim=-1, keepdim=True)
  
      retrieve_logits = torch.einsum(
          "atd,bvd->abtv", [text_feat, video_feat]
      )
      retrieve_logits = torch.einsum(
          "abtv,at->abtv", [retrieve_logits, text_mask]
      )
      retrieve_logits = torch.einsum(
          "abtv,bv->abtv", [retrieve_logits, video_mask]
      )
      text_sum = text_mask.sum(-1)
      video_sum = video_mask.sum(-1)
  
      # max for video token
      if self.config.interaction == "ti":  # token-wise interaction
          t2v_logits, max_idx1 = retrieve_logits.max(
              dim=-1
          )  # abtv -> abt
          v2t_logits, max_idx2 = retrieve_logits.max(
              dim=-2
          )  # abtv -> abv
          t2v_logits = torch.sum(t2v_logits, dim=2) / (
              text_sum.unsqueeze(1)
          )
          v2t_logits = torch.sum(v2t_logits, dim=2) / (
              video_sum.unsqueeze(0)
          )
          retrieve_logits = (t2v_logits + v2t_logits) / 2.0
  
      elif (
          self.config.interaction == "wti"
      ):  # weighted token-wise interaction
          t2v_logits, max_idx1 = retrieve_logits.max(
              dim=-1
          )  # abtv -> abt
          t2v_logits = torch.einsum(
              "abt,at->ab", [t2v_logits, text_weight]
          )
  
          v2t_logits, max_idx2 = retrieve_logits.max(
              dim=-2
          )  # abtv -> abv
          v2t_logits = torch.einsum(
              "abv,bv->ab", [v2t_logits, video_weight]
          )
          retrieve_logits = (t2v_logits + v2t_logits) / 2.0
  
      if self.training:
          logit_scale = self.clip.logit_scale.exp()
          retrieve_logits = logit_scale * retrieve_logits
  
          if self.config.cdcr == 1:
              # simple random
              text_feat = text_feat[
                  torch.arange(text_feat.shape[0]),
                  torch.randint_like(text_sum, 0, 10000) % text_sum,
                  :,
              ]
              video_feat = video_feat[
                  torch.arange(video_feat.shape[0]),
                  torch.randint_like(video_sum, 0, 10000) % video_sum,
                  :,
              ]
              z_a_norm = (
                  text_feat - text_feat.mean(0)
              ) / text_feat.std(
                  0
              )  # ndxnd_sxdd
              z_b_norm = (
                  video_feat - video_feat.mean(0)
              ) / video_feat.std(
                  0
              )  # ndxnd_txdd
  
              # cross-correlation matrix
              bd, dd = z_a_norm.shape
              c = (
                  torch.einsum("ac,ad->cd", z_a_norm, z_b_norm) / bd
              )  # ddxdd
              # loss
              on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
              off_diag = (
                  c.flatten()[1:]
                  .view(dd - 1, dd + 1)[:, :-1]
                  .pow_(2)
                  .sum()
              )
              cdcr_loss = (
                  on_diag * self.config.cdcr_alpha1
                  + off_diag * self.config.cdcr_alpha2
              )
              return retrieve_logits, retrieve_logits.T, cdcr_loss
          elif self.config.cdcr == 2:
              # selecet max
              max_idx1 = max_idx1[
                  torch.arange(max_idx1.shape[0]),
                  torch.arange(max_idx1.shape[1]),
              ]
              max_idx2 = max_idx2[
                  torch.arange(max_idx2.shape[0]),
                  torch.arange(max_idx2.shape[1]),
              ]
  
              max_t_feat = text_feat[
                  torch.arange(max_idx2.shape[0]).repeat_interleave(
                      max_idx2.shape[1]
                  ),
                  max_idx2.flatten(),
              ]
              max_v_feat = video_feat[
                  torch.arange(max_idx1.shape[0]).repeat_interleave(
                      max_idx1.shape[1]
                  ),
                  max_idx1.flatten(),
              ]
  
              t_feat = text_feat.reshape(-1, text_feat.shape[-1])
              t_mask = text_mask.flatten().type(torch.bool)
              v_feat = video_feat.reshape(-1, text_feat.shape[-1])
              v_mask = video_mask.flatten().type(torch.bool)
              t_feat = t_feat[t_mask]
              v_feat = v_feat[v_mask]
              max_t_feat = max_t_feat[v_mask]
              max_v_feat = max_v_feat[t_mask]
  
              z_a_norm = (t_feat - t_feat.mean(0)) / t_feat.std(
                  0
              )  # (bdxnd_t)xdd
              z_b_norm = (
                  max_v_feat - max_v_feat.mean(0)
              ) / max_v_feat.std(
                  0
              )  # (bdxnd_t)xdd
  
              x_a_norm = (v_feat - v_feat.mean(0)) / v_feat.std(
                  0
              )  # (bdxnd_v)xdd
              x_b_norm = (
                  max_t_feat - max_t_feat.mean(0)
              ) / max_t_feat.std(
                  0
              )  # (bdxnd_v)xdd
  
              # cross-correlation matrix
              nd, dd = z_a_norm.shape
              c1 = (
                  torch.einsum("ac,ad->cd", z_a_norm, z_b_norm) / nd
              )  # ddxdd
              nd, dd = x_a_norm.shape
              c2 = (
                  torch.einsum("ac,ad->cd", x_a_norm, x_b_norm) / nd
              )  # ddxdd
              c = (c1 + c2) / 2.0
              # loss
              on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
              off_diag = (
                  c.flatten()[1:]
                  .view(dd - 1, dd + 1)[:, :-1]
                  .pow_(2)
                  .sum()
              )
              cdcr_loss = (
                  on_diag * self.config.cdcr_alpha1
                  + off_diag * self.config.cdcr_alpha2
              )
              return retrieve_logits, retrieve_logits.T, cdcr_loss
          elif self.config.cdcr == 3:
              # selecet max
              max_idx1 = max_idx1[
                  torch.arange(max_idx1.shape[0]),
                  torch.arange(max_idx1.shape[1]),
              ]
              max_idx2 = max_idx2[
                  torch.arange(max_idx2.shape[0]),
                  torch.arange(max_idx2.shape[1]),
              ]
  
              max_t_feat = text_feat[
                  torch.arange(max_idx2.shape[0]).repeat_interleave(
                      max_idx2.shape[1]
                  ),
                  max_idx2.flatten(),
              ].squeeze(1)
              max_v_feat = video_feat[
                  torch.arange(max_idx1.shape[0]).repeat_interleave(
                      max_idx1.shape[1]
                  ),
                  max_idx1.flatten(),
              ].squeeze(1)
  
              t_feat = text_feat.reshape(-1, text_feat.shape[-1])
              t_mask = text_mask.flatten().type(torch.bool)
              v_feat = video_feat.reshape(-1, video_feat.shape[-1])
              v_mask = video_mask.flatten().type(torch.bool)
              t_feat = t_feat[t_mask]
              v_feat = v_feat[v_mask]
              max_t_feat = max_t_feat[v_mask]
              max_v_feat = max_v_feat[t_mask]
              text_weight = text_weight.flatten()[t_mask]
              video_weight = video_weight.flatten()[v_mask]
  
              z_a_norm = (t_feat - t_feat.mean(0)) / t_feat.std(
                  0
              )  # (bdxnd_t)xdd
              z_b_norm = (
                  max_v_feat - max_v_feat.mean(0)
              ) / max_v_feat.std(
                  0
              )  # (bdxnd_t)xdd
  
              x_a_norm = (v_feat - v_feat.mean(0)) / v_feat.std(
                  0
              )  # (bdxnd_v)xdd
              x_b_norm = (
                  max_t_feat - max_t_feat.mean(0)
              ) / max_t_feat.std(
                  0
              )  # (bdxnd_v)xdd
  
              # cross-correlation matrix
              nd, dd = z_a_norm.shape
              bd = text_feat.shape[0]
              c1 = (
                  torch.einsum(
                      "acd,a->cd",
                      torch.einsum("ac,ad->acd", z_a_norm, z_b_norm),
                      text_weight,
                  )
                  / bd
              )  # ddxdd
              c2 = (
                  torch.einsum(
                      "acd,a->cd",
                      torch.einsum("ac,ad->acd", x_a_norm, x_b_norm),
                      video_weight,
                  )
                  / bd
              )  # ddxdd
              c = (c1 + c2) / 2.0
              # loss
              on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
              off_diag = (
                  c.flatten()[1:]
                  .view(dd - 1, dd + 1)[:, :-1]
                  .pow_(2)
                  .sum()
              )
              cdcr_loss = (
                  on_diag * self.config.cdcr_alpha1
                  + off_diag * self.config.cdcr_alpha2
              )
              return retrieve_logits, retrieve_logits.T, cdcr_loss
          else:
              return retrieve_logits, retrieve_logits.T, 0.0
      else:
          return retrieve_logits, retrieve_logits.T, 0.0
  
imports: from collections import OrderedDict
  from typing import Optional
  from types import SimpleNamespace
  import torch
  import logging
  from torch import nn
  from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
  from towhee.models import clip
  from towhee.models.drl.until_module import convert_weights
  from towhee.models.drl.module_cross import CrossModel, Transformer
  from towhee.models.drl.until_module import LayerNorm, AllGather, AllGather2, CrossEn
label_desc: Use image node_id 12 for calling the DRL obj's underlying member method code with example usage: obj.wti_interaction(text_feat, video_feat, text_mask, video_mask) and returns: retrieve_logits,  retrieve_logits,  int, retrieve_logits,  retrieve_logits,  cdcr_loss, retrieve_logits,  retrieve_logits,  cdcr_loss, retrieve_logits,  retrieve_logits,  cdcr_loss, retrieve_logits,  retrieve_logits,  int