def get_sim(
    vision_proj: torch.Tensor,
    text_proj: torch.Tensor,
    temp=1.0,
    agg_method="mean",
):
    """calculate pair-wise video-text similarity.

    Args:
        vision_proj (torch.Tensor): The vision representation. Shape: [B,T,C].
        text_proj (torch.Tensor): The text representation. Shape: [B,C].
        temp (torch.Tensor): The temperature. Shape: [].

    Returns: The similarity between video and text. Shape: [B,B].

    """
    # vision_proj = F.normalize(vision_proj, dim=-1)
    # text_proj = F.normalize(text_proj, dim=-1)
    vision_proj = F.normalize(vision_proj, dim=-1).to(torch.float32)
    text_proj = F.normalize(text_proj, dim=-1).to(torch.float32)
    if vision_proj.ndim == 3:
        sim_v2t = torch.einsum("mld,nd->mln", vision_proj, text_proj) / temp  # [B, L, B]
        sim_t2v = torch.einsum("nd,mld->nlm", text_proj, vision_proj) / temp  # [B, L, B]
        if agg_method == "mean":
            sim_v2t = sim_v2t.mean(1)
            sim_t2v = sim_t2v.mean(1)
        elif agg_method == "max":
            sim_v2t = sim_v2t.max(1)[0]
            sim_t2v = sim_t2v.max(1)[0]
    elif text_proj.ndim == 3:
        sim_v2t = torch.einsum("nd,mld->nlm", vision_proj, text_proj) / temp  # [B, L, B]
        sim_t2v = torch.einsum("nld,md->nlm", text_proj, vision_proj) / temp  # [B, L, B]
        if agg_method == "mean":
            sim_v2t = sim_v2t.mean(1)
            sim_t2v = sim_t2v.mean(1)
        elif agg_method == "max":
            sim_v2t = sim_v2t.max(1)[0]
            sim_t2v = sim_t2v.max(1)[0]
    else:
        sim_v2t = vision_proj @ text_proj.T / temp
        sim_t2v = sim_v2t.T
    return sim_v2t, sim_t2v
calculate pair-wise video-text similarity.

Args:
vision_proj (torch.Tensor): The vision representation. Shape: [B,T,C].
text_proj (torch.Tensor): The text representation. Shape: [B,C].
temp (torch.Tensor): The temperature. Shape: [].

Returns: The similarity between video and text. Shape: [B,B].
videomamba/video_mm/models/criterions.py