Views
No views yet
This is a model to generated Chinese poetry with leading characters and certain tune of mood.
[CLS]之后的诗句。'忍看-窈窕-孤寝-勾带-嫩-黄昏《粉度》『二』[CLS]堞云齐，[CLS]清笳、愁入暮烟林杪。素艳透春，玉骨凄凉，勾带月痕生早。江天苍莽黄昏後，依然是、粉寒香瘦。动追感、西园嫩约，夜深人悄。记得东风窈窕。曾夜踏横斜，醉携娇小。惆怅旧欢，回首俱非，忍看绿笺红豆。香销纸帐人孤寝，相思恨、花还知否。梦回处，霜飞翠楼已晓。'1
2from transformers import (AutoTokenizer, AutoModelForCausalLM)
3tokenizer = AutoTokenizer.from_pretrained('raynardj/keywords-cangtou-chinese-poetry')
4model = AutoModelForCausalLM.from_pretrained('raynardj/keywords-cangtou-chinese-poetry')
5
6def inference(lead, keywords = []):
7 """
8 lead: 藏头的语句， 比如一个人的名字， 2，3 或4个字
9 keywords：关键词, 0~12个关键词比较好
10 """
11 leading = f"《{lead}》"
12 text = "-".join(keywords)+leading
13 input_ids = tokenizer(text, return_tensors='pt', ).input_ids[:,:-1]
14 lead_tok = tokenizer(lead, return_tensors='pt', ).input_ids[0,1:-1]
15
16 with torch.no_grad():
17 pred = model.generate(
18 input_ids,
19 max_length=256,
20 num_beams=5,
21 do_sample=True,
22 repetition_penalty=2.1,
23 top_p=.6,
24 bos_token_id=tokenizer.sep_token_id,
25 pad_token_id=tokenizer.pad_token_id,
26 eos_token_id=tokenizer.sep_token_id,
27 )[0,1:]
28
29 # 我们需要将[CLS] 字符， 也就是101, 逐个换回藏头的字符
30 mask = (pred==101)
31 while mask.sum()<len(lead_tok):
32 lead_tok = lead_tok[:mask.sum()]
33 while mask.sum()>len(lead_tok):
34 reversed_lead_tok = lead_tok.flip(0)
35 lead_tok = torch.cat([
36 lead_tok, reversed_lead_tok[:mask.sum()-len(lead_tok)]])
37 pred[mask] = lead_tok
38 # 从 token 编号解码成语句
39 generate = tokenizer.decode(pred, skip_special_tokens=True)
40 # 清理语句
41 generate = generate.replace("》","》\n").replace("。","。\n").replace(" ","")
42 return generate1>>> inference("上海",["高楼","虹光","灯红酒绿","华厦"])
2高楼-虹光-灯红酒绿-华厦《上海》
3『二』
4上台星月明如昼。
5海阁珠帘卷画堂。
6
7>>> inference("刘先生",["妆容","思","落花","空镜"])
8妆容-思-落花-空镜《刘先生》
9『三』
10刘郎何事不相逢，先把金尊酒未空。
11生意自知人薄命，多情只有月明中。