New KeyedVectors.vectors_for_all method for vectorizing all words in a dictionary

gensim/models/keyedvectors.py

@@ -171,7 +171,7 @@
 import itertools
 import warnings
 from numbers import Integral
-from typing import Iterable
+from typing import Iterable, Union
 
 from numpy import (
     dot, float32 as REAL, double, array, zeros, vstack,
@@ -1696,8 +1696,8 @@ def intersect_word2vec_format(self, fname, lockf=0.0, binary=False, encoding='ut
             msg=f"merged {overlap_count} vectors into {self.vectors.shape} matrix from {fname}",
         )
 
-    def vectors_for_all(self, keys: Iterable) -> 'KeyedVectors':
-        """Produces vectors for all keys in a given iterable.
+    def vectors_for_all(self, keys: Union[Iterable, Dictionary]) -> 'KeyedVectors':
+        """Produces vectors for all given keys.
 
         Notes
         -----
@@ -1713,7 +1713,7 @@ def vectors_for_all(self, keys: Iterable) -> 'KeyedVectors':
 
         Parameters
         ----------
-        keys : iterable of str
+        keys : {iterable of str, Dictionary}
             The keys that will be vectorized.
 
         Returns
@@ -1722,7 +1722,10 @@ def vectors_for_all(self, keys: Iterable) -> 'KeyedVectors':
             Vectors for all the given keys.
 
         """
-        vocabulary = sorted(set(filter(lambda key: key in self, keys)))
+        if isinstance(keys, Dictionary):
+            vocabulary = keys.token2id
+        else:
+            vocabulary = sorted(set(filter(lambda key: key in self, keys)))
         vocab_size = len(vocabulary)
         datatype = self.vectors.dtype
         kv = KeyedVectors(self.vector_size, vocab_size, dtype=datatype)

gensim/similarities/termsim.py

@@ -114,11 +114,10 @@ class WordEmbeddingSimilarityIndex(TermSimilarityIndex):
     >>> from gensim.models.word2vec import LineSentence
     >>> from gensim.similarities import WordEmbeddingSimilarityIndex
     >>>
-    >>> corpus = common_texts
-    >>> model = FastText(corpus, vector_size=20, min_count=1)  # train word-vectors on a corpus
+    >>> model = FastText(common_texts, vector_size=20, min_count=1)  # train word-vectors on a corpus
     >>> different_corpus = LineSentence(datapath('lee_background.cor'))
     >>> dictionary = Dictionary(different_corpus)  # construct a vocabulary on a different corpus
-    >>> word_vectors = model.wv.vectors_for_all(dictionary.token2id)  # remove OOV word-vectors and infer new words
+    >>> word_vectors = model.wv.vectors_for_all(dictionary)  # remove OOV word-vectors and infer new words
     >>> assert len(dictionary) == len(word_vectors)  # all words from our vocabulary received their word-vectors
     >>> termsim_index = WordEmbeddingSimilarityIndex(word_vectors)
 
@@ -433,7 +432,7 @@ class SparseTermSimilarityMatrix(SaveLoad):
     >>> model = Word2Vec(common_texts, vector_size=20, min_count=1)  # train word-vectors
     >>> annoy = AnnoyIndexer(model, num_trees=2)  # use annoy for faster word similarity lookups
     >>> dictionary = Dictionary(common_texts)
-    >>> word_vectors = model.wv.vectors_for_all(dictionary.token2id)
+    >>> word_vectors = model.wv.vectors_for_all(dictionary)
     >>> termsim_index = WordEmbeddingSimilarityIndex(word_vectors, kwargs={'indexer': annoy})
     >>> bow_corpus = [dictionary.doc2bow(document) for document in common_texts]
     >>> similarity_matrix = SparseTermSimilarityMatrix(termsim_index, dictionary, symmetric=True, dominant=True)