Fix the train method of TranslationMatrix

gensim/models/translation_matrix.py

@@ -365,23 +365,23 @@ class BackMappingTranslationMatrix(utils.SaveLoad):
     >>> src_model = Doc2Vec.load(datapath("small_tag_doc_5_iter50"))
     >>> dst_model = Doc2Vec.load(datapath("large_tag_doc_10_iter50"))
     >>>
-    >>> model_trans = BackMappingTranslationMatrix(data, src_model, dst_model)
+    >>> model_trans = BackMappingTranslationMatrix(src_model, dst_model)
     >>> trans_matrix = model_trans.train(data)
     >>>
     >>> result = model_trans.infer_vector(dst_model.docvecs[data[3].tags])
 
     """
-    def __init__(self, tagged_docs, source_lang_vec, target_lang_vec, random_state=None):
+    def __init__(self, source_lang_vec, target_lang_vec, tagged_docs=None, random_state=None):
         """
 
         Parameters
         ----------
-        tagged_docs : list of :class:`~gensim.models.doc2vec.TaggedDocument`, optional
-            Documents that will be used for training
         source_lang_vec : :class:`~gensim.models.doc2vec.Doc2Vec`
             Source Doc2Vec model.
         target_lang_vec : :class:`~gensim.models.doc2vec.Doc2Vec`
             Target Doc2Vec model.
+        tagged_docs : list of :class:`~gensim.models.doc2vec.TaggedDocument`, optional
+            Documents that will be used for training
         random_state : {None, int, array_like}, optional
             Seed for random state.
 
@@ -393,22 +393,24 @@ def __init__(self, tagged_docs, source_lang_vec, target_lang_vec, random_state=N
         self.random_state = utils.get_random_state(random_state)
         self.translation_matrix = None
 
+        if tagged_docs is not None:
+            self.train(tagged_docs)
+
     def train(self, tagged_docs):
         """Build the translation matrix that mapping from the source model's vector to target model's vector
 
         Parameters
         ----------
-        tagged_docs : list of :class:`~gensim.models.doc2vec.TaggedDocument`, optional
-            THIS ARGUMENT WILL BE IGNORED.
+        tagged_docs : list of :class:`~gensim.models.doc2vec.TaggedDocument`
 
         Returns
         -------
         numpy.ndarray
             Translation matrix that mapping from the source model's vector to target model's vector.
 
         """
-        m1 = [self.source_lang_vec.docvecs[item.tags].flatten() for item in self.tagged_docs]
-        m2 = [self.target_lang_vec.docvecs[item.tags].flatten() for item in self.tagged_docs]
+        m1 = [self.source_lang_vec.docvecs[item.tags].flatten() for item in tagged_docs]
+        m2 = [self.target_lang_vec.docvecs[item.tags].flatten() for item in tagged_docs]
 
         self.translation_matrix = np.linalg.lstsq(m2, m1, -1)[0]
         return self.translation_matrix

gensim/test/test_translation_matrix.py

@@ -99,14 +99,14 @@ def setUp(self):
 
     def test_translation_matrix(self):
         model = translation_matrix.BackMappingTranslationMatrix(
-            self.train_docs[:5], self.source_doc_vec, self.target_doc_vec
+            self.source_doc_vec, self.target_doc_vec, self.train_docs[:5]
         )
         transmat = model.train(self.train_docs[:5])
         self.assertEqual(transmat.shape, (100, 100))
 
     def test_infer_vector(self):
         model = translation_matrix.BackMappingTranslationMatrix(
-            self.train_docs[:5], self.source_doc_vec, self.target_doc_vec
+            self.source_doc_vec, self.target_doc_vec, self.train_docs[:5]
         )
         model.train(self.train_docs[:5])
         infered_vec = model.infer_vector(self.target_doc_vec.docvecs[self.train_docs[5].tags])