from gensim.models import Doc2Vec
import pickle

with open('./docs_proc.pkl', 'rb') as file:
    docs = pickle.load(file)

model = Doc2Vec(docs, vector_size = 500, window = 9, min_count = 20, workers=8, dm=0)
with open('./d2v_model.pkl', 'wb') as file:
    pickle.dump(model, file)

epochs = 5
start_alpha = 0.025
end_alpha = 0.0001
alpha_change = (start_alpha-end_alpha)/epochs
assert gensim.models.doc2vec.FAST_VERSION > -1
for i in range(epochs):
    print('Epoch', i)
    model.train(docs, total_examples=model.corpus_count, epochs=1)
    model.alpha -= alpha_change
    model.min_alpha = model.alpha

with open('./d2v_model_trained.pkl', 'wb') as file:
    pickle.dump(model, file)