code: def extract_bert_hidden_states(
      all_encoder_layers, bert_token_to_orig_map, weighted_avg=False
  ):
      """Extract BERT hidden states.
  
      Parameters
      ----------
      all_encoder_layers : torch.Tensor
          All layer outputs of the BERT encoder.
      bert_token_to_orig_map : torch.Tensor
          The mapping between BERT wordpiece tokens to original tokens.
      weighted_avg : boolean
        Specify whether to compute the weighted average (with leranable weight vector) or average of
        the embeddings of all hidden layers, default: ``False``.
  
      Returns
      -------
      torch.Tensor
          The output BERT embeddings.
      """
      (
          num_layers,
          batch_size,
          num_chunk,
          max_token_len,
          bert_dim,
      ) = all_encoder_layers.shape
      out_features = torch.matmul(
          bert_token_to_orig_map.unsqueeze(0).transpose(-1, -2),
          all_encoder_layers,
      )
      out_features = torch.sum(out_features, 2) / torch.clamp(
          torch.sum(bert_token_to_orig_map, (1, 2)), min=1
      ).unsqueeze(-1).unsqueeze(0)
  
      # Average through all layers
      if not weighted_avg:
          out_features = torch.mean(out_features, 0)
      return out_features
  
file_path: graph4nlp/graph4nlp/pytorch/modules/utils/bert_utils.py
imports: from collections import namedtuple
  import numpy
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: extract_bert_hidden_states(all_encoder_layers, bert_token_to_orig_map, weighted_avg) and returns: out_features