From 8c7fa11d4e2332fdd69879ea01526e86bbc244b5 Mon Sep 17 00:00:00 2001
From: Ben Frederickson <github@benfrederickson.com>
Date: Sat, 13 Jul 2019 09:35:12 -0700
Subject: [PATCH 1/2] Add pickle support for nearest neighbours models

---
 implicit/nearest_neighbours.py | 13 +++++++++++++
 tests/approximate_als_test.py  | 18 ++++++++++++++++++
 tests/recommender_base_test.py | 10 ++++++++++
 3 files changed, 41 insertions(+)

diff --git a/implicit/nearest_neighbours.py b/implicit/nearest_neighbours.py
index adda633e..98b0d695 100644
--- a/implicit/nearest_neighbours.py
+++ b/implicit/nearest_neighbours.py
@@ -87,6 +87,19 @@ def similar_items(self, itemid, N=10):
 
         return sorted(list(nonzeros(self.similarity, itemid)), key=lambda x: -x[1])[:N]
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # scorer isn't picklable
+        del state['scorer']
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        if self.similarity is not None:
+            self.scorer = NearestNeighboursScorer(self.similarity)
+        else:
+            self.scorer = None
+
     def save(self, filename):
         m = self.similarity
         numpy.savez(filename, data=m.data, indptr=m.indptr, indices=m.indices, shape=m.shape,
diff --git a/tests/approximate_als_test.py b/tests/approximate_als_test.py
index 3dc4475a..12836e90 100644
--- a/tests/approximate_als_test.py
+++ b/tests/approximate_als_test.py
@@ -15,6 +15,11 @@
     class AnnoyALSTest(unittest.TestCase, TestRecommenderBaseMixin):
         def _get_model(self):
             return AnnoyAlternatingLeastSquares(factors=2, regularization=0, use_gpu=False)
+
+        def test_pickle(self):
+            # pickle isn't supported on annoy indices
+            pass
+
 except ImportError:
     pass
 
@@ -25,6 +30,11 @@ class NMSLibALSTest(unittest.TestCase, TestRecommenderBaseMixin):
         def _get_model(self):
             return NMSLibAlternatingLeastSquares(factors=2, regularization=0,
                                                  index_params={'post': 2}, use_gpu=False)
+
+        def test_pickle(self):
+            # pickle isn't supported on nmslib indices
+            pass
+
 except ImportError:
     pass
 
@@ -36,6 +46,10 @@ def _get_model(self):
             return FaissAlternatingLeastSquares(nlist=1, nprobe=1, factors=2, regularization=0,
                                                 use_gpu=False)
 
+        def test_pickle(self):
+            # pickle isn't supported on faiss indices
+            pass
+
     if HAS_CUDA:
         class FaissALSGPUTest(unittest.TestCase, TestRecommenderBaseMixin):
             __regularization = 0
@@ -69,6 +83,10 @@ def test_large_recommend(self):
                 recs = model.recommend(0, plays.T.tocsr(), N=1050)
                 self.assertEqual(recs[0][0], 0)
 
+            def test_pickle(self):
+                # pickle isn't supported on faiss indices
+                pass
+
 except ImportError:
     pass
 
diff --git a/tests/recommender_base_test.py b/tests/recommender_base_test.py
index 381a348e..4b163f34 100644
--- a/tests/recommender_base_test.py
+++ b/tests/recommender_base_test.py
@@ -2,6 +2,8 @@
 
 from __future__ import print_function
 
+import pickle
+
 import numpy as np
 from scipy.sparse import csr_matrix
 
@@ -157,6 +159,14 @@ def test_rank_items(self):
                 wrong_item_list = selected_items + wrong_pos_items
                 model.rank_items(userid, user_items, wrong_item_list)
 
+    def test_pickle(self):
+        item_users = self.get_checker_board(50)
+        model = self._get_model()
+        model.fit(item_users, show_progress=False)
+
+        pickled = pickle.dumps(model)
+        pickle.loads(pickled)
+
     def get_checker_board(self, X):
         """ Returns a 'checkerboard' matrix: where every even userid has liked
         every even itemid and every odd userid has liked every odd itemid.

From c352d08b44f2f7321fac5c3e48a7f4fea100a38a Mon Sep 17 00:00:00 2001
From: Ben Frederickson <github@benfrederickson.com>
Date: Sat, 13 Jul 2019 09:39:35 -0700
Subject: [PATCH 2/2] travis ci fix

nmslib no longer supports python2, causing CI errors for us. Fix by only testing
out nmslib support for python 3.
---
 .travis.yml | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 7036efee..5ce9ae53 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -40,9 +40,13 @@ before_install:
     fi
   fi
 install:
-  - travis_wait travis_retry $PIP uninstall numpy -y
-  - travis_wait travis_retry $PIP install -r requirements.txt --ignore-installed flake8 isort cpplint nmslib faiss annoy
-  - travis_retry $PIP install -e .
+- |
+  travis_wait travis_retry $PIP uninstall numpy -y
+  travis_wait travis_retry $PIP install -r requirements.txt --ignore-installed flake8 isort cpplint faiss annoy
+  if  [ "${PYTHON:0:1}" = "3" ]; then
+     travis_wait travis_retry $PIP install nmslib
+  fi
+  travis_retry $PIP install -e .
 
 script:
   - flake8