doc2vec: allow indexing with np.int64 -- fixes #1231

gensim/models/doc2vec.py

@@ -55,7 +55,7 @@
 
 from numpy import zeros, random, sum as np_sum, add as np_add, concatenate, \
     repeat as np_repeat, array, float32 as REAL, empty, ones, memmap as np_memmap, \
-    sqrt, newaxis, ndarray, dot, vstack, dtype, divide as np_divide
+    sqrt, newaxis, ndarray, dot, vstack, dtype, divide as np_divide, integer
 
 
 from gensim.utils import call_on_class_only
@@ -319,7 +319,7 @@ def trained_item(self, indexed_tuple):
 
     def _int_index(self, index):
         """Return int index for either string or int index"""
-        if isinstance(index, int):
+        if isinstance(index, (int, integer)):
             return index
         else:
             return self.max_rawint + 1 + self.doctags[index].offset
@@ -347,7 +347,7 @@ def __getitem__(self, index):
         If a list, return designated tags' vector representations as a
         2D numpy array: #tags x #vector_size.
         """
-        if isinstance(index, string_types + (int,)):
+        if isinstance(index, string_types + (int, integer)):
             return self.doctag_syn0[self._int_index(index)]
 
         return vstack([self[i] for i in index])

gensim/test/test_doc2vec.py

@@ -114,6 +114,7 @@ def test_int_doctags(self):
         model.build_vocab(corpus)
         self.assertEqual(len(model.docvecs.doctag_syn0), 300)
         self.assertEqual(model.docvecs[0].shape, (100,))
+        self.assertEqual(model.docvecs[np.int64(0)].shape, (100,))
         self.assertRaises(KeyError, model.__getitem__, '_*0')
 
     def test_missing_string_doctag(self):