用gensim快速打开词向量

gensim是一个方便的nlp工具,特别是用来导入词向量,这里简单记录一下gensim导入词向量的方法

import gensim

w2v = gensim.models.KeyedVectors.load_word2vec_format("small_embedding.txt")
print(w2v['a'])

gensim导入词向量需要词向量文件的首行是:所有的单词数 词向量的维度
如果不是这个格式,需要做额外的操作

利用下面的代码,可以跨平台的自动添加第一行,生成符合要求的词向量文件,记录下来方便使用

def get_line_nums(filename):
    f = open(filename, 'r')
    count = 0
    for line in f:
        count += 1
    return count


# Linux或者Windows下打开词向量文件,在开始增加一行
def prepend_line(infile, outfile, line):
    with open(infile, 'r') as old:
        with open(outfile, 'w') as new:
            new.write(str(line) + "\n")
            shutil.copyfileobj(old, new)


def prepend_slow(infile, outfile, line):
    with open(infile, 'r') as fin:
        with open(outfile, 'w') as fout:
            fout.write(line + "\n")
            for line in fin:
                fout.write(line)
    

def add_first_line(filename):
    num_lines = get_line_nums(filename)
    gensim_file = 'glove_model.txt'
    gensim_first_line = "{} {}".format(num_lines, 300)
    # Prepends the line.
    if platform == "linux" or platform == "linux2":
        prepend_line(filename, gensim_file, gensim_first_line)
    else:
        prepend_slow(filename, gensim_file, gensim_first_line)

你可能感兴趣的:(用gensim快速打开词向量)