import gensim #引入gensim
import os
import re
import sys
import multiprocessing #引入多线程操作
from time import time
class getSentence(object):
#初始化，获取文件路径
    def __init__(self, dirname):
        self.dirname = dirname

#构建一个迭代器
    def __iter__(self):
        for root, dirs, files in os.walk(self.dirname):
            for filename in files:
                file_path = root + '/' + filename
                for line in open(file_path):
                    try:
                        #清除异常数据，主要是去除空白符以及长度为0的内容
                        s_line = line.strip()
                        if s_line== "":
                                continue
                       #把句子拆成词
                        word_line = [word for word in s_line.split( )]
                        yield word_line
                    except Exception:
                        print("catch exception")
                        yield ""
if __name__ == '__main__':
#记录一个起始时间
    begin = time()
#获取句子迭代器
    sentences = getSentence("traindata")
#训练word2vec模型 使用句子迭代器作为语料的输入，设定的最终向量长度为200维；窗口长度为15；词的最小计数为10，词频少于10的词不会进行计算；使用并行处理
    model = gensim.models.Word2Vec(sentences,size=200,window=15,min_count=10, workers=multiprocessing.cpu_count())
#模型存储，这块记得先预先新建一个model路径，或者也可以增加一段代码来识别是否已经创建，如果没有则新建一个路径
    model.save("model/word2vec_gensim")
    model.wv.save_word2vec_format("model/word2vec_org",
                                  "model/vocabulary",
                                  binary=False)
    end = time()
#输出运算所用时间
    print ("Total procesing time: %d seconds" % (end - begin))

import gensim
from sklearn.cluster import KMeans
from sklearn.externals import joblib
from time import time
#加载之前已经训练好的word2vec模型
def load_model():
    model = gensim.models.Word2Vec.load('../word2vec/model/word2vec_gensim')
    return model
#加载城市名称词库
def load_filterword():
    fd = open("mddwords.txt","r")
    filterword=[]
    for line in fd.readlines():
        line=line.strip()
        filterword.append(line)
    return filterword
if __name__=="__main__":
    start = time()
#加载word2vec模型
    model = load_model()
#加载词汇表
    filterword = load_filterword()
#输出词汇表长度
    print(len(filterword))
    wordvector = []
    filterkey={}
#获取我们的城市名称词库的词向量
    for word in filterword:
        wordvector.append(model[word])
        filterkey[word]=model[word]
   #输出词汇数量
    print(len(wordvector))
#训练K-means模型，这里代码设置的聚类数为2000，最大迭代次数为100，n_jobs设置的是有多少个任务同时在跑，这样可以进行多组实验来消除初始化点带来的影响
    clf = KMeans(n_clusters=2000,max_iter=100,n_jobs=10)
    s = clf.fit_predict(wordvector)
#把模型保存下来
    joblib.dump(clf,"kmeans_mdd2000.pkl")
    labels = clf.labels_
    labellist = labels.tolist()
    print(clf.inertia_)
#把所有城市名称的聚类标签保存下来
    fp = open("label_mdd2000",'w')
    fp.write(str(labellist))
    fp.close()
#把所有城市名称保存下来，其中顺序与聚类标签顺序一致
    fp1 = open("keys_mdd2000",'w')
    for key in filterkey:
        fp1.write(key+'\n')
    print("over")
    end = time()
    print("use time")
    print(end-start)