word2vec
以下代码用到了 gensim 包和 pandas 包，如果不用 pandas， 也可以使用基本的文本操作来实现文档迭代器。只需要做部分修改，这里不提供修改方式。关于迭代器的知识这里不做介绍。
# coding: utf-8
from gensim.models import Word2Vec
from data_helpers import load_csv
import pandas as pd
# 加载 csv 文件
def load_csv(file_name):
df = pd.read_csv(file_name, sep='\t|\n', encoding='utf8', header=None,
names=['id', 'head', 'content', 'label'], engine='python')
return df
class MySentences(object):
"""
Word2Vec 使用的文档迭代器
"""
def __init__(self, df_list=None):
"""
:param filename_list: 文件名列表
"""
# 定义要使用的文件名
# 其中文件为 csv 文件，没有头名称，基本结构如下：
# TEXT_ID 文章标题 文章内容 标签（POSITIVE/NEGATIVE）
train_words_clean_file = './data/train_words_clean.csv'
test_words_clean_file = './data/test_words_clean.csv'
if not df_list:
self.df_list = [load_csv(train_words_clean_file), load_csv(test_words_clean_file)]
def __iter__(self):
# 遍历所有要用到的文件
for df in self.df_list:
# 遍历每个文件的所有行
for line_num in range(df.shape[0]):
words = []
words.extend(df.iloc[line_num]['head'].split())
words.extend(df.iloc[line_num]['content'].split())
yield words
class W2VModelManager:
"""
Word2Vec 模型管理器
"""
def __init__(self):
# 定义模型名称
self.model_name = './model/sg.w2v' # sg=1
def train_model(self):
"""
训练模型
"""
# 文档迭代器
sens = MySentences()
w2v = Word2Vec(sens, size=200, window=5, sg=1, min_count=3, workers=12, iter=20)
w2v.save(self.model_name)
def load_model(self):
"""
加载模型
"""
return Word2Vec.load(self.model_name)
def model_test(self):
"""
测试模型
"""
# 加载模型
model = self.load_model()
# 打印模型基本信息
print(model)
# 定义词列表
words = ['经济', '中共', '政治', '篮球']
# 注意：如果词列表中的词不在词向量中，那么下列代码会抛出异常，这里没有进行异常处理，仅用来粗略的观察模型训练的结果
# 打印相近的词以及相似度
for word in words:
for w in (model.most_similar(positive=[word], topn=10)):
print(w)
print('-'*60)
# 输出某个词的向量
# print(model['经济'])
# print(model['的'])