Batch model queries for faster for entity prior feature lookup

semuar · Apr 27, 2016 · a366d15 · a366d15
1 parent b28958a
commit a366d15
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/nel/features/probability.py b/nel/features/probability.py
@@ -22,8 +22,12 @@ def __init__(self, entity_model_tag):
         self.tag = entity_model_tag
         self.em = disambiguation.EntityCounts(self.tag)
 
+    def compute_doc_state(self, doc):
+        candidates = set(c.id for chain in doc.chains for c in chain.candidates)
+        return dict(self.em.iter_counts(candidates))
+
     def compute_raw(self, doc, chain, candidate, state):
-        return self.em.count(candidate.id) + 0.1
+        return state[candidate.id] + 0.1
 
     @classmethod
     def add_arguments(cls, p):