增加doc2vec

src/doc2vec_data.py

@@ -0,0 +1,60 @@
+import gensim
+import numpy as np
+import jieba
+from gensim.models.doc2vec import Doc2Vec,TaggedDocument
+
+
+def jieba_tokenize(text):
+    """
+    文本分词
+    :param text: 文本
+    :return: 分词list
+    """
+    return jieba.lcut(text)
+
+
+def get_datasest():
+    """
+    获取doc2vec文本训练数据集
+    :return: 文本分词list，及id
+    """
+    x_train = []
+    for file in open('toutiao_cat_data.txt', encoding='utf8'):
+        file = file.split('_!_')
+        if len(file) > 3:
+            document = TaggededDocument(file[3], tags=[int(file[1])])
+            x_train.append(document)
+    return x_train
+
+
+def train(x_train, size=2000, epoch_num=10):
+    model_dm = Doc2Vec(x_train, min_count=1, window=3, size=size, sample=1e-3, negative=5, workers=4)
+    model_dm.train(x_train, total_examples=model_dm.corpus_count, epochs=epoch_num)
+    model_dm.save('model')
+    return model_dm
+
+
+def getVecs(model, corpus, size):
+    vecs = [np.array(model.docvecs[z.tags[0]].reshape(1, size)) for z in corpus]
+    return np.concatenate(vecs)
+
+
+def test():
+    model_dm = Doc2Vec.load("model")
+    test_text = ['想换个', '30', '万左右', '的', '车', '，', '现在', '开科鲁兹', '，', '有', '什么', '好', '推荐', '的', '？']
+    inferred_vector_dm = model_dm.infer_vector(test_text)
+    sims = model_dm.docvecs.most_similar([inferred_vector_dm], topn=10)
+    return sims
+
+
+if __name__ == '__main__':
+    x_train = get_datasest()
+    model_dm = train(x_train)
+
+    sims = test()
+    for count, sim in sims:
+        sentence = x_train[count]
+        words = ''
+        for word in sentence[0]:
+            words = words + word + ' '
+        print(words, sim, len(sentence[0]))