Make most_similar accept any integral topn

piskvorky · May 17, 2019 · b34b06e · b34b06e
1 parent d910fa9
commit b34b06e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/gensim/models/keyedvectors.py b/gensim/models/keyedvectors.py
@@ -162,6 +162,7 @@
 
 from itertools import chain
 import logging
+from numbers import Integral
 
 try:
     from queue import Queue, Empty
@@ -519,7 +520,7 @@ def most_similar(self, positive=None, negative=None, topn=10, restrict_vocab=Non
             one-dimensional numpy array with the size of the vocabulary.
 
         """
-        if isinstance(topn, int) and topn < 1:
+        if isinstance(topn, Integral) and topn < 1:
             return []
 
         if positive is None:
@@ -807,7 +808,7 @@ def most_similar_cosmul(self, positive=None, negative=None, topn=10):
             one-dimensional numpy array with the size of the vocabulary.
 
         """
-        if isinstance(topn, int) and topn < 1:
+        if isinstance(topn, Integral) and topn < 1:
             return []
 
         if positive is None:
@@ -1678,7 +1679,7 @@ def most_similar(self, positive=None, negative=None, topn=10, clip_start=0, clip
             Sequence of (doctag/index, similarity).
 
         """
-        if isinstance(topn, int) and topn < 1:
+        if isinstance(topn, Integral) and topn < 1:
             return []
 
         if positive is None:

diff --git a/gensim/models/poincare.py b/gensim/models/poincare.py
@@ -44,6 +44,7 @@
 
 import csv
 import logging
+from numbers import Integral
 import sys
 import time
 
@@ -1219,7 +1220,7 @@ def most_similar(self, node_or_vector, topn=10, restrict_vocab=None):
             [(u'kangaroo.n.01', 0.0), (u'marsupial.n.01', 0.26524229460827725)]
 
         """
-        if isinstance(topn, int) and topn < 1:
+        if isinstance(topn, Integral) and topn < 1:
             return []
 
         if not restrict_vocab:

diff --git a/gensim/test/test_keyedvectors.py b/gensim/test/test_keyedvectors.py
@@ -109,6 +109,9 @@ def test_most_similar_topn(self):
         predicted = self.vectors.most_similar('war', topn=0)
         self.assertEqual(len(predicted), 0)
 
+        predicted = self.vectors.most_similar('war', topn=np.uint8(0))
+        self.assertEqual(len(predicted), 0)
+
     def test_relative_cosine_similarity(self):
         """Test relative_cosine_similarity returns expected results with an input of a word pair and topn"""
         wordnet_syn = [