Merge branch 'master' of ssh://gogs@the-cake-is-a-lie.net:20022/jonas/NLP-LAB.git

Project/Tools/sklearn_doc2vec.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+from gensim.models import doc2vec
+from collections import namedtuple
+from gensim.utils import to_unicode
+from sklearn.base import BaseEstimator, TransformerMixin
+import numpy as np
+
+"""
+This is a litte helper module providing a doc2vec class
+which can be thrown into a sklearn pipeline. A little bit modified taken from:
+https://github.com/fanta-mnix/sklearn-doc2vec/blob/master/word_embeddings.py
+"""
+
+def documentize(X):
+    docs = []
+    analyzedDocument = namedtuple('AnalyzedDocument', 'words tags')
+    for i, text in enumerate(X):
+        words = text.lower().split()
+        tags = [i]
+        docs.append(analyzedDocument(words, tags))
+    return docs
+
+class Doc2VecTransformer(BaseEstimator, TransformerMixin):
+
+    def __init__(self, size=300, window=8, min_count=5):
+        self.size = size
+        self.window = window
+        self.min_count = min_count
+        self._model = None
+
+    def fit(self, X, y=None):
+        model = doc2vec.Doc2Vec(documentize(X), size=self.size, window=self.window, min_count=self.min_count)
+
+        self._model = model
+        return self
+
+    def fit_transform(self, X, y=None, **fit_params):
+        self.fit(X, y)
+        return self._model.docvecs
+
+    def transform(self, X, copy=True):
+        assert self._model is not None, 'model is not fitted'
+        return np.array([self._model.infer_vector(document.words) for document in documentize(X)])
+