w2v training script

w2v.py

"""
训练向量
python w2v.py corpus.txt w2v_char.txt
"""

import logging
import os.path
import sys
import multiprocessing

from gensim.corpora import WikiCorpus
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence

if __name__ == '__main__':
    program = os.path.basename(sys.argv[0])
    logger = logging.getLogger(program)

    logging.basicConfig(format='%(asctime)s: %(levelname)s: %(message)s')
    logging.root.setLevel(level=logging.INFO)
    logger.info("running %s" % ' '.join(sys.argv))

    # check and process input arguments
    if len(sys.argv) != 3:
        print(globals()['__doc__'] % locals())
        sys.exit(1)
    inp, outp = sys.argv[1:3]

    model = Word2Vec(LineSentence(inp), sg=1, size=200, window=5, min_count=5,
                      workers=multiprocessing.cpu_count())
    print("model: ", model)
    # trim unneeded model memory = use(much) less RAM
    # model.init_sims(replace=True)
    # model.save(outp1)
    model.wv.save_word2vec_format(outp, binary=False)