Views
No views yet
1from transformers import BertTokenizer, BertForMaskedLM
2import torch
3
4# 加载分词器
5tokenizer = BertTokenizer.from_pretrained("qixun/bert-chinese-poem")
6
7# 加载模型
8model = BertForMaskedLM.from_pretrained("qixun/bert-chinese-poem")
9
10# 输入文本
11text = "宵凉百念集孤[MASK],暗雨鸣廊睡未能。生计坐怜秋一叶,归程冥想浪千层。寒心国事浑难料,堆眼官资信可憎。此去梦中应不忘,顺承门内近觚棱。"
12
13# 分词
14inputs = tokenizer(text, return_tensors="pt")
15
16# 模型推理
17with torch.no_grad():
18 outputs = model(**inputs)
19
20# 获取[MASK]标记的位置
21mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
22
23# 获取预测的token_id
24predicted_token_id = outputs.logits[0, mask_token_index].argmax(axis=-1).item()
25
26# 获取预测的词
27predicted_token = tokenizer.decode([predicted_token_id])
28
29print(f"预测的词是:{predicted_token}")
30