code: def build_attention_mask(self, context_length: int, batch_size: int):
      # lazily create causal attention mask, with full attention between the vision tokens
      # pytorch uses additive attention mask; fill with -inf
      mask = torch.empty(context_length, context_length)
      mask.fill_(float("-inf"))
      mask.triu_(1)  # zero out the lower diagonal
      if not self.use_pytorch_mha:
          mask = mask.unsqueeze(0)  # add dummy batch dimension
          mask = mask.expand(batch_size, -1, -1)
      return mask
  
imports: import argparse
  import math
  from typing import Optional, Sequence
  import torch
  from torch import Tensor, nn
  from torch.nn import functional
  from torch.utils.checkpoint import checkpoint
  from cvnets.layers import Dropout, Embedding, PositionalEmbedding, get_normalization_layer
  from cvnets.modules import TransformerEncoder
  from cvnets.text_encoders import TEXT_ENCODER_REGISTRY, BaseTextEncoder
  from utils import logger
label_desc: Use image node_id 5 for calling the TextTransformer obj's underlying member method code with example usage: obj.build_attention_mask(context_length, batch_size) and returns: mask