class_docstr: Class for obtaining the BERT output for CLS-token and special entity token
  Args:
      pretrained_bert: pretrained Bert checkpoint path or key title (e.g. "bert-base-uncased")
      bert_config_file: path to Bert configuration file, or None, if `pretrained_bert` is a string name
      device: device to use
code: def forward(
      self,
      input_ids: Tensor,
      attention_mask: Tensor,
      entity_tokens_pos: List[int] = None,
  ) -> Union[Tuple[Any, Tensor], Tuple[Tensor]]:
      if entity_tokens_pos is not None:
          q_outputs = self.encoder(
              input_ids=input_ids, attention_mask=attention_mask
          )
          q_hidden_states = q_outputs.last_hidden_state
  
          entity_emb = []
          for i in range(len(entity_tokens_pos)):
              pos = entity_tokens_pos[i]
              entity_emb.append(q_hidden_states[i, pos])
  
          entity_emb = torch.stack(entity_emb, dim=0).to(self.device)
          return entity_emb
      else:
          c_outputs = self.encoder(
              input_ids=input_ids, attention_mask=attention_mask
          )
          c_cls_emb = c_outputs.last_hidden_state[:, :1, :].squeeze(1)
          return c_cls_emb
  
imports: from logging import getLogger
  from pathlib import Path
  from typing import List, Optional, Dict, Tuple, Union, Any
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from transformers import AutoConfig, AutoTokenizer, AutoModel
  from deeppavlov.core.commands.utils import expand_path
  from deeppavlov.core.common.errors import ConfigError
  from deeppavlov.core.common.registry import register
  from deeppavlov.core.models.torch_model import TorchModel
  from deeppavlov.models.preprocessors.torch_transformers_preprocessor import TorchTransformersEntityRankerPreprocessor
label_desc: Use image node_id 2 for calling the TextEncoder obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, entity_tokens_pos) and returns: entity_emb, c_cls_emb