Skip to content

Commit 27dbbb9

Browse files
authored
Bugfix/nel crossing sentence (#7630)
* ensure each entity gets a KB ID, even when it's not within a sentence * cleanup
1 parent 673e2bc commit 27dbbb9

2 files changed

Lines changed: 127 additions & 70 deletions

File tree

spacy/pipeline/entity_linker.py

Lines changed: 70 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -300,77 +300,77 @@ def predict(self, docs: Iterable[Doc]) -> List[str]:
300300
for i, doc in enumerate(docs):
301301
sentences = [s for s in doc.sents]
302302
if len(doc) > 0:
303-
# Looping through each sentence and each entity
304-
# This may go wrong if there are entities across sentences - which shouldn't happen normally.
305-
for sent_index, sent in enumerate(sentences):
306-
if sent.ents:
307-
# get n_neightbour sentences, clipped to the length of the document
308-
start_sentence = max(0, sent_index - self.n_sents)
309-
end_sentence = min(
310-
len(sentences) - 1, sent_index + self.n_sents
311-
)
312-
start_token = sentences[start_sentence].start
313-
end_token = sentences[end_sentence].end
314-
sent_doc = doc[start_token:end_token].as_doc()
315-
# currently, the context is the same for each entity in a sentence (should be refined)
316-
xp = self.model.ops.xp
317-
if self.incl_context:
318-
sentence_encoding = self.model.predict([sent_doc])[0]
319-
sentence_encoding_t = sentence_encoding.T
320-
sentence_norm = xp.linalg.norm(sentence_encoding_t)
321-
for ent in sent.ents:
322-
entity_count += 1
323-
if ent.label_ in self.labels_discard:
324-
# ignoring this entity - setting to NIL
325-
final_kb_ids.append(self.NIL)
326-
else:
327-
candidates = self.get_candidates(self.kb, ent)
328-
if not candidates:
329-
# no prediction possible for this entity - setting to NIL
330-
final_kb_ids.append(self.NIL)
331-
elif len(candidates) == 1:
332-
# shortcut for efficiency reasons: take the 1 candidate
333-
# TODO: thresholding
334-
final_kb_ids.append(candidates[0].entity_)
335-
else:
336-
random.shuffle(candidates)
337-
# set all prior probabilities to 0 if incl_prior=False
338-
prior_probs = xp.asarray(
339-
[c.prior_prob for c in candidates]
340-
)
341-
if not self.incl_prior:
342-
prior_probs = xp.asarray(
343-
[0.0 for _ in candidates]
344-
)
345-
scores = prior_probs
346-
# add in similarity from the context
347-
if self.incl_context:
348-
entity_encodings = xp.asarray(
349-
[c.entity_vector for c in candidates]
350-
)
351-
entity_norm = xp.linalg.norm(
352-
entity_encodings, axis=1
353-
)
354-
if len(entity_encodings) != len(prior_probs):
355-
raise RuntimeError(
356-
Errors.E147.format(
357-
method="predict",
358-
msg="vectors not of equal length",
359-
)
360-
)
361-
# cosine similarity
362-
sims = xp.dot(
363-
entity_encodings, sentence_encoding_t
364-
) / (sentence_norm * entity_norm)
365-
if sims.shape != prior_probs.shape:
366-
raise ValueError(Errors.E161)
367-
scores = (
368-
prior_probs + sims - (prior_probs * sims)
303+
# Looping through each entity (TODO: rewrite)
304+
for ent in doc.ents:
305+
sent = ent.sent
306+
sent_index = sentences.index(sent)
307+
assert sent_index >= 0
308+
# get n_neightbour sentences, clipped to the length of the document
309+
start_sentence = max(0, sent_index - self.n_sents)
310+
end_sentence = min(
311+
len(sentences) - 1, sent_index + self.n_sents
312+
)
313+
start_token = sentences[start_sentence].start
314+
end_token = sentences[end_sentence].end
315+
sent_doc = doc[start_token:end_token].as_doc()
316+
# currently, the context is the same for each entity in a sentence (should be refined)
317+
xp = self.model.ops.xp
318+
if self.incl_context:
319+
sentence_encoding = self.model.predict([sent_doc])[0]
320+
sentence_encoding_t = sentence_encoding.T
321+
sentence_norm = xp.linalg.norm(sentence_encoding_t)
322+
entity_count += 1
323+
if ent.label_ in self.labels_discard:
324+
# ignoring this entity - setting to NIL
325+
final_kb_ids.append(self.NIL)
326+
else:
327+
candidates = self.get_candidates(self.kb, ent)
328+
if not candidates:
329+
# no prediction possible for this entity - setting to NIL
330+
final_kb_ids.append(self.NIL)
331+
elif len(candidates) == 1:
332+
# shortcut for efficiency reasons: take the 1 candidate
333+
# TODO: thresholding
334+
final_kb_ids.append(candidates[0].entity_)
335+
else:
336+
random.shuffle(candidates)
337+
# set all prior probabilities to 0 if incl_prior=False
338+
prior_probs = xp.asarray(
339+
[c.prior_prob for c in candidates]
340+
)
341+
if not self.incl_prior:
342+
prior_probs = xp.asarray(
343+
[0.0 for _ in candidates]
344+
)
345+
scores = prior_probs
346+
# add in similarity from the context
347+
if self.incl_context:
348+
entity_encodings = xp.asarray(
349+
[c.entity_vector for c in candidates]
350+
)
351+
entity_norm = xp.linalg.norm(
352+
entity_encodings, axis=1
353+
)
354+
if len(entity_encodings) != len(prior_probs):
355+
raise RuntimeError(
356+
Errors.E147.format(
357+
method="predict",
358+
msg="vectors not of equal length",
369359
)
370-
# TODO: thresholding
371-
best_index = scores.argmax().item()
372-
best_candidate = candidates[best_index]
373-
final_kb_ids.append(best_candidate.entity_)
360+
)
361+
# cosine similarity
362+
sims = xp.dot(
363+
entity_encodings, sentence_encoding_t
364+
) / (sentence_norm * entity_norm)
365+
if sims.shape != prior_probs.shape:
366+
raise ValueError(Errors.E161)
367+
scores = (
368+
prior_probs + sims - (prior_probs * sims)
369+
)
370+
# TODO: thresholding
371+
best_index = scores.argmax().item()
372+
best_candidate = candidates[best_index]
373+
final_kb_ids.append(best_candidate.entity_)
374374
if not (len(final_kb_ids) == entity_count):
375375
err = Errors.E147.format(
376376
method="predict", msg="result variables not of equal length"

spacy/tests/regression/test_issue7065.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1+
from spacy.kb import KnowledgeBase
12
from spacy.lang.en import English
3+
from spacy.training import Example
24

35

46
def test_issue7065():
@@ -16,3 +18,58 @@ def test_issue7065():
1618
ent = doc.ents[0]
1719
assert ent.start < sent0.end < ent.end
1820
assert sentences.index(ent.sent) == 0
21+
22+
23+
def test_issue7065_b():
24+
# Test that the NEL doesn't crash when an entity crosses a sentence boundary
25+
nlp = English()
26+
vector_length = 3
27+
nlp.add_pipe("sentencizer")
28+
29+
text = "Mahler 's Symphony No. 8 was beautiful."
30+
entities = [(0, 6, "PERSON"), (10, 24, "WORK")]
31+
links = {(0, 6): {"Q7304": 1.0, "Q270853": 0.0},
32+
(10, 24): {"Q7304": 0.0, "Q270853": 1.0}}
33+
sent_starts = [1, -1, 0, 0, 0, 0, 0, 0, 0]
34+
doc = nlp(text)
35+
example = Example.from_dict(doc, {"entities": entities, "links": links, "sent_starts": sent_starts})
36+
train_examples = [example]
37+
38+
def create_kb(vocab):
39+
# create artificial KB
40+
mykb = KnowledgeBase(vocab, entity_vector_length=vector_length)
41+
mykb.add_entity(entity="Q270853", freq=12, entity_vector=[9, 1, -7])
42+
mykb.add_alias(
43+
alias="No. 8",
44+
entities=["Q270853"],
45+
probabilities=[1.0],
46+
)
47+
mykb.add_entity(entity="Q7304", freq=12, entity_vector=[6, -4, 3])
48+
mykb.add_alias(
49+
alias="Mahler",
50+
entities=["Q7304"],
51+
probabilities=[1.0],
52+
)
53+
return mykb
54+
55+
# Create the Entity Linker component and add it to the pipeline
56+
entity_linker = nlp.add_pipe("entity_linker", last=True)
57+
entity_linker.set_kb(create_kb)
58+
59+
# train the NEL pipe
60+
optimizer = nlp.initialize(get_examples=lambda: train_examples)
61+
for i in range(2):
62+
losses = {}
63+
nlp.update(train_examples, sgd=optimizer, losses=losses)
64+
65+
# Add a custom rule-based component to mimick NER
66+
patterns = [
67+
{"label": "PERSON", "pattern": [{"LOWER": "mahler"}]},
68+
{"label": "WORK", "pattern": [{"LOWER": "symphony"}, {"LOWER": "no"}, {"LOWER": "."}, {"LOWER": "8"}]}
69+
]
70+
ruler = nlp.add_pipe("entity_ruler", before="entity_linker")
71+
ruler.add_patterns(patterns)
72+
73+
# test the trained model - this should not throw E148
74+
doc = nlp(text)
75+
assert doc

0 commit comments

Comments
 (0)