| """Basic model. Predicts tags for every token""" |
| from typing import Dict, Optional, List, Any |
|
|
| import numpy |
| import torch |
| import torch.nn.functional as F |
| from allennlp.data import Vocabulary |
| from allennlp.models.model import Model |
| from allennlp.modules import TimeDistributed, TextFieldEmbedder |
| from allennlp.nn import InitializerApplicator, RegularizerApplicator |
| from allennlp.nn.util import get_text_field_mask, sequence_cross_entropy_with_logits |
| from allennlp.training.metrics import CategoricalAccuracy |
| from overrides import overrides |
| from torch.nn.modules.linear import Linear |
|
|
|
|
| @Model.register("seq2labels") |
| class Seq2Labels(Model): |
| """ |
| This ``Seq2Labels`` simply encodes a sequence of text with a stacked ``Seq2SeqEncoder``, then |
| predicts a tag (or couple tags) for each token in the sequence. |
| |
| Parameters |
| ---------- |
| vocab : ``Vocabulary``, required |
| A Vocabulary, required in order to compute sizes for input/output projections. |
| text_field_embedder : ``TextFieldEmbedder``, required |
| Used to embed the ``tokens`` ``TextField`` we get as input to the model. |
| encoder : ``Seq2SeqEncoder`` |
| The encoder (with its own internal stacking) that we will use in between embedding tokens |
| and predicting output tags. |
| calculate_span_f1 : ``bool``, optional (default=``None``) |
| Calculate span-level F1 metrics during training. If this is ``True``, then |
| ``label_encoding`` is required. If ``None`` and |
| label_encoding is specified, this is set to ``True``. |
| If ``None`` and label_encoding is not specified, it defaults |
| to ``False``. |
| label_encoding : ``str``, optional (default=``None``) |
| Label encoding to use when calculating span f1. |
| Valid options are "BIO", "BIOUL", "IOB1", "BMES". |
| Required if ``calculate_span_f1`` is true. |
| labels_namespace : ``str``, optional (default=``labels``) |
| This is needed to compute the SpanBasedF1Measure metric, if desired. |
| Unless you did something unusual, the default value should be what you want. |
| verbose_metrics : ``bool``, optional (default = False) |
| If true, metrics will be returned per label class in addition |
| to the overall statistics. |
| initializer : ``InitializerApplicator``, optional (default=``InitializerApplicator()``) |
| Used to initialize the model parameters. |
| regularizer : ``RegularizerApplicator``, optional (default=``None``) |
| If provided, will be used to calculate the regularization penalty during training. |
| """ |
|
|
| def __init__(self, vocab: Vocabulary, |
| text_field_embedder: TextFieldEmbedder, |
| predictor_dropout=0.0, |
| labels_namespace: str = "labels", |
| detect_namespace: str = "d_tags", |
| verbose_metrics: bool = False, |
| label_smoothing: float = 0.0, |
| confidence: float = 0.0, |
| del_confidence: float = 0.0, |
| initializer: InitializerApplicator = InitializerApplicator(), |
| regularizer: Optional[RegularizerApplicator] = None) -> None: |
| super(Seq2Labels, self).__init__(vocab, regularizer) |
|
|
| self.label_namespaces = [labels_namespace, |
| detect_namespace] |
| self.text_field_embedder = text_field_embedder |
| self.num_labels_classes = self.vocab.get_vocab_size(labels_namespace) |
| self.num_detect_classes = self.vocab.get_vocab_size(detect_namespace) |
| self.label_smoothing = label_smoothing |
| self.confidence = confidence |
| self.del_conf = del_confidence |
| self.incorr_index = self.vocab.get_token_index("INCORRECT", |
| namespace=detect_namespace) |
|
|
| self._verbose_metrics = verbose_metrics |
| self.predictor_dropout = TimeDistributed(torch.nn.Dropout(predictor_dropout)) |
|
|
| self.tag_labels_projection_layer = TimeDistributed( |
| Linear(text_field_embedder._token_embedders['bert'].get_output_dim(), self.num_labels_classes)) |
|
|
| self.tag_detect_projection_layer = TimeDistributed( |
| Linear(text_field_embedder._token_embedders['bert'].get_output_dim(), self.num_detect_classes)) |
|
|
| self.metrics = {"accuracy": CategoricalAccuracy()} |
|
|
| initializer(self) |
|
|
| @overrides |
| def forward(self, |
| tokens: Dict[str, torch.LongTensor], |
| labels: torch.LongTensor = None, |
| d_tags: torch.LongTensor = None, |
| metadata: List[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]: |
| |
| """ |
| Parameters |
| ---------- |
| tokens : Dict[str, torch.LongTensor], required |
| The output of ``TextField.as_array()``, which should typically be passed directly to a |
| ``TextFieldEmbedder``. This output is a dictionary mapping keys to ``TokenIndexer`` |
| tensors. At its most basic, using a ``SingleIdTokenIndexer`` this is: ``{"tokens": |
| Tensor(batch_size, num_tokens)}``. This dictionary will have the same keys as were used |
| for the ``TokenIndexers`` when you created the ``TextField`` representing your |
| sequence. The dictionary is designed to be passed directly to a ``TextFieldEmbedder``, |
| which knows how to combine different word representations into a single vector per |
| token in your input. |
| labels : torch.LongTensor, optional (default = None) |
| A torch tensor representing the sequence of integer gold class labels of shape |
| ``(batch_size, num_tokens)``. |
| d_tags : torch.LongTensor, optional (default = None) |
| A torch tensor representing the sequence of integer gold class labels of shape |
| ``(batch_size, num_tokens)``. |
| metadata : ``List[Dict[str, Any]]``, optional, (default = None) |
| metadata containing the original words in the sentence to be tagged under a 'words' key. |
| |
| Returns |
| ------- |
| An output dictionary consisting of: |
| logits : torch.FloatTensor |
| A tensor of shape ``(batch_size, num_tokens, tag_vocab_size)`` representing |
| unnormalised log probabilities of the tag classes. |
| class_probabilities : torch.FloatTensor |
| A tensor of shape ``(batch_size, num_tokens, tag_vocab_size)`` representing |
| a distribution of the tag classes per word. |
| loss : torch.FloatTensor, optional |
| A scalar loss to be optimised. |
| |
| """ |
| encoded_text = self.text_field_embedder(tokens) |
| batch_size, sequence_length, _ = encoded_text.size() |
| mask = get_text_field_mask(tokens) |
| logits_labels = self.tag_labels_projection_layer(self.predictor_dropout(encoded_text)) |
| logits_d = self.tag_detect_projection_layer(encoded_text) |
|
|
| class_probabilities_labels = F.softmax(logits_labels, dim=-1).view( |
| [batch_size, sequence_length, self.num_labels_classes]) |
| class_probabilities_d = F.softmax(logits_d, dim=-1).view( |
| [batch_size, sequence_length, self.num_detect_classes]) |
| error_probs = class_probabilities_d[:, :, self.incorr_index] * mask |
| incorr_prob = torch.max(error_probs, dim=-1)[0] |
|
|
| probability_change = [self.confidence, self.del_conf] + [0] * (self.num_labels_classes - 2) |
| class_probabilities_labels += torch.FloatTensor(probability_change).repeat( |
| (batch_size, sequence_length, 1)).to(class_probabilities_labels.device) |
|
|
| output_dict = {"logits_labels": logits_labels, |
| "logits_d_tags": logits_d, |
| "class_probabilities_labels": class_probabilities_labels, |
| "class_probabilities_d_tags": class_probabilities_d, |
| "max_error_probability": incorr_prob} |
| if labels is not None and d_tags is not None: |
| loss_labels = sequence_cross_entropy_with_logits(logits_labels, labels, mask, |
| label_smoothing=self.label_smoothing) |
| loss_d = sequence_cross_entropy_with_logits(logits_d, d_tags, mask) |
| for metric in self.metrics.values(): |
| metric(logits_labels, labels, mask.float()) |
| metric(logits_d, d_tags, mask.float()) |
| output_dict["loss"] = loss_labels + loss_d |
|
|
| if metadata is not None: |
| output_dict["words"] = [x["words"] for x in metadata] |
| return output_dict |
|
|
| @overrides |
| def decode(self, output_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: |
| """ |
| Does a simple position-wise argmax over each token, converts indices to string labels, and |
| adds a ``"tags"`` key to the dictionary with the result. |
| """ |
| for label_namespace in self.label_namespaces: |
| all_predictions = output_dict[f'class_probabilities_{label_namespace}'] |
| all_predictions = all_predictions.cpu().data.numpy() |
| if all_predictions.ndim == 3: |
| predictions_list = [all_predictions[i] for i in range(all_predictions.shape[0])] |
| else: |
| predictions_list = [all_predictions] |
| all_tags = [] |
|
|
| for predictions in predictions_list: |
| argmax_indices = numpy.argmax(predictions, axis=-1) |
| tags = [self.vocab.get_token_from_index(x, namespace=label_namespace) |
| for x in argmax_indices] |
| all_tags.append(tags) |
| output_dict[f'{label_namespace}'] = all_tags |
| return output_dict |
|
|
| @overrides |
| def get_metrics(self, reset: bool = False) -> Dict[str, float]: |
| metrics_to_return = {metric_name: metric.get_metric(reset) for |
| metric_name, metric in self.metrics.items()} |
| return metrics_to_return |
|
|
