python基於詞向量的古詩生成器

python基於詞向量的古詩生成器

from gensim.models import Word2Vec  # 詞向量
from random import choice
from os.path import exists
import warnings

warnings.filterwarnings('ignore')  # 不打印警告

class CONF:
    path = '古詩詞.txt'
    window = 16          # 滑窗大小
    min_count = 60       # 過濾低頻字
    size = 125           # 詞向量維度
    topn = 14            # 生成詩詞的開放度
    model_path = 'word2vec'

class Model:
    def __init__(self, window, topn, model):
        self.window = window
        self.topn = topn
        self.model = model                   # 詞向量模型
        self.chr_dict = model.wv.index2word  # 字典

    """模型初始化"""
    @classmethod
    def initialize(cls, config):
        if exists(config.model_path):
            # 模型讀取
            model = Word2Vec.load(config.model_path)
        else:
            # 語料讀取
            with open(config.path, encoding='utf-8') as f:
                ls_of_ls_of_c = [list(line.strip()) for line in f]
            # 模型訓練和保存
            model = Word2Vec(ls_of_ls_of_c, config.size,
                             window=config.window, min_count=config.min_count)
            model.save(config.model_path)
        return cls(config.window, config.topn, model)

    """古詩詞生成"""
    def poem_generator(self, title, form):
        filter = lambda lst: [t[0] for t in lst if t[0] not in ['，', '。']]
        # 標題補全
        if len(title) < 4:
            if not title:
                title += choice(self.chr_dict)
            for _ in range(4 - len(title)):
                similar_chr = self.model.similar_by_word(title[-1], self.topn // 2)
                similar_chr = filter(similar_chr)
                char = choice([c for c in similar_chr if c not in title])
                title += char
        # 文本生成
        poem = list(title)
        for i in range(form[0]):
            for _ in range(form[1]):
                predict_chr = self.model.predict_output_word(
                    poem[-self.window:], max(self.topn, len(poem) + 1))
                predict_chr = filter(predict_chr)
                char = choice([c for c in predict_chr if c not in poem[len(title):]])
                poem.append(char)
            poem.append('，' if i % 2 == 0 else '。')
        length = form[0] * (form[1] + 1)
        return '《%s》' % ''.join(poem[:-length]) + '\n' + ''.join(poem[-length:])

def main(config=CONF):
    form = {'五言絕句': (4, 5), '七言絕句': (4, 7), '對聯': (2, 9)}
    m = Model.initialize(config)
    while True:
        title = input('輸入標題：').strip()
        try:
            poem = m.poem_generator(title, form['五言絕句'])
            print('\033[031m%s\033[0m' % poem)             # red
            poem = m.poem_generator(title, form['七言絕句'])
            print('\033[033m%s\033[0m' % poem)             # yellow
            poem = m.poem_generator(title, form['對聯'])
            print('\033[036m%s\033[0m' % poem)             # purple
            print()
        except:
            pass

if __name__ == '__main__':
    main()

PS：源碼、word2vec和古詩詞.txt文件將在Github上提供下載。

Github地址：https://github.com/zhangziqi1999/python-poems-writer

2018年5月1日