def __init__(self,
vocab: Vocabulary,
mention_feedforward: FeedForward,
feature_size: int,
initializer: InitializerApplicator = InitializerApplicator(),
regularizer: Optional[RegularizerApplicator] = None) -> None:
super(NERTagger_Has_None, self).__init__(vocab, regularizer)
# Number of classes determine the output dimension of the final layer
self._n_labels = vocab.get_vocab_size('ner_labels')
# TODO(dwadden) think of a better way to enforce this.
# Null label is needed to keep track of when calculating the metrics
null_label = vocab.get_token_index("", "ner_labels")
assert null_label == 0 # If not, the dummy class won't correspond to the null label.
self._ner_scorer = torch.nn.Sequential(
TimeDistributed(mention_feedforward),
TimeDistributed(
torch.nn.Linear(mention_feedforward.get_output_dim(),
self._n_labels)))
self._ner_metrics = NERMetrics(self._n_labels, null_label)
self._loss = torch.nn.CrossEntropyLoss(reduction="sum")
initializer(self)
def __init__(self,
vocab: Vocabulary,
make_feedforward: Callable,
span_emb_dim: int,
regularizer: Optional[RegularizerApplicator] = None) -> None:
super(NERTagger, self).__init__(vocab, regularizer)
self._namespaces = [
entry for entry in vocab.get_namespaces() if "ner_labels" in entry
]
# Number of classes determine the output dimension of the final layer
self._n_labels = {
name: vocab.get_vocab_size(name)
for name in self._namespaces
}
# Null label is needed to keep track of when calculating the metrics
for namespace in self._namespaces:
null_label = vocab.get_token_index("", namespace)
assert null_label == 0 # If not, the dummy class won't correspond to the null label.
# The output dim is 1 less than the number of labels because we don't score the null label;
# we just give it a score of 0 by default.
# Create a separate scorer and metric for each dataset we're dealing with.
self._ner_scorers = torch.nn.ModuleDict()
self._ner_metrics = {}
for namespace in self._namespaces:
mention_feedforward = make_feedforward(input_dim=span_emb_dim)
self._ner_scorers[namespace] = torch.nn.Sequential(
TimeDistributed(mention_feedforward),
TimeDistributed(
torch.nn.Linear(mention_feedforward.get_output_dim(),
self._n_labels[namespace] - 1)))
self._ner_metrics[namespace] = NERMetrics(
self._n_labels[namespace], null_label)
self._active_namespace = None
self._loss = torch.nn.CrossEntropyLoss(reduction="sum")
class NERTagger_Has_None(Model):
"""
Named entity recognition module of DyGIE model.
Parameters
----------
mention_feedforward : ``FeedForward``
This feedforward network is applied to the span representations which is then scored
by a linear layer.
feature_size: ``int``
The embedding size for all the embedded features, such as distances or span widths.
lexical_dropout: ``int``
The probability of dropping out dimensions of the embedded text.
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,
mention_feedforward: FeedForward,
feature_size: int,
initializer: InitializerApplicator = InitializerApplicator(),
regularizer: Optional[RegularizerApplicator] = None) -> None:
super(NERTagger_Has_None, self).__init__(vocab, regularizer)
# Number of classes determine the output dimension of the final layer
self._n_labels = vocab.get_vocab_size('ner_labels')
# TODO(dwadden) think of a better way to enforce this.
# Null label is needed to keep track of when calculating the metrics
null_label = vocab.get_token_index("", "ner_labels")
assert null_label == 0 # If not, the dummy class won't correspond to the null label.
self._ner_scorer = torch.nn.Sequential(
TimeDistributed(mention_feedforward),
TimeDistributed(
torch.nn.Linear(mention_feedforward.get_output_dim(),
self._n_labels)))
self._ner_metrics = NERMetrics(self._n_labels, null_label)
self._loss = torch.nn.CrossEntropyLoss(reduction="sum")
initializer(self)
@overrides
def forward(
self, # type: ignore
spans: torch.IntTensor,
span_mask: torch.IntTensor,
span_embeddings: torch.IntTensor,
sentence_lengths: torch.Tensor,
ner_labels: torch.IntTensor = None,
metadata: List[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
"""
TODO(dwadden) Write documentation.
"""
# Shape: (Batch size, Number of Spans, Span Embedding Size)
# span_embeddings
ner_scores = self._ner_scorer(span_embeddings)
# Give large negative scores to masked-out elements.
mask = span_mask.unsqueeze(-1)
ner_scores = util.replace_masked_values(ner_scores, mask, -1e20)
ner_scores[:, :, 0] *= span_mask
_, predicted_ner = ner_scores.max(2)
output_dict = {
"spans": spans,
"span_mask": span_mask,
"ner_scores": ner_scores,
"predicted_ner": predicted_ner
}
if ner_labels is not None:
self._ner_metrics(predicted_ner, ner_labels, span_mask)
ner_scores_flat = ner_scores.view(-1, self._n_labels)
ner_labels_flat = ner_labels.view(-1)
mask_flat = span_mask.view(-1).bool()
loss = self._loss(ner_scores_flat[mask_flat],
ner_labels_flat[mask_flat])
output_dict["loss"] = loss
if metadata is not None:
output_dict["document"] = [x["sentence"] for x in metadata]
return output_dict
@overrides
def decode(self, output_dict: Dict[str, torch.Tensor]):
predicted_ner_batch = output_dict["predicted_ner"].detach().cpu()
spans_batch = output_dict["spans"].detach().cpu()
span_mask_batch = output_dict["span_mask"].detach().cpu().bool()
res_list = []
res_dict = []
for spans, span_mask, predicted_NERs in zip(spans_batch,
span_mask_batch,
predicted_ner_batch):
entry_list = []
entry_dict = {}
for span, ner in zip(spans[span_mask], predicted_NERs[span_mask]):
ner = ner.item()
if ner > 0:
the_span = (span[0].item(), span[1].item())
the_label = self.vocab.get_token_from_index(
ner, "ner_labels")
entry_list.append((the_span[0], the_span[1], the_label))
entry_dict[the_span] = the_label
res_list.append(entry_list)
res_dict.append(entry_dict)
output_dict["decoded_ner"] = res_list
output_dict["decoded_ner_dict"] = res_dict
return output_dict
@overrides
def get_metrics(self, reset: bool = False) -> Dict[str, float]:
ner_precision, ner_recall, ner_f1 = self._ner_metrics.get_metric(reset)
return {
"ner_precision": ner_precision,
"ner_recall": ner_recall,
"ner_f1": ner_f1
}