Merge branch 'master' of ssh://gogs@the-cake-is-a-lie.net:20022/jonas/NLP-LAB.git
This commit is contained in:
commit
9f596781ad
44
Project/Tools/sklearn_doc2vec.py
Normal file
44
Project/Tools/sklearn_doc2vec.py
Normal file
@ -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)])
|
||||||
|
|
Loading…
Reference in New Issue
Block a user