Fix KeyedVectors.add matrix type

gensim/models/keyedvectors.py

@@ -214,8 +214,8 @@ def __str__(self):
 
 class BaseKeyedVectors(utils.SaveLoad):
     """Abstract base class / interface for various types of word vectors."""
-    def __init__(self, vector_size):
-        self.vectors = zeros((0, vector_size))
+    def __init__(self, vector_size, dtype=REAL):
+        self.vectors = zeros((0, vector_size), dtype=dtype)
         self.vocab = {}
         self.vector_size = vector_size
         self.index2entity = []
@@ -308,7 +308,7 @@ def add(self, entities, weights, replace=False):
             self.index2entity.append(entity)
 
         # add vectors for new entities
-        self.vectors = vstack((self.vectors, weights[~in_vocab_mask]))
+        self.vectors = vstack((self.vectors, weights[~in_vocab_mask].astype(self.vectors.dtype)))
 
         # change vectors for in_vocab entities if `replace` flag is specified
         if replace:
@@ -376,8 +376,8 @@ def rank(self, entity1, entity2):
 
 class WordEmbeddingsKeyedVectors(BaseKeyedVectors):
     """Class containing common methods for operations over word vectors."""
-    def __init__(self, vector_size):
-        super(WordEmbeddingsKeyedVectors, self).__init__(vector_size=vector_size)
+    def __init__(self, vector_size, dtype=REAL):
+        super(WordEmbeddingsKeyedVectors, self).__init__(vector_size=vector_size, dtype=REAL)
         self.vectors_norm = None
         self.index2word = []
 
@@ -1550,8 +1550,8 @@ def load(cls, fname_or_handle, **kwargs):
 
 class Doc2VecKeyedVectors(BaseKeyedVectors):
 
-    def __init__(self, vector_size, mapfile_path):
-        super(Doc2VecKeyedVectors, self).__init__(vector_size=vector_size)
+    def __init__(self, vector_size, mapfile_path, dtype=REAL):
+        super(Doc2VecKeyedVectors, self).__init__(vector_size=vector_size, dtype=REAL)
         self.doctags = {}  # string -> Doctag (only filled if necessary)
         self.max_rawint = -1  # highest rawint-indexed doctag
         self.offset2doctag = []  # int offset-past-(max_rawint+1) -> String (only filled if necessary)

gensim/models/poincare.py

@@ -866,8 +866,8 @@ class PoincareKeyedVectors(BaseKeyedVectors):
     Used to perform operations on the vectors such as vector lookup, distance calculations etc.
 
     """
-    def __init__(self, vector_size):
-        super(PoincareKeyedVectors, self).__init__(vector_size)
+    def __init__(self, vector_size, dtype=REAL):
+        super(PoincareKeyedVectors, self).__init__(vector_size, dtype=REAL)
         self.max_distance = 0
         self.index2word = []
         self.vocab = {}

gensim/test/test_keyedvectors.py

@@ -15,7 +15,7 @@
 import numpy as np
 
 from gensim.corpora import Dictionary
-from gensim.models.keyedvectors import KeyedVectors as EuclideanKeyedVectors, WordEmbeddingSimilarityIndex, \
+from gensim.models.keyedvectors import KeyedVectors, WordEmbeddingSimilarityIndex, \
     FastTextKeyedVectors
 from gensim.test.utils import datapath
 
@@ -27,7 +27,7 @@
 
 class TestWordEmbeddingSimilarityIndex(unittest.TestCase):
     def setUp(self):
-        self.vectors = EuclideanKeyedVectors.load_word2vec_format(
+        self.vectors = KeyedVectors.load_word2vec_format(
             datapath('euclidean_vectors.bin'), binary=True, datatype=np.float64)
 
     def test_most_similar(self):
@@ -70,9 +70,9 @@ def test_most_similar(self):
         self.assertTrue(np.allclose(first_similarities**2.0, second_similarities))
 
 
-class TestEuclideanKeyedVectors(unittest.TestCase):
+class TestKeyedVectors(unittest.TestCase):
     def setUp(self):
-        self.vectors = EuclideanKeyedVectors.load_word2vec_format(
+        self.vectors = KeyedVectors.load_word2vec_format(
             datapath('euclidean_vectors.bin'), binary=True, datatype=np.float64)
 
     def test_similarity_matrix(self):
@@ -227,7 +227,7 @@ def test_add_single(self):
             self.assertTrue(np.allclose(self.vectors[ent], vector))
 
         # Test `add` on empty kv.
-        kv = EuclideanKeyedVectors(self.vectors.vector_size)
+        kv = KeyedVectors(self.vectors.vector_size)
         for ent, vector in zip(entities, vectors):
             kv.add(ent, vector)
 
@@ -248,13 +248,23 @@ def test_add_multiple(self):
             self.assertTrue(np.allclose(self.vectors[ent], vector))
 
         # Test `add` on empty kv.
-        kv = EuclideanKeyedVectors(self.vectors.vector_size)
+        kv = KeyedVectors(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_add_type(self):
+        dtype = np.float32
+        kv = KeyedVectors(2, dtype=dtype)
+        words, vectors = ["a"], np.array([1., 1.], dtype=np.float16).reshape(1, -1)
+
+        assert kv.vectors.dtype == dtype
+
+        kv.add(words, vectors)
+        assert kv.vectors.dtype == dtype
+
     def test_set_item(self):
         """Test that __setitem__ works correctly."""
         vocab_size = len(self.vectors.vocab)
@@ -287,7 +297,7 @@ def test_set_item(self):
             self.assertTrue(np.allclose(self.vectors[ent], vector))
 
     def test_ft_kv_backward_compat_w_360(self):
-        kv = EuclideanKeyedVectors.load(datapath("ft_kv_3.6.0.model.gz"))
+        kv = KeyedVectors.load(datapath("ft_kv_3.6.0.model.gz"))
         ft_kv = FastTextKeyedVectors.load(datapath("ft_kv_3.6.0.model.gz"))
 
         expected = ['trees', 'survey', 'system', 'graph', 'interface']