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)