def test_widths_are_embedded_correctly(self):
input_dim = 7
max_span_width = 5
span_width_embedding_dim = 3
output_dim = input_dim + span_width_embedding_dim
extractor = SelfAttentiveSpanExtractor(
input_dim=input_dim,
num_width_embeddings=max_span_width,
span_width_embedding_dim=span_width_embedding_dim,
)
assert extractor.get_output_dim() == output_dim
assert extractor.get_input_dim() == input_dim
sequence_tensor = torch.randn([2, max_span_width, input_dim])
indices = torch.LongTensor([[[1, 3], [0, 4], [0, 0]],
[[0, 2], [1, 4],
[2, 2]]]) # smaller span tests masking.
span_representations = extractor(sequence_tensor, indices)
assert list(span_representations.size()) == [2, 3, output_dim]
width_embeddings = extractor._span_width_embedding.weight.data.numpy()
widths_minus_one = indices[..., 1] - indices[..., 0]
for element in range(indices.size(0)):
for span in range(indices.size(1)):
width = widths_minus_one[element, span].item()
width_embedding = span_representations[element, span,
input_dim:]
numpy.testing.assert_array_almost_equal(
width_embedding.data.numpy(), width_embeddings[width])
def test_attention_is_normalised_correctly(self):
input_dim = 7
sequence_tensor = torch.randn([2, 5, input_dim])
extractor = SelfAttentiveSpanExtractor(input_dim=input_dim)
assert extractor.get_output_dim() == input_dim
assert extractor.get_input_dim() == input_dim
# In order to test the attention, we'll make the weight which computes the logits
# zero, so the attention distribution is uniform over the sentence. This lets
# us check that the computed spans are just the averages of their representations.
extractor._global_attention._module.weight.data.fill_(0.0)
extractor._global_attention._module.bias.data.fill_(0.0)
indices = torch.LongTensor(
[[[1, 3], [2, 4]], [[0, 2], [3, 4]]]
) # smaller span tests masking.
span_representations = extractor(sequence_tensor, indices)
assert list(span_representations.size()) == [2, 2, input_dim]
# First element in the batch.
batch_element = 0
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 1:4, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span.
mean_embeddings = sequence_tensor[batch_element, 2:5, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), mean_embeddings.data.numpy())
# Now the second element in the batch.
batch_element = 1
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 0:3, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span.
mean_embeddings = sequence_tensor[batch_element, 3:5, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), mean_embeddings.data.numpy())
# Now test the case in which we have some masked spans in our indices.
indices_mask = torch.BoolTensor([[True, True], [True, False]])
span_representations = extractor(sequence_tensor, indices, span_indices_mask=indices_mask)
# First element in the batch.
batch_element = 0
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 1:4, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span.
mean_embeddings = sequence_tensor[batch_element, 2:5, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), mean_embeddings.data.numpy())
# Now the second element in the batch.
batch_element = 1
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 0:3, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span was masked, so should be completely zero.
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), numpy.zeros([input_dim]))
def test_attention_is_normalised_correctly(self):
input_dim = 7
sequence_tensor = torch.randn([2, 5, input_dim])
extractor = SelfAttentiveSpanExtractor(input_dim=input_dim)
assert extractor.get_output_dim() == input_dim
assert extractor.get_input_dim() == input_dim
# In order to test the attention, we'll make the weight which computes the logits
# zero, so the attention distribution is uniform over the sentence. This lets
# us check that the computed spans are just the averages of their representations.
extractor._global_attention._module.weight.data.fill_(0.0)
extractor._global_attention._module.bias.data.fill_(0.0)
indices = torch.LongTensor([[[1, 3],
[2, 4]],
[[0, 2],
[3, 4]]]) # smaller span tests masking.
span_representations = extractor(sequence_tensor, indices)
assert list(span_representations.size()) == [2, 2, input_dim]
# First element in the batch.
batch_element = 0
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 1:4, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span.
mean_embeddings = sequence_tensor[batch_element, 2:5, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), mean_embeddings.data.numpy())
# Now the second element in the batch.
batch_element = 1
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 0:3, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span.
mean_embeddings = sequence_tensor[batch_element, 3:5, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), mean_embeddings.data.numpy())
# Now test the case in which we have some masked spans in our indices.
indices_mask = torch.LongTensor([[1, 1], [1, 0]])
span_representations = extractor(sequence_tensor, indices, span_indices_mask=indices_mask)
# First element in the batch.
batch_element = 0
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 1:4, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span.
mean_embeddings = sequence_tensor[batch_element, 2:5, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), mean_embeddings.data.numpy())
# Now the second element in the batch.
batch_element = 1
spans = span_representations[batch_element]
# First span.
mean_embeddings = sequence_tensor[batch_element, 0:3, :].mean(0)
numpy.testing.assert_array_almost_equal(spans[0].data.numpy(), mean_embeddings.data.numpy())
# Second span was masked, so should be completely zero.
numpy.testing.assert_array_almost_equal(spans[1].data.numpy(), numpy.zeros([input_dim]))
class DyGIE(Model):
"""
TODO(dwadden) document me.
Parameters
----------
vocab : ``Vocabulary``
text_field_embedder : ``TextFieldEmbedder``
Used to embed the ``text`` ``TextField`` we get as input to the model.
context_layer : ``Seq2SeqEncoder``
This layer incorporates contextual information for each word in the document.
feature_size: ``int``
The embedding size for all the embedded features, such as distances or span widths.
submodule_params: ``TODO(dwadden)``
A nested dictionary specifying parameters to be passed on to initialize submodules.
max_span_width: ``int``
The maximum width of candidate spans.
max_trigger_span_width: ``int``
The maximum width of candidate trigger spans.
target_task: ``str``:
The task used to make early stopping decisions.
initializer : ``InitializerApplicator``, optional (default=``InitializerApplicator()``)
Used to initialize the model parameters.
module_initializer : ``InitializerApplicator``, optional (default=``InitializerApplicator()``)
Used to initialize the individual modules.
regularizer : ``RegularizerApplicator``, optional (default=``None``)
If provided, will be used to calculate the regularization penalty during training.
display_metrics: ``List[str]``. A list of the metrics that should be printed out during model
training.
"""
def __init__(
self,
vocab: Vocabulary,
embedder: TextFieldEmbedder,
context_layer: Seq2SeqEncoder,
modules, # TODO(dwadden) Add type.
feature_size: int,
max_span_width: int,
max_trigger_span_width: int,
target_task: str,
feedforward_params: Dict[str, Union[int, float]],
loss_weights: Dict[str, float],
lexical_dropout: float = 0.2,
use_attentive_span_extractor: bool = False,
initializer: InitializerApplicator = InitializerApplicator(),
module_initializer: InitializerApplicator = InitializerApplicator(
),
regularizer: Optional[RegularizerApplicator] = None,
display_metrics: List[str] = None) -> None:
super(DyGIE, self).__init__(vocab, regularizer)
####################
# Create span extractor.
self._endpoint_span_extractor = EndpointSpanExtractor(
context_layer.get_output_dim(),
combination="x,y",
num_width_embeddings=max_span_width,
span_width_embedding_dim=feature_size,
bucket_widths=False)
self._endpoint_trigger_span_extractor = EndpointSpanExtractor(
context_layer.get_output_dim(),
combination="x,y",
num_width_embeddings=max_trigger_span_width,
span_width_embedding_dim=feature_size,
bucket_widths=False)
####################
if lexical_dropout > 0:
self._lexical_dropout = torch.nn.Dropout(p=lexical_dropout)
else:
self._lexical_dropout = lambda x: x
if use_attentive_span_extractor:
self._attentive_span_extractor = SelfAttentiveSpanExtractor(
input_dim=context_layer.get_output_dim())
else:
self._attentive_span_extractor = None
# Set parameters.
self._embedder = embedder
self._context_layer = context_layer
self._loss_weights = loss_weights
self._max_span_width = max_span_width
self._max_trigger_span_width = max_trigger_span_width
self._display_metrics = self._get_display_metrics(target_task)
trigger_emb_dim = self._endpoint_trigger_span_extractor.get_output_dim(
)
span_emb_dim = self._endpoint_span_extractor.get_output_dim()
if self._attentive_span_extractor is not None:
span_emb_dim += self._attentive_span_extractor.get_output_dim()
trigger_emb_dim += self._attentive_span_extractor.get_output_dim()
####################
# Create submodules.
modules = Params(modules)
# Helper function to create feedforward networks.
def make_feedforward(input_dim):
return FeedForward(input_dim=input_dim,
num_layers=feedforward_params["num_layers"],
hidden_dims=feedforward_params["hidden_dims"],
activations=torch.nn.ReLU(),
dropout=feedforward_params["dropout"])
# Submodules
self._ner = NERTagger.from_params(vocab=vocab,
make_feedforward=make_feedforward,
span_emb_dim=span_emb_dim,
feature_size=feature_size,
params=modules.pop("ner"))
self._coref = CorefResolver.from_params(
vocab=vocab,
make_feedforward=make_feedforward,
span_emb_dim=span_emb_dim,
feature_size=feature_size,
params=modules.pop("coref"))
self._relation = RelationExtractor.from_params(
vocab=vocab,
make_feedforward=make_feedforward,
span_emb_dim=span_emb_dim,
feature_size=feature_size,
params=modules.pop("relation"))
self._events = EventExtractor.from_params(
vocab=vocab,
make_feedforward=make_feedforward,
text_emb_dim=self._embedder.get_output_dim(),
trigger_emb_dim=trigger_emb_dim,
span_emb_dim=span_emb_dim,
feature_size=feature_size,
params=modules.pop("events"))
####################
# Initialize text embedder and all submodules
for module in [self._ner, self._coref, self._relation, self._events]:
module_initializer(module)
initializer(self)
@staticmethod
def _get_display_metrics(target_task):
"""
The `target` is the name of the task used to make early stopping decisions. Show metrics
related to this task.
"""
lookup = {
"ner": [
f"MEAN__{name}"
for name in ["ner_precision", "ner_recall", "ner_f1"]
],
"relation": [
f"MEAN__{name}" for name in
["relation_precision", "relation_recall", "relation_f1"]
],
"coref": [
"coref_precision", "coref_recall", "coref_f1",
"coref_mention_recall"
],
"events":
[f"MEAN__{name}" for name in ["trig_class_f1", "arg_class_f1"]]
}
if target_task not in lookup:
raise ValueError(
f"Invalied value {target_task} has been given as the target task."
)
return lookup[target_task]
@staticmethod
def _debatch(x):
# TODO(dwadden) Get rid of this when I find a better way to do it.
return x if x is None else x.squeeze(0)
@overrides
def forward(self,
text,
trigger_spans,
spans,
metadata,
ner_labels=None,
coref_labels=None,
relation_labels=None,
trigger_labels=None,
argument_labels=None):
"""
TODO(dwadden) change this.
"""
# In AllenNLP, AdjacencyFields are passed in as floats. This fixes it.
if relation_labels is not None:
relation_labels = relation_labels.long()
if argument_labels is not None:
argument_labels = argument_labels.long()
# TODO(dwadden) Multi-document minibatching isn't supported yet. For now, get rid of the
# extra dimension in the input tensors. Will return to this once the model runs.
if len(metadata) > 1:
raise NotImplementedError(
"Multi-document minibatching not yet supported.")
metadata = metadata[0]
spans = self._debatch(spans) # (n_sents, max_n_spans, 2)
trigger_spans = self._debatch(
trigger_spans) # (n_sents, max_n_spans, 2)
ner_labels = self._debatch(ner_labels) # (n_sents, max_n_spans)
coref_labels = self._debatch(coref_labels) # (n_sents, max_n_spans)
relation_labels = self._debatch(
relation_labels) # (n_sents, max_n_spans, max_n_spans)
trigger_labels = self._debatch(trigger_labels) # TODO(dwadden)
argument_labels = self._debatch(argument_labels) # TODO(dwadden)
# Encode using BERT, then debatch.
# Since the data are batched, we use `num_wrapping_dims=1` to unwrap the document dimension.
# (1, n_sents, max_sententence_length, embedding_dim)
# TODO(dwadden) Deal with the case where the input is longer than 512.
text_embeddings = self._embedder(text, num_wrapping_dims=1)
# (n_sents, max_n_wordpieces, embedding_dim)
text_embeddings = self._debatch(text_embeddings)
# apply lexical dropout
text_embeddings = self._lexical_dropout(text_embeddings)
# (n_sents, max_sentence_length)
text_mask = self._debatch(
util.get_text_field_mask(text, num_wrapping_dims=1).float())
sentence_lengths = text_mask.sum(dim=1).long() # (n_sents)
# contextualize text embeddings
text_embeddings = self._context_layer(text_embeddings, text_mask)
# Create spans, i.e. span_embeddings, masks and span_indices
span_mask = (spans[:, :, 0] >= 0).float() # (n_sents, max_n_spans)
# SpanFields return -1 when they are used as padding. As we do some comparisons based on
# span widths when we attend over the span representations that we generate from these
# indices, we need them to be <= 0. This is only relevant in edge cases where the number of
# spans we consider after the pruning stage is >= the total number of spans, because in this
# case, it is possible we might consider a masked span.
spans = F.relu(spans.float()).long() # (n_sents, max_n_spans, 2)
# Shape: (batch_size, num_spans, 2 * encoding_dim + feature_size)
span_embeddings = self._endpoint_span_extractor(text_embeddings, spans)
trigger_mask = (trigger_spans[:, :, 0] >= 0).float()
trigger_spans = F.relu(trigger_spans.float()).long()
trigger_embeddings = self._endpoint_trigger_span_extractor(
text_embeddings, trigger_spans)
# Make attented spans embeddings
if self._attentive_span_extractor is not None:
# Shape: (batch_size, num_spans, embedding_size)
attended_span_embeddings = self._attentive_span_extractor(
text_embeddings, spans)
attended_trigger_span_embeddings = self._attentive_span_extractor(
text_embeddings, trigger_spans)
# Shape: (batch_size, num_spans, embedding_size + 2 * encoding_dim + feature_size)
span_embeddings = torch.cat(
[span_embeddings, attended_span_embeddings], -1)
trigger_embeddings = torch.cat(
[trigger_embeddings, attended_trigger_span_embeddings], -1)
# Make calls out to the modules to get results.
output_coref = {'loss': 0}
output_ner = {'loss': 0}
output_relation = {'loss': 0}
output_events = {'loss': 0}
# Prune and compute span representations for coreference module
if self._loss_weights["coref"] > 0 or self._coref.coref_prop > 0:
output_coref, coref_indices = self._coref.compute_representations(
spans, span_mask, span_embeddings, sentence_lengths,
coref_labels, metadata)
# Propagation of global information to enhance the span embeddings
if self._coref.coref_prop > 0:
output_coref = self._coref.coref_propagation(output_coref)
span_embeddings = self._coref.update_spans(output_coref,
span_embeddings,
coref_indices)
# Make predictions and compute losses for each module
if self._loss_weights['ner'] > 0:
output_ner = self._ner(spans, span_mask, span_embeddings,
sentence_lengths, ner_labels, metadata)
if self._loss_weights['coref'] > 0:
output_coref = self._coref.predict_labels(output_coref, metadata)
if self._loss_weights['relation'] > 0:
output_relation = self._relation(spans, span_mask, span_embeddings,
sentence_lengths, relation_labels,
metadata)
if self._loss_weights['events'] > 0:
output_events = self._events(trigger_spans, trigger_mask,
trigger_embeddings, spans, span_mask,
span_embeddings, text_mask,
text_embeddings, sentence_lengths,
trigger_labels, argument_labels,
ner_labels, metadata)
# Use `get` since there are some cases where the output dict won't have a loss - for
# instance, when doing prediction.
loss = (
self._loss_weights['coref'] * output_coref.get("loss", 0) +
self._loss_weights['ner'] * output_ner.get("loss", 0) +
self._loss_weights['relation'] * output_relation.get("loss", 0) +
self._loss_weights['events'] * output_events.get("loss", 0))
# Multiply the loss by the weight multiplier for this document.
weight = metadata.weight if metadata.weight is not None else 1.0
loss *= torch.tensor(weight)
output_dict = dict(coref=output_coref,
relation=output_relation,
ner=output_ner,
events=output_events)
output_dict['loss'] = loss
output_dict["metadata"] = metadata
return output_dict
def update_span_embeddings(self, span_embeddings, span_mask,
top_span_embeddings, top_span_mask,
top_span_indices):
# TODO(Ulme) Speed this up by tensorizing
new_span_embeddings = span_embeddings.clone()
for sample_nr in range(len(top_span_mask)):
for top_span_nr, span_nr in enumerate(top_span_indices[sample_nr]):
if top_span_mask[sample_nr,
top_span_nr] == 0 or span_mask[sample_nr,
span_nr] == 0:
break
new_span_embeddings[sample_nr,
span_nr] = top_span_embeddings[sample_nr,
top_span_nr]
return new_span_embeddings
@overrides
def make_output_human_readable(self, output_dict: Dict[str, torch.Tensor]):
"""
Converts the list of spans and predicted antecedent indices into clusters
of spans for each element in the batch.
Parameters
----------
output_dict : ``Dict[str, torch.Tensor]``, required.
The result of calling :func:`forward` on an instance or batch of instances.
Returns
-------
The same output dictionary, but with an additional ``clusters`` key:
clusters : ``List[List[List[Tuple[int, int]]]]``
A nested list, representing, for each instance in the batch, the list of clusters,
which are in turn comprised of a list of (start, end) inclusive spans into the
original document.
"""
doc = copy.deepcopy(output_dict["metadata"])
if self._loss_weights["coref"] > 0:
# TODO(dwadden) Will need to get rid of the [0] when batch training is enabled.
decoded_coref = self._coref.make_output_human_readable(
output_dict["coref"])["predicted_clusters"][0]
sentences = doc.sentences
sentence_starts = [sent.sentence_start for sent in sentences]
predicted_clusters = [
document.Cluster(entry, i, sentences, sentence_starts)
for i, entry in enumerate(decoded_coref)
]
doc.predicted_clusters = predicted_clusters
# TODO(dwadden) update the sentences with cluster information.
if self._loss_weights["ner"] > 0:
for predictions, sentence in zip(output_dict["ner"]["predictions"],
doc):
sentence.predicted_ner = predictions
if self._loss_weights["relation"] > 0:
for predictions, sentence in zip(
output_dict["relation"]["predictions"], doc):
sentence.predicted_relations = predictions
if self._loss_weights["events"] > 0:
for predictions, sentence in zip(
output_dict["events"]["predictions"], doc):
sentence.predicted_events = predictions
return doc
def get_metrics(self, reset: bool = False) -> Dict[str, float]:
"""
Get all metrics from all modules. For the ones that shouldn't be displayed, prefix their
keys with an underscore.
"""
metrics_coref = self._coref.get_metrics(reset=reset)
metrics_ner = self._ner.get_metrics(reset=reset)
metrics_relation = self._relation.get_metrics(reset=reset)
metrics_events = self._events.get_metrics(reset=reset)
# Make sure that there aren't any conflicting names.
metric_names = (list(metrics_coref.keys()) + list(metrics_ner.keys()) +
list(metrics_relation.keys()) +
list(metrics_events.keys()))
assert len(set(metric_names)) == len(metric_names)
all_metrics = dict(
list(metrics_coref.items()) + list(metrics_ner.items()) +
list(metrics_relation.items()) + list(metrics_events.items()))
# If no list of desired metrics given, display them all.
if self._display_metrics is None:
return all_metrics
# Otherwise only display the selected ones.
res = {}
for k, v in all_metrics.items():
if k in self._display_metrics:
res[k] = v
else:
new_k = "_" + k
res[new_k] = v
return res
class SCIIE(Model):
"""
Parameters
----------
vocab : ``Vocabulary``
text_field_embedder : ``TextFieldEmbedder``
Used to embed the ``text`` ``TextField`` we get as input to the model.
context_layer : ``Seq2SeqEncoder``
This layer incorporates contextual information for each word in the document.
mention_feedforward : ``FeedForward``
This feedforward network is applied to the span representations which is then scored
by a linear layer.
antecedent_feedforward: ``FeedForward``
This feedforward network is applied to pairs of span representation, along with any
pairwise features, 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.
max_span_width: ``int``
The maximum width of candidate spans.
spans_per_word: float, required.
A multiplier between zero and one which controls what percentage of candidate mention
spans we retain with respect to the number of words in the document.
max_antecedents: int, required.
For each mention which survives the pruning stage, we consider this many antecedents.
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,
embedding_dim: int,
feature_size: int,
max_span_width: int,
spans_per_word: float,
lexical_dropout: float = 0.2,
mlp_dropout: float = 0.4,
embedder_type=None,
regularizer: Optional[RegularizerApplicator] = None) -> None:
super(SCIIE, self).__init__(vocab, regularizer)
self.class_num = self.vocab.get_vocab_size('labels')
word_embeddings = get_embeddings(embedder_type, self.vocab,
embedding_dim, True)
embedding_dim = word_embeddings.get_output_dim()
self._text_field_embedder = word_embeddings
context_layer = PytorchSeq2SeqWrapper(
torch.nn.LSTM(embedding_dim,
feature_size,
batch_first=True,
bidirectional=True))
self._context_layer = context_layer
endpoint_span_extractor_input_dim = context_layer.get_output_dim()
attentive_span_extractor_input_dim = word_embeddings.get_output_dim()
self._endpoint_span_extractor = EndpointSpanExtractor(
endpoint_span_extractor_input_dim,
combination="x,y",
num_width_embeddings=max_span_width,
span_width_embedding_dim=feature_size,
bucket_widths=False)
self._attentive_span_extractor = SelfAttentiveSpanExtractor(
input_dim=attentive_span_extractor_input_dim)
# self._span_extractor = PoolingSpanExtractor(embedding_dim,
# num_width_embeddings=max_span_width,
# span_width_embedding_dim=feature_size,
# bucket_widths=False)
entity_feedforward = FeedForward(
self._endpoint_span_extractor.get_output_dim() +
self._attentive_span_extractor.get_output_dim(), 2, 150, F.relu,
mlp_dropout)
# entity_feedforward = FeedForward(self._span_extractor.get_output_dim(), 2, 150,
# F.relu, mlp_dropout)
feedforward_scorer = torch.nn.Sequential(
TimeDistributed(entity_feedforward),
TimeDistributed(
torch.nn.Linear(entity_feedforward.get_output_dim(), 1)))
self._mention_pruner = Pruner(feedforward_scorer)
self._entity_scorer = torch.nn.Sequential(
TimeDistributed(entity_feedforward),
TimeDistributed(
torch.nn.Linear(entity_feedforward.get_output_dim(),
self.class_num - 1)))
self._max_span_width = max_span_width
self._spans_per_word = spans_per_word
if lexical_dropout > 0:
self._lexical_dropout = torch.nn.Dropout(p=lexical_dropout)
else:
self._lexical_dropout = lambda x: x
self._metric_all = FBetaMeasure()
self._metric_avg = NERF1Metric()
@overrides
def forward(
self, # type: ignore
text: Dict[str, torch.LongTensor],
spans: torch.IntTensor,
labels: torch.IntTensor = None,
metadata: List[Dict[str, Any]] = None,
**kwargs) -> Dict[str, torch.Tensor]:
# Shape: (batch_size, document_length, embedding_size)
text_embeddings = self._lexical_dropout(
self._text_field_embedder(text))
document_length = text_embeddings.size(1)
num_spans = spans.size(1)
# Shape: (batch_size, document_length)
text_mask = util.get_text_field_mask(text).float()
# Shape: (batch_size, num_spans)
span_mask = (spans[:, :, 0] >= 0).squeeze(-1).float()
# SpanFields return -1 when they are used as padding. As we do
# some comparisons based on span widths when we attend over the
# span representations that we generate from these indices, we
# need them to be <= 0. This is only relevant in edge cases where
# the number of spans we consider after the pruning stage is >= the
# total number of spans, because in this case, it is possible we might
# consider a masked span.
# Shape: (batch_size, num_spans, 2)
spans = F.relu(spans.float()).long()
# Shape: (batch_size, document_length, encoding_dim)
contextualized_embeddings = self._context_layer(
text_embeddings, text_mask)
# Shape: (batch_size, num_spans, 2 * encoding_dim + feature_size)
endpoint_span_embeddings = self._endpoint_span_extractor(
contextualized_embeddings, spans)
# Shape: (batch_size, num_spans, emebedding_size)
attended_span_embeddings = self._attentive_span_extractor(
text_embeddings, spans)
# Shape: (batch_size, num_spans, emebedding_size + 2 * encoding_dim + feature_size)
span_embeddings = torch.cat(
[endpoint_span_embeddings, attended_span_embeddings], -1)
# span_embeddings = self._span_extractor(text_embeddings, spans, span_indices_mask=span_mask)
# Prune based on mention scores.
num_spans_to_keep = int(
math.floor(self._spans_per_word * document_length))
num_spans_to_keep = min(num_spans_to_keep, span_embeddings.shape[1])
# Shape: (batch_size, num_spans_to_keep, emebedding_size + 2 * encoding_dim + feature_size)
# (batch_size, num_spans_to_keep)
# (batch_size, num_spans_to_keep)
# (batch_size, num_spans_to_keep, 1)
(top_span_embeddings, top_span_mask, top_span_indices,
top_span_mention_scores) = self._mention_pruner(
span_embeddings, span_mask, num_spans_to_keep)
# (batch_size, num_spans_to_keep, 1)
top_span_mask = top_span_mask.unsqueeze(-1)
# Shape: (batch_size * num_spans_to_keep)
# torch.index_select only accepts 1D indices, but here
# we need to select spans for each element in the batch.
# This reformats the indices to take into account their
# index into the batch. We precompute this here to make
# the multiple calls to util.batched_index_select below more efficient.
flat_top_span_indices = util.flatten_and_batch_shift_indices(
top_span_indices, num_spans)
# Compute final predictions for which spans to consider as mentions.
# Shape: (batch_size, num_spans_to_keep, 2)
top_spans = util.batched_index_select(spans, top_span_indices,
flat_top_span_indices)
# Shape: (batch_size, num_spans_to_keep, class_num + 1)
ne_scores = self._compute_named_entity_scores(top_span_embeddings)
# Shape: (batch_size, num_spans_to_keep)
_, predicted_named_entities = ne_scores.max(2)
output_dict = {
"top_spans": top_spans,
"predicted_named_entities": predicted_named_entities
}
if labels is not None:
# Find the gold labels for the spans which we kept.
# Shape: (batch_size, num_spans_to_keep, 1)
pruned_gold_labels = util.batched_index_select(
labels.unsqueeze(-1), top_span_indices,
flat_top_span_indices).squeeze(-1)
negative_log_likelihood = F.cross_entropy(
ne_scores.reshape(-1, self.class_num),
pruned_gold_labels.reshape(-1))
pruner_loss = F.binary_cross_entropy_with_logits(
top_span_mention_scores.reshape(-1),
(pruned_gold_labels.reshape(-1) != 0).float())
loss = negative_log_likelihood + pruner_loss
output_dict["loss"] = loss
output_dict["pruner_loss"] = pruner_loss
batch_size, _ = labels.shape
all_scores = ne_scores.new_zeros(
[batch_size * num_spans, self.class_num])
all_scores[:, 0] = 1
all_scores[flat_top_span_indices] = ne_scores.reshape(
-1, self.class_num)
all_scores = all_scores.reshape(
[batch_size, num_spans, self.class_num])
self._metric_all(all_scores, labels)
self._metric_avg(all_scores, labels)
return output_dict
@overrides
def get_metrics(self, reset: bool = False, prefix=""):
metric = self._metric_all.get_metric(reset)
metric2 = self._metric_avg.get_metric(reset)
metric.update(metric2)
return metric
def _compute_named_entity_scores(
self, span_embeddings: torch.FloatTensor) -> torch.Tensor:
"""
Parameters
----------
span_embeddings: ``torch.FloatTensor``, required.
Embedding representations of spans. Has shape
(batch_size, num_spans_to_keep, encoding_dim)
"""
# Shape: (batch_size, num_spans_to_keep, class_num)
scores = self._entity_scorer(span_embeddings)
# Shape: (batch_size, num_spans_to_keep, 1)
shape = [scores.size(0), scores.size(1), 1]
dummy_scores = scores.new_full(shape, 0)
ne_scores = torch.cat([dummy_scores, scores], -1)
return ne_scores
class SrlE2e(Model):
"""
# Parameters
vocab : `Vocabulary`, required
A Vocabulary, required in order to compute sizes for input/output projections.
model : `Union[str, BertModel]`, required.
A string describing the BERT model to load or an already constructed BertModel.
initializer : `InitializerApplicator`, optional (default=`InitializerApplicator()`)
Used to initialize the model parameters.
label_smoothing : `float`, optional (default = `0.0`)
Whether or not to use label smoothing on the labels when computing cross entropy loss.
ignore_span_metric : `bool`, optional (default = `False`)
Whether to calculate span loss, which is irrelevant when predicting BIO for Open Information Extraction.
srl_eval_path : `str`, optional (default=`DEFAULT_SRL_EVAL_PATH`)
The path to the srl-eval.pl script. By default, will use the srl-eval.pl included with allennlp,
which is located at allennlp/tools/srl-eval.pl . If `None`, srl-eval.pl is not used.
"""
def __init__(
self,
vocab: Vocabulary,
bert_model: Union[str, BertModel],
mention_feedforward: FeedForward,
context_layer: Seq2SeqEncoder = None,
embedding_dropout: float = 0.0,
initializer: InitializerApplicator = InitializerApplicator(),
max_span_width: int = 30,
feature_size: int = 10,
spans_per_word: float = 100,
label_smoothing: float = None,
ignore_span_metric: bool = False,
srl_eval_path: str = DEFAULT_SRL_EVAL_PATH,
**kwargs,
) -> None:
super().__init__(vocab, **kwargs)
if isinstance(bert_model, str):
self.bert_model = BertModel.from_pretrained(bert_model)
else:
self.bert_model = bert_model
self.num_classes = self.vocab.get_vocab_size("span_labels")
if srl_eval_path is not None:
# For the span based evaluation, we don't want to consider labels
# for verb, because the verb index is provided to the model.
self.span_metric = SrlEvalScorer(srl_eval_path,
ignore_classes=["V"])
else:
self.span_metric = None
self.tag_projection_layer = Linear(self.bert_model.config.hidden_size,
self.num_classes)
self.embedding_dropout = Dropout(p=embedding_dropout)
self._label_smoothing = label_smoothing
self.ignore_span_metric = ignore_span_metric
self._mention_feedforward = TimeDistributed(mention_feedforward)
self._mention_scorer = TimeDistributed(
torch.nn.Linear(mention_feedforward.get_output_dim(), 1))
self._attentive_span_extractor = SelfAttentiveSpanExtractor(
input_dim=self.bert_model.config.hidden_size)
self.span_representation_dim = self._attentive_span_extractor.get_output_dim(
)
self._context_layer = context_layer
if context_layer is not None:
self._endpoint_span_extractor = EndpointSpanExtractor(
context_layer.get_output_dim(),
combination="x,y",
num_width_embeddings=max_span_width,
span_width_embedding_dim=feature_size,
bucket_widths=False,
)
self.span_representation_dim = self._endpoint_span_extractor.get_output_dim(
)
self.hidden_layer = torch.nn.Sequential(
torch.nn.Linear(self.span_representation_dim +
self.bert_model.config.hidden_size,
self.span_representation_dim,
bias=False), torch.nn.ReLU())
self.output_layer = torch.nn.Linear(self.span_representation_dim,
self.num_classes - 1,
bias=False)
self._max_span_width = max_span_width
self._spans_per_word = spans_per_word
self._ce_loss = torch.nn.CrossEntropyLoss(reduction='none')
self._bce_loss = torch.nn.BCEWithLogitsLoss(reduction='none')
initializer(self)
def forward( # type: ignore
self,
tokens: TextFieldTensors,
verb_indicator: torch.Tensor,
sentence_end: torch.LongTensor,
spans: torch.LongTensor,
span_labels: torch.LongTensor,
metadata: List[Any],
tags: torch.LongTensor = None,
):
"""
# Parameters
tokens : `TextFieldTensors`, required
The output of `TextField.as_array()`, which should typically be passed directly to a
`TextFieldEmbedder`. For this model, this must be a `SingleIdTokenIndexer` which
indexes wordpieces from the BERT vocabulary.
verb_indicator: `torch.LongTensor`, required.
An integer `SequenceFeatureField` representation of the position of the verb
in the sentence. This should have shape (batch_size, num_tokens) and importantly, can be
all zeros, in the case that the sentence has no verbal predicate.
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 containg the original words in the sentence, the verb to compute the
frame for, and start offsets for converting wordpieces back to a sequence of words,
under 'words', 'verb' and 'offsets' keys, respectively.
# 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.
"""
mask = get_text_field_mask(tokens)
start = time.time()
bert_embeddings, _ = self.bert_model(
input_ids=util.get_token_ids_from_text_field_tensors(tokens),
# token_type_ids=verb_indicator,
attention_mask=mask,
)
# Shape: (batch_size, num_spans)
span_mask = (spans[:, :, 0] >= 0).squeeze(-1)
# SpanFields return -1 when they are used as padding. As we do
# some comparisons based on span widths when we attend over the
# span representations that we generate from these indices, we
# need them to be <= 0. This is only relevant in edge cases where
# the number of spans we consider after the pruning stage is >= the
# total number of spans, because in this case, it is possible we might
# consider a masked span.
# Shape: (batch_size, num_spans, 2)
spans = F.relu(spans.float()).long()
embedded_text_input = self.embedding_dropout(bert_embeddings)
batch_size, sequence_length, _ = embedded_text_input.size()
# Shape: (batch_size, num_spans, emebedding_size)
attended_span_embeddings = self._attentive_span_extractor(
bert_embeddings, spans)
if self._context_layer is not None:
contextualized_embeddings = self._context_layer(
embedded_text_input, mask)
# Shape: (batch_size, num_spans, 2 * encoding_dim + feature_size)
endpoint_span_embeddings = self._endpoint_span_extractor(
contextualized_embeddings, spans)
# Shape: (batch_size, num_spans, emebedding_size + 2 * encoding_dim + feature_size)
# span_embeddings = torch.cat([endpoint_span_embeddings, attended_span_embeddings], -1)
span_embeddings = endpoint_span_embeddings
else:
span_embeddings = attended_span_embeddings
# Prune based on mention scores.
num_spans_to_keep = int(
math.floor(self._spans_per_word * sequence_length))
num_spans = spans.shape[1]
num_spans_to_keep = min(num_spans_to_keep, num_spans)
# Shape: (batch_size, num_spans)
span_mention_scores = self._mention_scorer(
self._mention_feedforward(span_embeddings)).squeeze(-1)
# Shape: (batch_size, num_spans) for all 3 tensors
top_span_mention_scores, top_span_mask, top_span_indices = util.masked_topk(
span_mention_scores, span_mask, num_spans_to_keep)
verb_index = verb_indicator.argmax(1).unsqueeze(1).unsqueeze(2).repeat(
1, 1, embedded_text_input.shape[-1])
verb_embeddings = torch.gather(embedded_text_input, 1, verb_index)
assert len(
verb_embeddings.shape) == 3 and verb_embeddings.shape[1] == 1
verb_embeddings = verb_embeddings.squeeze(1)
# print(verb_indicator.sum(1, keepdim=True) > 0)
verb_embeddings = torch.where(
(verb_indicator.sum(1, keepdim=True) > 0).repeat(
1, verb_embeddings.shape[-1]), verb_embeddings,
torch.zeros_like(verb_embeddings))
# print(verb_embeddings)
flat_top_span_indices = util.flatten_and_batch_shift_indices(
top_span_indices, spans.shape[1])
span_embeddings = util.batched_index_select(span_embeddings,
top_span_indices,
flat_top_span_indices)
top_spans = util.batched_index_select(spans, top_span_indices,
flat_top_span_indices)
top_span_labels = util.batched_index_select(
span_labels.unsqueeze(-1), top_span_indices,
flat_top_span_indices).squeeze(-1)
concatenated_span_embeddings = torch.cat(
(span_embeddings, verb_embeddings.unsqueeze(1).repeat(
1, span_embeddings.shape[1], 1)),
dim=2)
# print(concatenated_span_embeddings[:,:,:])
hidden = self.hidden_layer(concatenated_span_embeddings)
# print(hidden[1,:,:])
# print(top_span_indices)
# print([[span_mention_scores[i,top_span_indices[i,j]].item() for j in range(top_span_indices.shape[1])] for i in range(top_span_labels.shape[0])])
# print(top_span_mention_scores, self.vocab.get_token_index("O", namespace="span_labels"))
predictions = self.output_layer(hidden)
# predictions += top_span_mention_scores.unsqueeze(-1).repeat(1, 1, self.num_classes-1)
predictions = torch.cat(
(torch.zeros_like(predictions[:, :, :1]), predictions), dim=-1)
# print(top_span_mention_scores.unsqueeze(-1).repeat(1, 1, self.num_classes-1))
output_dict = {}
# We need to retain the mask in the output dictionary
# so that we can crop the sequences to remove padding
# when we do viterbi inference in self.make_output_human_readable.
output_dict["mask"] = mask
# We add in the offsets here so we can compute the un-wordpieced tags.
words, verbs, offsets = zip(*[(x["words"], x["verb"], x["offsets"])
for x in metadata])
output_dict["words"] = list(words)
output_dict["verb"] = list(verbs)
output_dict["wordpiece_offsets"] = list(offsets)
if tags is not None:
loss = (self._ce_loss(predictions.view(-1, predictions.shape[-1]),
top_span_labels.view(-1)) *
top_span_mask.float().view(-1)
).sum() / top_span_mask.float().sum()
# print(top_span_labels)
# print(predictions.argmax(-1))
if not self.ignore_span_metric and self.span_metric is not None and not self.training:
batch_verb_indices = [
example_metadata["verb_index"]
for example_metadata in metadata
]
batch_sentences = [
example_metadata["words"] for example_metadata in metadata
]
# Get the BIO tags from make_output_human_readable()
# TODO (nfliu): This is kind of a hack, consider splitting out part
# of make_output_human_readable() to a separate function.
batch_bio_predicted_tags = self.get_tags(
top_spans, predictions, mask.shape[1], top_span_mask,
output_dict)
from allennlp_models.structured_prediction.models.srl import (
convert_bio_tags_to_conll_format, )
batch_conll_predicted_tags = [
convert_bio_tags_to_conll_format(tags)
for tags in batch_bio_predicted_tags
]
batch_bio_gold_tags = [
example_metadata["gold_tags"]
for example_metadata in metadata
]
# print('G', batch_bio_gold_tags)
batch_conll_gold_tags = [
convert_bio_tags_to_conll_format(tags)
for tags in batch_bio_gold_tags
]
self.span_metric(
batch_verb_indices,
batch_sentences,
batch_conll_predicted_tags,
batch_conll_gold_tags,
)
output_dict["loss"] = loss
return output_dict
def get_tags(self, spans, logits, sequence_length, span_mask, output_dict):
predicted_tag_ids = logits.argmax(2)
predicted_tags = []
for i in range(spans.shape[0]):
sequence = ["O" for _ in range(sequence_length)]
for j in range(spans.shape[1]):
if span_mask[i, j].item() == 0:
continue
tag = predicted_tag_ids[i, j].item()
if tag != self.vocab.get_token_index("O",
namespace="span_labels"):
start = spans[i, j, 0].item()
end = spans[i, j, 1].item()
if all([el == "O" for el in sequence[start:end + 1]]):
tag = self.vocab.get_token_from_index(
tag, namespace="span_labels")
sequence[start] = "B-" + tag
for index in range(start + 1, end + 1):
sequence[index] = "I-" + tag
predicted_tags.append(
[sequence[ind] for ind in output_dict["wordpiece_offsets"][i]])
print(predicted_tags)
return predicted_tags
@overrides
def make_output_human_readable(
self, output_dict: Dict[str,
torch.Tensor]) -> Dict[str, torch.Tensor]:
"""
Does constrained viterbi decoding on class probabilities output in :func:`forward`. The
constraint simply specifies that the output tags must be a valid BIO sequence. We add a
`"tags"` key to the dictionary with the result.
NOTE: First, we decode a BIO sequence on top of the wordpieces. This is important; viterbi
decoding produces low quality output if you decode on top of word representations directly,
because the model gets confused by the 'missing' positions (which is sensible as it is trained
to perform tagging on wordpieces, not words).
Secondly, it's important that the indices we use to recover words from the wordpieces are the
start_offsets (i.e offsets which correspond to using the first wordpiece of words which are
tokenized into multiple wordpieces) as otherwise, we might get an ill-formed BIO sequence
when we select out the word tags from the wordpiece tags. This happens in the case that a word
is split into multiple word pieces, and then we take the last tag of the word, which might
correspond to, e.g, I-V, which would not be allowed as it is not preceeded by a B tag.
"""
all_predictions = output_dict["class_probabilities"]
sequence_lengths = get_lengths_from_binary_sequence_mask(
output_dict["mask"]).data.tolist()
if all_predictions.dim() == 3:
predictions_list = [
all_predictions[i].detach().cpu()
for i in range(all_predictions.size(0))
]
else:
predictions_list = [all_predictions]
wordpiece_tags = []
word_tags = []
transition_matrix = self.get_viterbi_pairwise_potentials()
start_transitions = self.get_start_transitions()
# **************** Different ********************
# We add in the offsets here so we can compute the un-wordpieced tags.
for predictions, length, offsets in zip(
predictions_list, sequence_lengths,
output_dict["wordpiece_offsets"]):
max_likelihood_sequence, _ = viterbi_decode(
predictions[:length],
transition_matrix,
allowed_start_transitions=start_transitions)
tags = [
self.vocab.get_token_from_index(x, namespace="labels")
for x in max_likelihood_sequence
]
wordpiece_tags.append(tags)
word_tags.append([tags[i] for i in offsets])
output_dict["wordpiece_tags"] = wordpiece_tags
output_dict["tags"] = word_tags
return output_dict
def get_metrics(self, reset: bool = False):
if self.ignore_span_metric:
# Return an empty dictionary if ignoring the
# span metric
return {}
else:
metric_dict = self.span_metric.get_metric(reset=reset)
# This can be a lot of metrics, as there are 3 per class.
# we only really care about the overall metrics, so we filter for them here.
return {x: y for x, y in metric_dict.items() if "overall" in x}
def get_viterbi_pairwise_potentials(self):
"""
Generate a matrix of pairwise transition potentials for the BIO labels.
The only constraint implemented here is that I-XXX labels must be preceded
by either an identical I-XXX tag or a B-XXX tag. In order to achieve this
constraint, pairs of labels which do not satisfy this constraint have a
pairwise potential of -inf.
# Returns
transition_matrix : `torch.Tensor`
A `(num_labels, num_labels)` matrix of pairwise potentials.
"""
all_labels = self.vocab.get_index_to_token_vocabulary("labels")
num_labels = len(all_labels)
transition_matrix = torch.zeros([num_labels, num_labels])
for i, previous_label in all_labels.items():
for j, label in all_labels.items():
# I labels can only be preceded by themselves or
# their corresponding B tag.
if i != j and label[
0] == "I" and not previous_label == "B" + label[1:]:
transition_matrix[i, j] = float("-inf")
return transition_matrix
def get_start_transitions(self):
"""
In the BIO sequence, we cannot start the sequence with an I-XXX tag.
This transition sequence is passed to viterbi_decode to specify this constraint.
# Returns
start_transitions : `torch.Tensor`
The pairwise potentials between a START token and
the first token of the sequence.
"""
all_labels = self.vocab.get_index_to_token_vocabulary("labels")
num_labels = len(all_labels)
start_transitions = torch.zeros(num_labels)
for i, label in all_labels.items():
if label[0] == "I":
start_transitions[i] = float("-inf")
return start_transitions
default_predictor = "semantic_role_labeling"