use StringStore

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).

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 = <int64_t>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 = <int64_t>self._alias_index.get(alias_hash)
         return self._aliases_table[alias_index]