diff --git a/spacy/kb.pyx b/spacy/kb.pyx index 969b43f6d..ea23e5373 100644 --- a/spacy/kb.pyx +++ b/spacy/kb.pyx @@ -39,13 +39,15 @@ cdef class KnowledgeBase: cdef vector[float] probs for entity, prob in zip(entities, probabilities): - entity_hash = self.strings.add(entity) + entity_hash = self.strings[entity] + if not entity_hash in self._entry_index: + raise ValueError("Alias '" + alias + "' defined for unknown entity '" + entity + "'") + entry_index = self._entry_index.get(entity_hash) entry_indices.push_back(int(entry_index)) probs.push_back(float(prob)) # TODO: check that alias hadn't been defined before - # TODO: check that entity is already in this KB (entity_index is OK) # TODO: check sum(probabilities) <= 1 # TODO: check len(entities) == len(probabilities) diff --git a/spacy/sandbox_test_sofie/testing_el.py b/spacy/sandbox_test_sofie/testing_el.py index b6255f9f9..b5b529d4b 100644 --- a/spacy/sandbox_test_sofie/testing_el.py +++ b/spacy/sandbox_test_sofie/testing_el.py @@ -7,6 +7,10 @@ def create_kb(): print("kb size", len(mykb), mykb.get_size_entities(), mykb.get_size_aliases()) # adding entities + entity_0 = "Q0" # douglas adams + mykb.add_entity(entity_id=entity_0, prob=0.5) + print(" adding entity", entity_0) + entity_42 = "Q42" # douglas adams mykb.add_entity(entity_id=entity_42, prob=0.5) print(" adding entity", entity_42) @@ -18,7 +22,7 @@ def create_kb(): print("kb size", len(mykb), mykb.get_size_entities(), mykb.get_size_aliases()) # adding aliases - alias = "douglas" + alias = "douglassss" print(" adding alias", alias) mykb.add_alias(alias=alias, entities=["Q42", "Q5301561"], probabilities=[0.8, 0.2]) diff --git a/spacy/tests/pipeline/test_el.py b/spacy/tests/pipeline/test_el.py new file mode 100644 index 000000000..ed88076ce --- /dev/null +++ b/spacy/tests/pipeline/test_el.py @@ -0,0 +1,29 @@ +import pytest + +from spacy.kb import KnowledgeBase + + +def test_kb_valid_entities(): + mykb = KnowledgeBase() + + # adding entities + mykb.add_entity(entity_id="Q1", prob=0.5) + mykb.add_entity(entity_id="Q2", prob=0.5) + mykb.add_entity(entity_id="Q3", prob=0.5) + + # adding aliases + mykb.add_alias(alias="douglassss", entities=["Q2", "Q3"], probabilities=[0.8, 0.2]) + + +def test_kb_invalid_entities(): + mykb = KnowledgeBase() + + # adding entities + mykb.add_entity(entity_id="Q1", prob=0.5) + mykb.add_entity(entity_id="Q2", prob=0.5) + mykb.add_entity(entity_id="Q3", prob=0.5) + + # adding aliases - should fail because one of the given IDs is not valid + with pytest.raises(ValueError): + mykb.add_alias(alias="douglassss", entities=["Q2", "Q342"], probabilities=[0.8, 0.2]) +