Optimize FastText.load_fasttext_model

gensim/models/_fasttext_bin.py

@@ -79,6 +79,49 @@ def _yield_field_names():
 
 _FIELD_NAMES = sorted(set(_yield_field_names()))
 Model = collections.namedtuple('Model', _FIELD_NAMES)
+"""Holds data loaded from the Facebook binary.
+
+Fields
+------
+dim : int
+    The dimensionality of the vectors.
+ws : int
+    The window size.
+epoch : int
+    The number of training epochs.
+neg : int
+    If non-zero, indicates that the model uses negative sampling.
+loss : int
+    If equal to 1, indicates that the model uses hierarchical sampling.
+model : int
+    If equal to 2, indicates that the model uses skip-grams.
+bucket : int
+    The number of buckets.
+min_count : int
+    The threshold below which the model ignores terms.
+t : float
+    The sample threshold.
+minn : int
+    The minimum ngram length.
+maxn : int
+    The maximum ngram length.
+raw_vocab : collections.OrderedDict
+    A map from words (str) to their frequency (int).  The order in the dict
+    corresponds to the order of the words in the Facebook binary.
+nwords : int
+    The number of words.
+vocab_size : int
+    The size of the vocabulary.
+vectors_ngrams : numpy.array
+    This is a matrix that contains vectors learned by the model.
+    Each row corresponds to a vector.
+    The number of vectors is equal to the number of words plus the number of buckets.
+    The number of columns is equal to the vector dimensionality.
+hidden_output : numpy.array
+    This is a matrix that contains the shallow neural network output.
+    This array has the same dimensions as vectors_ngrams.
+    May be None - in that case, it is impossible to continue training the model.
+"""
 
 
 def _struct_unpack(fin, fmt):
@@ -177,7 +220,7 @@ def _load_matrix(fin, new_format=True):
     return matrix
 
 
-def load(fin, encoding='utf-8'):
+def load(fin, encoding='utf-8', full_model=True):
     """Load a model from a binary stream.
 
     Parameters
@@ -186,6 +229,9 @@ def load(fin, encoding='utf-8'):
         The readable binary stream.
     encoding : str, optional
         The encoding to use for decoding text
+    full_model : boolean, optional
+        If False, skips loading the hidden output matrix.  This saves a fair bit
+        of time, but prevents training continuation.
 
     Returns
     -------
@@ -209,10 +255,12 @@ def load(fin, encoding='utf-8'):
 
     vectors_ngrams = _load_matrix(fin, new_format=new_format)
 
-    hidden_output = _load_matrix(fin, new_format=new_format)
-    model.update(vectors_ngrams=vectors_ngrams, hidden_output=hidden_output)
-
-    assert fin.read() == b'', 'expected to reach EOF'
+    if not full_model:
+        hidden_output = None
+    else:
+        hidden_output = _load_matrix(fin, new_format=new_format)
+        assert fin.read() == b'', 'expected to reach EOF'
 
+    model.update(vectors_ngrams=vectors_ngrams, hidden_output=hidden_output)
     model = {k: v for k, v in model.items() if k in _FIELD_NAMES}
     return Model(**model)

gensim/models/fasttext.py

@@ -704,6 +704,13 @@ def train(self, sentences=None, corpus_file=None, total_examples=None, total_wor
             >>> model.train(sentences, total_examples=model.corpus_count, epochs=model.epochs)
 
         """
+        cant_train = hasattr(self.trainables, 'syn1neg') and self.trainables.syn1neg is None
+        if cant_train:
+            raise ValueError(
+                'this model cannot be trained any further, '
+                'if this is a native model, try loading it with full_model=True'
+            )
+
         super(FastText, self).train(
             sentences=sentences, corpus_file=corpus_file, total_examples=total_examples, total_words=total_words,
             epochs=epochs, start_alpha=start_alpha, end_alpha=end_alpha, word_count=word_count,
@@ -754,7 +761,7 @@ def __contains__(self, word):
         return self.wv.__contains__(word)
 
     @classmethod
-    def load_fasttext_format(cls, model_file, encoding='utf8'):
+    def load_fasttext_format(cls, model_file, encoding='utf8', full_model=True):
         """Load the input-hidden weight matrix from Facebook's native fasttext `.bin` and `.vec` output files.
 
         Notes
@@ -770,14 +777,17 @@ def load_fasttext_format(cls, model_file, encoding='utf8'):
             as Gensim requires only `.bin` file to the load entire fastText model.
         encoding : str, optional
             Specifies the file encoding.
+        full_model : boolean, optional
+            If False, skips loading the hidden output matrix.  This saves a fair bit
+            of time, but prevents training continuation.
 
         Returns
         -------
         :class: `~gensim.models.fasttext.FastText`
             The loaded model.
 
         """
-        return _load_fasttext_format(model_file, encoding=encoding)
+        return _load_fasttext_format(model_file, encoding=encoding, full_model=full_model)
 
     def load_binary_data(self, encoding='utf8'):
         """Load data from a binary file created by Facebook's native FastText.
@@ -959,7 +969,7 @@ def _pad_ones(m, new_shape):
     return vstack([m, suffix])
 
 
-def _load_fasttext_format(model_file, encoding='utf-8'):
+def _load_fasttext_format(model_file, encoding='utf-8', full_model=True):
     """Load the input-hidden weight matrix from Facebook's native fasttext `.bin` and `.vec` output files.
 
     Parameters
@@ -971,6 +981,9 @@ def _load_fasttext_format(model_file, encoding='utf-8'):
         as Gensim requires only `.bin` file to the load entire fastText model.
     encoding : str, optional
         Specifies the file encoding.
+    full_model : boolean, optional
+        If False, skips loading the hidden output matrix.  This saves a fair bit
+        of time, but prevents training continuation.
 
     Returns
     -------
@@ -980,7 +993,7 @@ def _load_fasttext_format(model_file, encoding='utf-8'):
     if not model_file.endswith('.bin'):
         model_file += '.bin'
     with smart_open(model_file, 'rb') as fin:
-        m = gensim.models._fasttext_bin.load(fin, encoding=encoding)
+        m = gensim.models._fasttext_bin.load(fin, encoding=encoding, full_model=full_model)
 
     model = FastText(
         size=m.dim,

gensim/test/test_fasttext.py

@@ -59,6 +59,21 @@ def setUp(self):
         self.test_model = FT_gensim.load_fasttext_format(self.test_model_file)
         self.test_new_model_file = datapath('lee_fasttext_new')
 
+    def test_native_partial_model(self):
+        """Can we skip loading the NN and still get a working model?"""
+        model = FT_gensim.load_fasttext_format(self.test_model_file, full_model=False)
+
+        #
+        # Training continuation should be impossible
+        #
+        self.assertIsNone(model.trainables.syn1neg)
+        self.assertRaises(ValueError, model.train, sentences,
+                          total_examples=model.corpus_count, epochs=model.epochs)
+
+        model.wv['green']
+        model.wv['foobar']
+        model.wv['thisworddoesnotexist']
+
     def test_training(self):
         model = FT_gensim(size=10, min_count=1, hs=1, negative=0, seed=42, workers=1)
         model.build_vocab(sentences)