diff --git a/spacy/kb.pxd b/spacy/kb.pxd index 9f0a5e68d..f4f60d478 100644 --- a/spacy/kb.pxd +++ b/spacy/kb.pxd @@ -3,8 +3,9 @@ from cymem.cymem cimport Pool from preshed.maps cimport PreshMap from libcpp.vector cimport vector from libc.stdint cimport int32_t, int64_t + +from spacy.strings cimport StringStore from .typedefs cimport hash_t -from .strings cimport hash_string # Internal struct, for storage and disambiguation. This isn't what we return @@ -41,6 +42,7 @@ cdef struct _AliasC: cdef class KnowledgeBase: cdef Pool mem + cpdef readonly StringStore strings # This maps 64bit keys (hash of unique entity string) # to 64bit values (position of the _EntryC struct in the _entries vector). diff --git a/spacy/kb.pyx b/spacy/kb.pyx index 854feb069..969b43f6d 100644 --- a/spacy/kb.pyx +++ b/spacy/kb.pyx @@ -1,7 +1,5 @@ # cython: profile=True # coding: utf8 -from preshed.maps import PreshMap - cdef class KnowledgeBase: @@ -9,6 +7,7 @@ cdef class KnowledgeBase: self._entry_index = PreshMap() self._alias_index = PreshMap() self.mem = Pool() + self.strings = StringStore() def __len__(self): return self.get_size_entities() @@ -20,13 +19,12 @@ cdef class KnowledgeBase: return self._aliases_table.size() def add_entity(self, unicode entity_id, float prob, vectors=None, features=None): - cdef hash_t id_hash = hash_string(entity_id) + cdef hash_t id_hash = self.strings.add(entity_id) # TODO: more friendly check for non-unique name if id_hash in self._entry_index: return - cdef int32_t dummy_value = 342 self.c_add_entity(entity_key=id_hash, prob=prob, vector_rows=&dummy_value, feats_row=dummy_value) # TODO self._vectors_table.get_pointer(vectors), @@ -34,14 +32,14 @@ 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.""" - cdef hash_t alias_hash = hash_string(alias) + cdef hash_t alias_hash = self.strings.add(alias) cdef hash_t entity_hash cdef vector[int64_t] entry_indices cdef vector[float] probs for entity, prob in zip(entities, probabilities): - entity_hash = hash_string(entity) + entity_hash = self.strings.add(entity) entry_index = self._entry_index.get(entity_hash) entry_indices.push_back(int(entry_index)) probs.push_back(float(prob)) @@ -54,7 +52,7 @@ cdef class KnowledgeBase: self.c_add_aliases(alias_key=alias_hash, entry_indices=entry_indices, probs=probs) def get_candidates(self, unicode alias): - cdef hash_t alias_hash = hash_string(alias) + cdef hash_t alias_hash = self.strings.add(alias) alias_index = self._alias_index.get(alias_hash) return self._aliases_table[alias_index]