topic_model.py 文件源码

def train_model(model_name, corpus, id2word, num_topics):
    """
    Train specified model
    """
    # LDA
    if model_name == 'lda':
        model = models.LdaModel(
            corpus,
            id2word=id2word,
            num_topics=num_topics,
            alpha='auto',
            eval_every=5,
        )
        return model
    # LSI
    elif model_name == 'lsi':
        model = models.LsiModel(
            corpus,
            id2word=id2word,
            num_topics=num_topics,
        )
        return model
    else:
        print('Invalid model name')
    return None