Repository: piskvorky/gensim

Issue:
Support custom pickle version in utils.SaveLoad
I was trying to load a word2vec model saved in Python 3.4 at Python 2.7 environment. Currently it fails to load due to the unrecognized pickle protocol version 4.  

After digging it down, the `Word2Vec.save()` is defined by `utils.SaveLoad` and it always calls the `utils.pickle` with default argument, which pass the protocol version `-1` implying the highest protocol available.

``` python
# utils.py
def pickle(obj, fname, protocol=-1):
     """Pickle object `obj` to file `fname`."""
     with smart_open(fname, 'wb') as fout: # 'b' for binary, needed on Windows
         _pickle.dump(obj, fout, protocol=protocol)

class SaveLoad(object):
    def save(self, fname, separately=None, sep_limit=10 * 1024**2, ignore=frozenset()):
        # ...
        pickle(self, fname)
```

To make saved models work on both Python 2.x and 3.x, the pickle protocol version should be set at 2, the highest version Python 2.x. If user can supply their desired protocol version, it would be easier to switch between Python versions.

Naively, I came up with the following api change:

``` python
# utils.py
class SaveLoad(object):
    def save(..., pickle_protocol=-1):
         pickle(self, fname, pickle_protocol)  # same for all places calling pickle
```

Rests of the model implementation (subclass of `SaveLoad`) remains unchanged. Take word2vec as example,

``` python
# models/word2vec.py
class Word2Vec(utils.SaveLoad):
     # ...
     def save(self, *args, **kwargs):
         kwargs['ignore'] = kwargs.get('ignore', ['syn0norm']) # don't bother storing the cached normalized vectors
         super(Word2Vec, self).save(*args, **kwargs)
     save.__doc__ = utils.SaveLoad.save.__doc__
```

Whether passing `word2vec.save(..., pickle_protocol=2)` or not doesn't break existed code.

Hope this ease the transition between Python 2 and 3. I could file a pull request if this is feasible.
### EDIT: Current solution

In Python 3.4, run the following code before saving will help convert the existed models

``` python
import gensim
import pickle as _pickle

def pickle_py2k_compat(obj, fname, protocol=2):  # change the version number you like
    """Pickle object `obj` to file `fname`."""
    print('Saving %s' % fname)
    with gensim.utils.smart_open(fname, 'wb') as fout: # 'b' for binary, needed on Windows
        _pickle.dump(obj, fout, protocol=protocol)

gensim.utils.pickle = pickle_py2k_compat
```



Discussion:
Makes sense, thanks for reporting @ccwang002 !

It's probably enough to change the protocol default from `-1` to `2`. There are no issues with pickle performance / memory, so whatever `3` or `4` bring cannot be critical for gensim, I prefer the safe compatibility of `2`.

What do you think? Can you craft a PR?

Some comparison of my custom wiki model using different pickle version is shown below. Newer pickle version creates smaller model file, though they take roughly same time to load on my machine. 

``` bash
# pickle ver 4
144M    twwiki_400/model
1.3G    twwiki_400/model.syn0.npy
1.3G    twwiki_400/model.syn1.npy

# pickle ver 2
261M    twwiki_py2k_400/model
1.3G    twwiki_py2k_400/model.syn0.npy
1.3G    twwiki_py2k_400/model.syn1.npy
```

Like you said, a normal user should not be able to notice the change. So I guess changing the default pickle version to `2` should work for most people. 

However I hope user can specify their pickle version manually by `model.save(..., pickle_protocol)`. It should only require the change of their base class `utils.SaveLoad`. I haven't thought thoroughly about the API change, not sure if it breaks anything.

I can make a PR in the next few days :)

Alright, either option is fine by me (change default in `pickle` / allow override in `save`), so whatever works for you.

Looking forward to the PR :)


Resolving patch:
diff --git a/gensim/utils.py b/gensim/utils.py
--- a/gensim/utils.py
+++ b/gensim/utils.py
@@ -299,7 +299,7 @@ def load(cls, fname, mmap=None):
         return obj
 
     def _smart_save(self, fname, separately=None, sep_limit=10 * 1024**2,
-                    ignore=frozenset()):
+                    ignore=frozenset(), pickle_protocol=2):
         """
         Save the object to file (also see `load`).
 
@@ -316,6 +316,9 @@ def _smart_save(self, fname, separately=None, sep_limit=10 * 1024**2,
         handles, caches etc). On subsequent load() these attributes will
         be set to None.
 
+        `pickle_protocol` defaults to 2 so the pickled object can be imported
+        in both Python 2 and 3.
+
         """
         logger.info(
             "saving %s object under %s, separately %s" % (
@@ -376,7 +379,8 @@ def _smart_save(self, fname, separately=None, sep_limit=10 * 1024**2,
                     val.data, val.indptr, val.indices = None, None, None
 
                     try:
-                        pickle(val, subname(attrib)) # store array-less object
+                        # store array-less object
+                        pickle(val, subname(attrib), protocol=pickle_protocol)
                     finally:
                         val.data, val.indptr, val.indices = data, indptr, indices
                 else:
@@ -386,14 +390,14 @@ def _smart_save(self, fname, separately=None, sep_limit=10 * 1024**2,
             self.__dict__['__numpys'] = numpys
             self.__dict__['__scipys'] = scipys
             self.__dict__['__ignoreds'] = ignoreds
-            pickle(self, fname)
+            pickle(self, fname, protocol=pickle_protocol)
         finally:
             # restore the attributes
             for attrib, val in iteritems(tmp):
                 setattr(self, attrib, val)
 
     def save(self, fname_or_handle, separately=None, sep_limit=10 * 1024**2,
-             ignore=frozenset()):
+             ignore=frozenset(), pickle_protocol=2):
         """
         Save the object to file (also see `load`).
 
@@ -415,12 +419,16 @@ def save(self, fname_or_handle, separately=None, sep_limit=10 * 1024**2,
         handles, caches etc). On subsequent load() these attributes will
         be set to None.
 
+        `pickle_protocol` defaults to 2 so the pickled object can be imported
+        in both Python 2 and 3.
+
         """
         try:
-            _pickle.dump(self, fname_or_handle, protocol=_pickle.HIGHEST_PROTOCOL)
+            _pickle.dump(self, fname_or_handle, protocol=pickle_protocol)
             logger.info("saved %s object" % self.__class__.__name__)
         except TypeError:  # `fname_or_handle` does not have write attribute
-            self._smart_save(fname_or_handle, separately, sep_limit, ignore)
+            self._smart_save(fname_or_handle, separately, sep_limit, ignore,
+                             pickle_protocol=pickle_protocol)
 #endclass SaveLoad
 
 


Test patch: