docstr: Returns token embeddings.
  
  :param text_tokens: a tensor of token indices. ([Batch, Seq_len])
  :param key_padding_mask: a tensor of boolean values as the padding mask.
  :param return_all_tokens: a boolean flag to return all tokens, defaults to False
      to return only EOT token embedding.
  :return: a tensor of [Batch, Seq_len, hidden_dim] if return_all_tokens is True,
      otherwise a tensor of [Batch, hidden_dim].
docstr_tok: ["Returns","token","embeddings",".",":","param","text_tokens",":","a","tensor","of","token","indices",".","(","[","Batch",",","Seq_len","]",")",":","param","key_padding_mask",":","a","tensor","of","boolean","values","as","the","padding","mask",".",":","param","return_all_tokens",":","a","boolean","flag","to","return","all","tokens",",","defaults","to","False","to","return","only","EOT","token","embedding",".",":","return",":","a","tensor","of","[","Batch",",","Seq_len",",","hidden_dim","]","if","return_all_tokens","is","True",",","otherwise","a","tensor","of","[","Batch",",","hidden_dim","]","."]
code: def encode_text(
      self,
      text_tokens: Tensor,
      key_padding_mask: Optional[Tensor] = None,
      return_all_tokens: bool = False,
      *args,
      **kwargs,
  ) -> Tensor:
      """
      Returns token embeddings.
  
      :param text_tokens: a tensor of token indices. ([Batch, Seq_len])
      :param key_padding_mask: a tensor of boolean values as the padding mask.
      :param return_all_tokens: a boolean flag to return all tokens, defaults to False
          to return only EOT token embedding.
      :return: a tensor of [Batch, Seq_len, hidden_dim] if return_all_tokens is True,
          otherwise a tensor of [Batch, hidden_dim].
      """
      # discrete tokens to continuous embeddings
      # [Batch, Seq_len] --> [Batch, Seq_len, hidden_dim]
      token_emb = self.forward_embedding(text_tokens)
  
      # [1, Seq_len, Seq_len]
      attn_mask = None
      if self.causal_masking:
          attn_mask = self.build_attention_mask(
              context_length=text_tokens.shape[1],
              batch_size=text_tokens.shape[0],
          )
          attn_mask = attn_mask.to(
              device=token_emb.device, dtype=token_emb.dtype
          )
          key_padding_mask = None
  
      if self.use_pytorch_mha:
          # [Batch, Seq_len, hidden_dim] --> [Seq_len, Batch, hidden_dim]
          # we will use PyTorch's multi-head attention, which uses sequence_first format
          token_emb = token_emb.transpose(0, 1)
  
      for layer in self.transformer:
          if self.gradient_ckpt:
              token_emb = gradient_checkpoint_fn(
                  layer, token_emb, None, key_padding_mask, attn_mask
              )
          else:
              token_emb = layer(
                  token_emb,
                  key_padding_mask=key_padding_mask,
                  attn_mask=attn_mask,
                  use_pytorch_mha=self.use_pytorch_mha,
              )
  
      # Apply layer norm
      token_emb = self.final_layer_norm(token_emb)
  
      if return_all_tokens:
          if self.use_pytorch_mha:
              # [Seq_len, Batch, hidden_dim] --> [Batch, Seq_len, hidden_dim]
              token_emb = token_emb.transpose(0, 1)
  
          return token_emb
  
      # take features from the eot embedding (eot_token is the highest number in each sequence)
      if self.use_pytorch_mha:
          token_emb = token_emb[
              text_tokens.argmax(dim=-1),
              torch.arange(text_tokens.shape[0]),
          ]
      else:
          token_emb = token_emb[
              torch.arange(text_tokens.shape[0]),
              text_tokens.argmax(dim=-1),
          ]
  
      token_emb = token_emb @ self.projection_layer
      # normalize text features
      token_emb = F.normalize(token_emb, dim=-1)
      return token_emb
  
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 6 for calling the TextTransformer obj's underlying member method code with example usage: obj.encode_text(text_tokens, key_padding_mask, return_all_tokens) and returns: token_emb, token_emb