check the length of entities and probabilities vector + unit test

This commit is contained in:
svlandeg 2019-03-19 21:55:10 +01:00
parent d133ffaff9
commit a9074e0886
2 changed files with 22 additions and 4 deletions

View File

@ -36,11 +36,18 @@ cdef class KnowledgeBase:
def add_alias(self, unicode alias, entities, probabilities):
"""For a given alias, add its potential entities and prior probabilies to the KB."""
# Throw an error if the length of entities and probabilities are not the same
if not len(entities) == len(probabilities):
raise ValueError("The vectors for entities and probabilities for alias '" + alias
+ "' should have equal length, but found "
+ str(len(entities)) + " and " + str(len(probabilities)) + "respectively.")
# Throw an error if the probabilities sum up to more than 1
prob_sum = sum(probabilities)
if prob_sum > 1:
raise ValueError("The sum of prior probabilities for alias '" + alias + "' should not exceed 1, "
"but found " + str(prob_sum))
+ "but found " + str(prob_sum))
cdef hash_t alias_hash = self.strings.add(alias)
@ -63,9 +70,6 @@ cdef class KnowledgeBase:
entry_indices.push_back(int(entry_index))
probs.push_back(float(prob))
# TODO: check sum(probabilities) <= 1
# TODO: check len(entities) == len(probabilities)
self.c_add_aliases(alias_key=alias_hash, entry_indices=entry_indices, probs=probs)

View File

@ -49,3 +49,17 @@ def test_kb_invalid_probabilities():
with pytest.raises(ValueError):
mykb.add_alias(alias="douglassss", entities=["Q2", "Q3"], probabilities=[0.8, 0.4])
def test_kb_invalid_combination():
"""Test the invalid construction of a KB with non-matching entity and probability lists"""
mykb = KnowledgeBase()
# adding entities
mykb.add_entity(entity_id="Q1", prob=0.9)
mykb.add_entity(entity_id="Q2", prob=0.2)
mykb.add_entity(entity_id="Q3", prob=0.5)
# adding aliases - should fail because the entities and probabilities vectors are not of equal length
with pytest.raises(ValueError):
mykb.add_alias(alias="douglassss", entities=["Q2", "Q3"], probabilities=[0.3, 0.4, 0.1])