chinese poem 中文 写诗 唐诗 宋词

一个好玩的中文AI写诗模型V2

架构

数据来源

语言支持

训练

运行代码示例

# 安装以下2个包方便文字处理和模型生成
# !pip install -q simplet5
# !pip install -q chinese-converter

# 具体代码
import torch
from simplet5 import SimpleT5
from transformers import T5Tokenizer, T5ForConditionalGeneration
import chinese_converter

MODELS = {
    # id: (hf_path, max_len)
    '2022-v1': ("hululuzhu/chinese-poem-t5-mengzi-finetune", 64),
    '2023-v2': ("hululuzhu/chinese-poem-t5-v2", 32)
}
MODEL_VERSION = '2023-v2'  # @param ["2023-v2", "2022-v1"]
# Huggingface model card
MODEL_PATH = MODELS[MODEL_VERSION][0]
class PoemModel(SimpleT5):
  def __init__(self) -> None:
    super().__init__()
    self.device = torch.device("cuda")

  def load_my_model(self):
    self.tokenizer = T5Tokenizer.from_pretrained(MODEL_PATH)
    self.model = T5ForConditionalGeneration.from_pretrained(MODEL_PATH)


AUTHOR_PROMPT = "模仿:"
TITLE_PROMPT = "作诗:"
EOS_TOKEN = '</s>'

poem_model = PoemModel()
poem_model.load_my_model()
poem_model.model = poem_model.model.to('cuda')

MAX_AUTHOR_CHAR = 4
MAX_TITLE_CHAR = 12
MIN_CONTENT_CHAR = 10
MAX_CONTENT_CHAR = MODELS[MODEL_VERSION][1]

def poem(title_str, opt_author=None, model=poem_model,
         is_input_traditional_chinese=False,
         num_beams=2):
  model.model = model.model.to('cuda')
  if opt_author:
    in_request = TITLE_PROMPT + title_str[:MAX_TITLE_CHAR] + EOS_TOKEN + AUTHOR_PROMPT + opt_author[:MAX_AUTHOR_CHAR]
  else:
    in_request = TITLE_PROMPT + title_str[:MAX_TITLE_CHAR]
  if is_input_traditional_chinese:
    in_request = chinese_converter.to_simplified(in_request)
  out = model.predict(in_request,
                      max_length=MAX_CONTENT_CHAR,
                      num_beams=num_beams)[0].replace(",", ",")
  if is_input_traditional_chinese:
    out = chinese_converter.to_traditional(out)
    print(f"標題: {in_request.replace('</s>', ' ')}\n詩歌: {out}")
  else:
    print(f"标题: {in_request.replace('</s>', ' ')}\n诗歌: {out}")

简体中文示例

for title in ['秋思', "百花", '佳人有约']:
  # Empty author means general style
  for author in ['', "杜甫", "李白", "李清照", "苏轼"]:
    poem(title, author)
  print()
标题: 作诗:秋思
诗歌: 秋风吹我衣,落叶满庭除。老去心更苦,愁来鬓已疎。
标题: 作诗:秋思 模仿:杜甫
诗歌: 秋风吹我衣,落叶满庭除。客子思乡泪,故人伤远书。
标题: 作诗:秋思 模仿:李白
诗歌: 秋风吹我衣,飒飒满庭树。忆得故园花,今朝已零落。
标题: 作诗:秋思 模仿:李清照
诗歌: 秋风吹我衣,落叶满庭除。天高鸿雁少,日短萤火疎。
标题: 作诗:秋思 模仿:苏轼
诗歌: 秋风吹我衣,飒飒吹我衣。出门无所诣,但觉天宇低。

标题: 作诗:百花
诗歌: 百花头上开,春色爲谁来。欲识春风面,先教花上开。
标题: 作诗:百花 模仿:杜甫
诗歌: 百花开尽见春归,红紫纷纷照眼稀。莫道花时无赏处,且留樽酒对芳菲。
标题: 作诗:百花 模仿:李白
诗歌: 百花开尽见春归,谁把芳菲比玉池。若使东君无别意,春风应解惜花枝。
标题: 作诗:百花 模仿:李清照
诗歌: 百花头上开,春色爲谁来。欲识春风面,先教桃李开。
标题: 作诗:百花 模仿:苏轼
诗歌: 百花头上开,百草头边出。春风吹不断,尽逐东风去。

标题: 作诗:佳人有约
诗歌: 佳人有约在烟汀,相约花前共醉醒。莫道人间春色晚,隔帘应是笑谈声。
标题: 作诗:佳人有约 模仿:杜甫
诗歌: 佳人有约在江干,万里相随入夢寒。玉笛夜吹明月下,金杯春泛水晶盘。
标题: 作诗:佳人有约 模仿:李白
诗歌: 佳人有约在烟汀,玉颜金面映红英。天边月下吹笙处,疑是瑶池旧主人。
标题: 作诗:佳人有约 模仿:李清照
诗歌: 佳人有约在烟汀,相约花前共醉醒。莫道春来无约到,隔帘应是月中听。
标题: 作诗:佳人有约 模仿:苏轼
诗歌: 佳人有约在烟汀,玉佩金鱼照碧浔。应是仙家好风景,夜来花下弄潺湲。

# Try different beams
for title in ['冬雪']:
  for author in  ['', "杜甫"]:
    for num_beams in (2, 3, 5, 10, 20, 50, 100, 200):    
      print(f"num beams: {num_beams}")
      poem(title, author, num_beams=num_beams)
    print("-"*80)
num beams: 2
标题: 作诗:冬雪
诗歌: 冬雪未全消,春寒犹未回。山空云气重,天阔水光开。
num beams: 3
标题: 作诗:冬雪
诗歌: 冬雪未成雪,春寒犹未回。山空云气重,江阔水声来。
num beams: 5
标题: 作诗:冬雪
诗歌: 冬雪未全消,春寒犹未回。山空云气重,江阔水声来。
num beams: 10
标题: 作诗:冬雪
诗歌: 冬雪未成雪,北风先着人。寒威欺病骨,老色逼衰身。
num beams: 20
标题: 作诗:冬雪
诗歌: 冬雪未成雪,北风先作威。山高云气重,江阔水声微。
num beams: 50
标题: 作诗:冬雪
诗歌: 冻云凝不散,寒日淡无光。夜半风号屋,朝来雪满堂。
num beams: 100
标题: 作诗:冬雪
诗歌: 朔风吹雪满山城,万壑千岩冻不鸣。夜半忽闻檐溜响,晓来还见瓦沟平。
num beams: 200
标题: 作诗:冬雪
诗歌: 去年冬雪未全消,今岁春冰犹未消。山色不随人意改,江声长送雁声遥。
--------------------------------------------------------------------------------
num beams: 2
标题: 作诗:冬雪 模仿:杜甫
诗歌: 冬雪未全消,春寒犹未销。山城迷远道,江路入重霄。
num beams: 3
标题: 作诗:冬雪 模仿:杜甫
诗歌: 冬雪未全落,春寒犹未回。江城风日好,山寺雨声来。
num beams: 5
标题: 作诗:冬雪 模仿:杜甫
诗歌: 冬雪未成雪,春寒犹着人。江天无定色,风日有微尘。
num beams: 10
标题: 作诗:冬雪 模仿:杜甫
诗歌: 朔风吹雪满江城,客子衣裘不自温。夜半忽闻檐溜响,晓来还见瓦沟浑。
num beams: 20
标题: 作诗:冬雪 模仿:杜甫
诗歌: 朔风吹雪满江城,万木号风急霰声。老去不知身是客,乱来唯觉鬓成丝。
num beams: 50
标题: 作诗:冬雪 模仿:杜甫
诗歌: 朔风吹雪满江城,万壑千岩冻未平。夜半忽惊飞霰急,晓来还作打窗声。
num beams: 100
标题: 作诗:冬雪 模仿:杜甫
诗歌: 朔风吹雪满江城,万壑千岩冻未平。夜半忽惊飞霰急,晓来还作打窗声。
num beams: 200
标题: 作诗:冬雪 模仿:杜甫
诗歌: 朔风吹雪满江城,万壑千岩冻未平。夜半忽惊飞霰急,晓来还作打窗声。

繁体中文

for title in ['春節', "中秋", "春秋战国"]:
  # Empty author means general style
  for author in ['', "杜甫", "李白", "李清照", "蘇軾"]:
    poem(title, author, is_input_traditional_chinese=True)
  print()

標題: 作诗:春节
詩歌: 節物今朝是,年光此際同。年華驚歲換,身世逐時窮。
標題: 作诗:春节 模仿:杜甫
詩歌: 節物今朝是,春光此夜同。江城聞鼓角,野寺見燒紅。
標題: 作诗:春节 模仿:李白
詩歌: 節物今朝是,春光此夜同。柳條初弄色,梅蕊未藏紅。
標題: 作诗:春节 模仿:李清照
詩歌: 節物今朝是,年光此際同。柳條新染綠,梅蕊未藏紅。
標題: 作诗:春节 模仿:苏轼
詩歌: 節物今朝是,春光此夜同。老來多感事,老去少知功。

標題: 作诗:中秋
詩歌: 秋色今宵半,江天此夜深。雲收山吐月,風送水浮金。
標題: 作诗:中秋 模仿:杜甫
詩歌: 秋色今宵盡,江天萬里長。雲收山氣白,風送水聲涼。
標題: 作诗:中秋 模仿:李白
詩歌: 月色秋來好,人言此夜奇。桂華清似水,桂魄冷於泥。
標題: 作诗:中秋 模仿:李清照
詩歌: 秋色今宵半,清光此夜分。月從天上出,人向世間聞。
標題: 作诗:中秋 模仿:苏轼
詩歌: 中秋月色好,況復是中秋。露重珠猶溼,風高葉未收。

標題: 作诗:春秋战国
詩歌: 國破人亡國亦亡,君王何事獨稱王。當時若使無張許,誰信賢哉是魯王。
標題: 作诗:春秋战国 模仿:杜甫
詩歌: 國破人亡國亦亡,君王何事獨稱王。當時若使無張許,天下安知有範滂。
標題: 作诗:春秋战国 模仿:李白
詩歌: 吳越爭雄勢已分,君王何事更相君。若教國士輕天下,肯信賢人有異聞。
標題: 作诗:春秋战国 模仿:李清照
詩歌: 東門西去是通津,誰信君王不識真。若使魯儒輕國士,肯教吳客作諸侯。
標題: 作诗:春秋战国 模仿:苏轼
詩歌: 天下兵戈尚未休,豈知今日是良謀。君王若問當時事,不道今朝有許愁。