class_docstr: Adds a optional transformer layer and classification layers on top of BERT.
  Args:
      bert_model: pretrained BERT model
      output_dim: dimension of the output layer for defult 1 linear layer classifier. Either output_dim or classifier_layer must be specified
      add_transformer_layer: if additional transformer layer should be added on top of the pretrained model
      layer_pulled: which layer should be pulled from pretrained model
      aggregation: embeddings aggregation (pooling) strategy. Available options are:
          (default)"first" - [CLS] representation,
          "mean" - average of all embeddings except CLS,
          "max" - max of all embeddings except CLS
      classifier_layer: classification layers, can be a signle layer, or list of layers (for ex, torch.nn.Sequential)
docstr: Forward pass.
docstr_tok: ["Forward","pass","."]
code: def forward(self, token_ids, segment_ids, attention_mask):
      """
      Forward pass.
      """
      output_bert, output_pooler = self.bert_model(
          token_ids, segment_ids, attention_mask
      )
      # output_bert is a list of 12 (for bert base) layers.
      layer_of_interest = output_bert[self.layer_pulled]
      dtype = next(self.parameters()).dtype
      if self.add_transformer_layer:
          # Follow up by yet another transformer layer
          extended_attention_mask = attention_mask.unsqueeze(
              1
          ).unsqueeze(2)
          extended_attention_mask = (~extended_attention_mask).to(
              dtype
          ) * neginf(dtype)
          embedding_layer = self.additional_transformer_layer(
              layer_of_interest, extended_attention_mask
          )
      else:
          embedding_layer = layer_of_interest
  
      if self.aggregation == "mean":
          #  consider the average of all the output except CLS.
          # obviously ignores masked elements
          outputs_of_interest = embedding_layer[:, 1:, :]
          mask = (
              attention_mask[:, 1:]
              .type_as(embedding_layer)
              .unsqueeze(2)
          )
          sumed_embeddings = torch.sum(
              outputs_of_interest * mask, dim=1
          )
          nb_elems = torch.sum(
              attention_mask[:, 1:].type(dtype), dim=1
          ).unsqueeze(1)
          embeddings = sumed_embeddings / nb_elems
      elif self.aggregation == "max":
          #  consider the max of all the output except CLS
          outputs_of_interest = embedding_layer[:, 1:, :]
          mask = (~attention_mask[:, 1:]).type(dtype).unsqueeze(
              2
          ) * neginf(dtype)
          embeddings, _ = torch.max(outputs_of_interest + mask, dim=1)
      else:
          # easiest, we consider the output of "CLS" as the embedding
          embeddings = embedding_layer[:, 0, :]
  
      # We need this in case of dimensionality reduction
      result = self.additional_linear_layer(embeddings)
  
      # Sort of hack to make it work with distributed: this way the pooler layer
      # is used for grad computation, even though it does not change anything...
      # in practice, it just adds a very (768*768) x (768*batchsize) matmul
      result = result + 0 * torch.sum(output_pooler)
      return result
  
file_path: ParlAI/parlai/agents/bert_ranker/helpers.py
imports: from parlai.core.torch_ranker_agent import TorchRankerAgent
  from parlai.utils.torch import neginf
  import torch
label_desc: Use image node_id 2 for calling the BertWrapper obj's underlying member method code with example usage: obj.forward(token_ids, segment_ids, attention_mask) and returns: result