train_word2vec.py 文件源码

def main(in_dir, out_loc, task=1, size=128, window=5, min_count=10,
    n_workers=4, hs=1, nr_iter=5):
    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)
    model = Word2Vec(
        sg=task,
        size=size,
        window=window,
        min_count=min_count,
        workers=n_workers,
        hs=1,
        iter=nr_iter
    )
    corpus = Corpus(in_dir)
    total_words = 0
    total_sents = 0
    for text_no, text_loc in enumerate(iter_dir(corpus.directory)):
        with io.open(text_loc, 'r', encoding='utf8') as file_:
            try:
                text = file_.read()
            except UnicodeDecodeError:
                print(text_loc)
        total_sents += text.count('\n')
        total_words += corpus.count_doc(text.split())  
        logger.info("PROGRESS: at batch #%i, processed %i words, keeping %i word types",
                    text_no, total_words, len(corpus.strings))
    model.corpus_count = total_sents
    model.raw_vocab = defaultdict(int)
    for key, string in corpus.strings.items():
        model.raw_vocab[string] = corpus.counts[key]
    model.scale_vocab()
    model.finalize_vocab()
    model.iter = nr_iter
    model.train(corpus)

    # Trims down model
    model.init_sims(replace=True)

    model.save(out_loc)