Add gensim.models.BaseKeyedVectors.add_entity method for fill KeyedVectors in manual way. Fix #1942

gensim/models/keyedvectors.py

@@ -154,6 +154,75 @@ def get_vector(self, entity):
         else:
             raise KeyError("'%s' not in vocabulary" % entity)
 
+    def add_entity(self, entity, weights, replace=False):
+        """Add entity vector in a manual way.
+        If `entity` is already in the vocabulary, old vector is keeped unless `replace` flag is True.
+
+        Parameters
+        ----------
+        entity : str
+            Entity specified by string tag.
+        weights : np.array
+            1D numpy array with shape (`vector_size`,)
+        replace: bool, optional
+            Boolean flag indicating whether to replace old vector if entity is already in the vocabulary.
+            Default, False, means that old vector is keeped.
+        """
+        self.add_entities([entity], weights.reshape(1, -1), replace=replace)
+
+    def add_entities(self, entities, weights, replace=False):
+        """Add entities and theirs vectors in a manual way.
+        If some entity is already in the vocabulary, old vector is keeped unless `replace` flag is True.
+
+        Parameters
+        ----------
+        entities : list of str
+            Entities specified by string tags.
+        weights: list of np.array or np.array
+            List of 1D np.array vectors or 2D np.array of vectors.
+        replace: bool, optional
+            Boolean flag indicating whether to replace vectors for entities which are already in the vocabulary.
+            Default, False, means that old vectors for those entities are keeped.
+        """
+        if isinstance(weights, list):
+            weights = np.array(weights)
+
+        in_vocab_mask = np.zeros(len(entities), dtype=np.bool)
+        in_vocab_idxs = []
+        out_vocab_entities = []
+
+        for idx, entity in zip(range(len(entities)), entities):
+            if entity in self.vocab:
+                in_vocab_mask[idx] = True
+                in_vocab_idxs.append(self.vocab[entity].index)
+            else:
+                out_vocab_entities.append(entity)
+
+        # add new entities to the vocab
+        for entity in out_vocab_entities:
+            entity_id = len(self.vocab)
+            self.vocab[entity] = Vocab(index=entity_id, count=1)
+            self.index2entity.append(entity)
+
+        # add vectors for new entities
+        if len(self.vectors) == 0:
+            self.vectors = weights[~in_vocab_mask]
+        else:
+            self.vectors = vstack((self.vectors, weights[~in_vocab_mask]))
+
+        # change vectors for in_vocab entities if `replace` flag is specified
+        if replace:
+            self.vectors[in_vocab_idxs] = weights[in_vocab_mask]
+
+    def __setitem__(self, entities, weights):
+        """Idiomatic way to call `add_entities` with `replace=True`.
+        """
+        if not isinstance(entities, list):
+            entities = [entities]
+            weights = weights.reshape(1, -1)
+
+        self.add_entities(entities, weights, replace=True)
+
     def __getitem__(self, entities):
         """
         Accept a single entity (string tag) or list of entities as input.

gensim/test/test_keyedvectors.py

@@ -153,6 +153,78 @@ def test_wv_property(self):
         """Test that the deprecated `wv` property returns `self`. To be removed in v4.0.0."""
         self.assertTrue(self.vectors is self.vectors.wv)
 
+    def test_add_entity(self):
+        """Test that adding entity in a manual way works correctly."""
+        entities = ['___some_entity{}_not_present_in_keyed_vectors___'.format(i) for i in range(5)]
+        vectors = [np.random.randn(self.vectors.vector_size) for _ in range(5)]
+
+        # Test `add_entity` on already filled kv.
+        for ent, vector in zip(entities, vectors):
+            self.vectors.add_entity(ent, vector)
+
+        for ent, vector in zip(entities, vectors):
+            self.assertTrue(np.allclose(self.vectors[ent], vector))
+
+        # Test `add_entity` on empty kv.
+        kv = EuclideanKeyedVectors(self.vectors.vector_size)
+        for ent, vector in zip(entities, vectors):
+            kv.add_entity(ent, vector)
+
+        for ent, vector in zip(entities, vectors):
+            self.assertTrue(np.allclose(kv[ent], vector))
+
+    def test_add_entities(self):
+        """Test that adding a bulk of entities in a manual way works correctly."""
+        entities = ['___some_entity{}_not_present_in_keyed_vectors___'.format(i) for i in range(5)]
+        vectors = [np.random.randn(self.vectors.vector_size) for _ in range(5)]
+
+        # Test `add_entities` on already filled kv.
+        vocab_size = len(self.vectors.vocab)
+        self.vectors.add_entities(entities, vectors, replace=False)
+        self.assertEqual(vocab_size + len(entities), len(self.vectors.vocab))
+
+        for ent, vector in zip(entities, vectors):
+            self.assertTrue(np.allclose(self.vectors[ent], vector))
+
+        # Test `add_entities` on empty kv.
+        kv = EuclideanKeyedVectors(self.vectors.vector_size)
+        kv[entities] = vectors
+        self.assertEqual(len(kv.vocab), len(entities))
+
+        for ent, vector in zip(entities, vectors):
+            self.assertTrue(np.allclose(kv[ent], vector))
+
+    def test_set_item(self):
+        """Test that __setitem__ works correctly."""
+        vocab_size = len(self.vectors.vocab)
+
+        # Add new entity.
+        entity = '___some_new_entity___'
+        vector = np.random.randn(self.vectors.vector_size)
+        self.vectors[entity] = vector
+
+        self.assertEqual(len(self.vectors.vocab), vocab_size + 1)
+        self.assertTrue(np.allclose(self.vectors[entity], vector))
+
+        # Replace vector for entity in vocab.
+        vocab_size = len(self.vectors.vocab)
+        vector = np.random.randn(self.vectors.vector_size)
+        self.vectors['war'] = vector
+
+        self.assertEqual(len(self.vectors.vocab), vocab_size)
+        self.assertTrue(np.allclose(self.vectors['war'], vector))
+
+        # __setitem__ on several entities.
+        vocab_size = len(self.vectors.vocab)
+        entities = ['war', '___some_new_entity1___', '___some_new_entity2___', 'terrorism', 'conflict']
+        vectors = [np.random.randn(self.vectors.vector_size) for _ in range(len(entities))]
+
+        self.vectors[entities] = vectors
+
+        self.assertEqual(len(self.vectors.vocab), vocab_size + 2)
+        for ent, vector in zip(entities, vectors):
+            self.assertTrue(np.allclose(self.vectors[ent], vector))
+
 
 if __name__ == '__main__':
     logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.DEBUG)