class_docstr: BERT model for multiple choice,sequence labeling, ner, classification or regression
  This module is composed of the BERT model with a linear layer on top of
  the pooled output.
  Params:
  task_num_classes
  task_types
  backbone_model - na
class_docstr_tok: ["BERT","model","for","multiple","choice",",","sequence","labeling",",","ner",",","classification","or","regression","This","module","is","composed","of","the","BERT","model","with","a","linear","layer","on","top","of","the","pooled","output",".","Params",":","task_num_classes","task_types","backbone_model","-","na"]
code: def predict_on_top(self, task_id, last_hidden_state, labels=None):
      name = self.task_types[task_id]
      if name == "sequence_labeling":
          #  last hidden state is all token tensor
          final_output = self.dropout(last_hidden_state)
          logits = self.bert.final_classifier[task_id](final_output)
          if labels is not None:
              active_logits = logits.view(-1, self.classes[task_id])
              if self.multilabel[task_id]:
                  loss_fct = BCEWithLogitsLoss()
                  loss = loss_fct(active_logits, labels)
              elif not self.multilabel[task_id]:
                  loss_fct = CrossEntropyLoss()
                  loss = loss_fct(active_logits, labels.view(-1))
              return loss, logits
          else:
              return logits
      elif name in ["classification", "regression", "multiple_choice"]:
          #  last hidden state is a first token tensor
          if self.new_model:  # or True:
              pooled_output = self.bert.pooling_layer(last_hidden_state)
          else:
              pooled_output = self.bert.pooler(last_hidden_state)
          pooled_output = self.activation(pooled_output)
          pooled_output = self.dropout(pooled_output)
          logits = self.bert.final_classifier[task_id](pooled_output)
          if name == "multiple_choice":
              logits = logits.view((-1, self.classes[task_id]))
              if labels is not None:
                  l1, l2 = len(logits), len(labels)
                  if len(logits) != len(labels):
                      raise Exception(
                          f"Len of logits {l1} and labels {l2} not match"
                      )
          if labels is not None:
              if name != "regression":
                  if self.multilabel[task_id]:
                      loss_fct = BCEWithLogitsLoss()
                      loss = loss_fct(logits, labels)
                  elif not self.multilabel[task_id]:
                      if self.focal:
                          if self.weights[task_id] is None:
                              loss_fct = FocalLoss()
                          else:
                              loss_fct = FocalLoss(
                                  weight=torch.tensor(
                                      [self.weights[task_id]]
                                  ).cuda()
                              )
                          loss = loss_fct(logits, labels.view(-1))
                      else:
                          if self.weights[task_id] is None:
                              loss_fct = CrossEntropyLoss()
                          else:
                              loss_fct = CrossEntropyLoss(
                                  weight=torch.Tensor(
                                      [self.weights[task_id]]
                                  ).cuda()
                              )
                          loss = loss_fct(logits, labels.view(-1))
                  return loss, logits
              elif name == "regression":
                  loss_fct = MSELoss()
                  loss = loss_fct(logits, labels.unsqueeze(1))
                  return loss, logits
          else:
              return logits
      elif name == "binary_head":
          last_hidden_state = self.dropout(last_hidden_state)
          pooled_output = self.bert.pooler(last_hidden_state)
          pooled_output = self.activation(pooled_output)
          pooled_output = self.dropout(pooled_output)
          logits = self.bert.final_classifier[task_id](pooled_output)
          if labels is not None:
              if self.focal:
                  if self.weights[task_id] is None:
                      loss_fct = FocalLoss()
                  else:
                      loss_fct = FocalLoss(
                          weight=torch.tensor(
                              [self.weights[task_id]]
                          ).cuda()
                      )
              else:
                  if self.weights[task_id] is None:
                      loss_fct = BCEWithLogitsLoss()
                  else:
                      loss_fct = BCEWithLogitsLoss(
                          weight=torch.Tensor(
                              [self.weights[task_id]]
                          ).cuda()
                      )
              if len(labels.shape) == 1 and len(logits.shape) == 2:
                  labels = labels.unsqueeze(1)
              loss = loss_fct(logits, labels)
              return loss, logits
          else:
              return logits
      else:
          raise Exception(f"Unsupported name {name}")
  
imports: from collections.abc import Iterable
  from logging import getLogger
  from pathlib import Path
  from typing import Dict, Optional
  import numpy
  import torch
  import torch.nn
  from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
  from transformers import AutoConfig, AutoModel
  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.torch_bert.torch_transformers_sequence_tagger import token_from_subtoken, token_labels_to_subtoken_labels
label_desc: Use image node_id 3 for calling the BertForMultiTask obj's underlying member method code with example usage: obj.predict_on_top(task_id, last_hidden_state, labels) and returns: loss,  logits, logits, logits, loss,  logits, loss,  logits, logits, loss,  logits