class DocEmbedder:
def __init__(self, embedder='tfidf', vector_size=16, window=2, workers=4):
self.documents = []
self.n_docs = 1
self.vector_size = vector_size
self.window = window
self.workers = workers
self.map = {}
self.embedder = embedder
def add_doc(self, doc):
if 'doc2vec' in self.embedder:
self.documents.append(TaggedDocument(doc, str(self.n_docs)))
else:
self.documents.append(doc)
self.map[doc] = self.n_docs
self.n_docs = self.n_docs+1
def train(self):
if 'doc2vec' in self.embedder:
self.model = Doc2Vec(self.documents, vector_size=self.vector_size, window=self.window, min_count=1, workers=self.workers)
else:
self.model = TfidfVectorizer()
self.model.fit_transform(self.documents)
return self.model
def embed(self, doc):
if 'doc2vec' in self.embedder:
return self.model.infer_vector([doc])
else:
result = self.model.transform([doc])
return np.array(result.toarray()[0])
