Views
No views yet
对联：您的上联，比如 对联：北国风光，千里冰封，万里雪飘对联：前缀，长度限制32字符is_input_traditional_chinese=True1# 安装以下2个包方便文字处理和模型生成
2# !pip install -q simplet5
3# !pip install -q chinese-converter
4
5# 具体代码
6import torch
7from simplet5 import SimpleT5
8from transformers import T5Tokenizer, T5ForConditionalGeneration
9import chinese_converter
10
11MODEL_PATH = "hululuzhu/chinese-couplet-t5-mengzi-finetune"
12class PoemModel(SimpleT5):
13 def __init__(self) -> None:
14 super().__init__()
15 self.device = torch.device("cuda")
16
17 def load_my_model(self):
18 self.tokenizer = T5Tokenizer.from_pretrained(MODEL_PATH)
19 self.model = T5ForConditionalGeneration.from_pretrained(MODEL_PATH)
20
21COUPLET_PROMPOT = '对联：'
22MAX_SEQ_LEN = 32
23MAX_OUT_TOKENS = MAX_SEQ_LEN
24
25def couplet(in_str, model=couplet_model,
26 is_input_traditional_chinese=False,
27 num_beams=2):
28 model.model = model.model.to('cuda')
29 in_request = f"{COUPLET_PROMPOT}{in_str[:MAX_SEQ_LEN]}"
30 if is_input_traditional_chinese:
31 # model only knows s chinese
32 in_request = chinese_converter.to_simplified(in_request)
33 # Note default sampling is turned off for consistent result
34 out = model.predict(in_request,
35 max_length=MAX_OUT_TOKENS,
36 num_beams=num_beams)[0].replace(",", "，")
37 if is_input_traditional_chinese:
38 out = chinese_converter.to_traditional(out)
39 print(f"上： {in_str}\n下： {out}")1for pre in ['欢天喜地度佳节',
2 '不待鸣钟已汗颜，重来试手竟何艰',
3 '当年欲跃龙门去，今日真披马革还',
4 '北国风光，千里冰封，万里雪飘',
5 '寂寞寒窗空守寡',
6 '烟锁池塘柳',
7 '五科五状元，金木水火土',
8 '望江楼，望江流，望江楼上望江流，江楼千古，江流千古']:
9 couplet(pre)
10
11上： 欢天喜地度佳节
12下： 笑语欢歌迎新春
13上： 不待鸣钟已汗颜，重来试手竟何艰
14下： 何堪击鼓频催泪?一别伤心更枉然
15上： 当年欲跃龙门去，今日真披马革还
16下： 此日当登虎榜来，他年又见龙图新
17上： 北国风光，千里冰封，万里雪飘
18下： 南疆气象，五湖浪涌，三江潮来
19上： 寂寞寒窗空守寡
20下： 逍遥野渡醉吟诗
21上： 烟锁池塘柳
22下： 云封岭上松
23上： 五科五状元，金木水火土
24下： 三才三进士，诗书礼乐诗
25上： 望江楼，望江流，望江楼上望江流，江楼千古，江流千古
26下： 听雨阁，听雨落，听雨阁中听雨落，雨阁万重，雨落万重1for pre in ['飛龍在天', '臺北風光好']:
2 couplet(pre, is_input_traditional_chinese=True, num_beams=10)
3
4上： 飛龍在天
5下： 飛鳳於天
6上： 臺北風光好
7下： 神州氣象新