Skip to content





class ClassifierHead(Head):
 | def __init__(
 |     self,
 |     vocab: Vocabulary,
 |     seq2vec_encoder: Seq2VecEncoder,
 |     feedforward: Optional[FeedForward] = None,
 |     input_dim: int = None,
 |     dropout: float = None,
 |     num_labels: int = None,
 |     label_namespace: str = "labels"
 | ) -> None

A classification Head. Takes encoded text, gets a single vector out of it, runs an optional feedforward layer on that vector, then classifies it into some label space.

Registered as a Head with name "classifier".


  • vocab : Vocabulary
    Used to get the number of labels, if num_labels is not provided, and to translate label indices to strings in make_output_human_readable.
  • seq2vec_encoder : Seq2VecEncoder
    The input to this module is assumed to be a sequence of encoded vectors. We use a Seq2VecEncoder to compress this into a single vector on which we can perform classification.
  • feedforward : FeedForward, optional (default = None)
    An optional feedforward layer to apply on the pooled output before performing the classification.
  • input_dim : int, optional (default = None)
    We need to know how many dimensions to use for the final classification weight matrix. If you have provided either a seq2vec_encoder or a feedforward module, we can get the correct size from those objects. If you use default values for both of those parameters, then you must provide this parameter, so that we know the size of that encoding.
  • dropout : float, optional (default = None)
    Dropout percentage to use.
  • num_labels : int, optional (default = None)
    Number of labels to project to in classification layer. By default, the classification layer will project to the size of the vocabulary namespace corresponding to labels.
  • label_namespace : str, optional (default = "labels")
    Vocabulary namespace corresponding to labels. By default, we use the "labels" namespace.


class ClassifierHead(Head):
 | ...
 | def forward(
 |     self,
 |     encoded_text: torch.FloatTensor,
 |     encoded_text_mask: torch.BoolTensor,
 |     label: torch.IntTensor = None
 | ) -> Dict[str, torch.Tensor]


class ClassifierHead(Head):
 | ...
 | def make_output_human_readable(
 |     self,
 |     output_dict: Dict[str, torch.Tensor]
 | ) -> Dict[str, torch.Tensor]

Does a simple argmax over the probabilities, converts index to string label, and add "label" key to the dictionary with the result.


class ClassifierHead(Head):
 | ...
 | def get_metrics(self, reset: bool = False) -> Dict[str, float]