Views
No views yet
1import torch
2import torch.nn.functional as F
3from transformers import GPT2LMHeadModel, CpmTokenizer
4
5def generate_next_token(input_ids):
6input_ids = input_ids[:, -context_len:]
7outputs = model(input_ids=input_ids)
8logits = outputs.logits
9next_token_logits = logits[0, -1, :]
10next_token_logits = next_token_logits / temperature
11next_token_logits[unk_id] = -float('Inf')
12filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=topk, top_p=topp)
13next_token_id = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)
14return next_token_id
15
16def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
17assert logits.dim() == 1
18top_k = min(top_k, logits.size(-1))
19if top_k > 0:
20 indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
21 logits[indices_to_remove] = filter_value
22
23if top_p > 0.0:
24 sorted_logits, sorted_indices = torch.sort(logits, descending=True)
25 cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
26 sorted_indices_to_remove = cumulative_probs > top_p
27 sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
28 sorted_indices_to_remove[..., 0] = 0
29 indices_to_remove = sorted_indices[sorted_indices_to_remove]
30 logits[indices_to_remove] = filter_value
31return logits
32
33def set_logger(log_path):
34logger = logging.getLogger(__name__)
35logger.setLevel(logging.INFO)
36formatter = logging.Formatter(
37 '%(asctime)s - %(levelname)s - %(message)s')
38file_handler = logging.FileHandler(
39 filename=log_path)
40file_handler.setFormatter(formatter)
41file_handler.setLevel(logging.INFO)
42logger.addHandler(file_handler)
43console = logging.StreamHandler()
44console.setLevel(logging.DEBUG)
45console.setFormatter(formatter)
46logger.addHandler(console)
47return logger
48
49# 参数
50device = 'cuda' if torch.cuda.is_available() else 'cpu'
51temperature = 1
52topk = 0
53topp = 0.85
54context_len = 800
55model_path = "model"
56model_vocab = 'model/chinese_vocab.model'
57# 加载tokenizer
58tokenizer = CpmTokenizer(vocab_file=model_vocab)
59eod_id = tokenizer.convert_tokens_to_ids("<eod>") # 文档结束符
60sep_id = tokenizer.sep_token_id
61unk_id = tokenizer.unk_token_id
62
63# 加载模型
64model = GPT2LMHeadModel.from_pretrained(model_path)
65model.eval()
66model = model.to(device)
67
68def generate_zuowen(title, context, max_len):
69title_ids = tokenizer.encode(title, add_special_tokens=False)
70context_ids = tokenizer.encode(context, add_special_tokens=False)
71input_ids = title_ids + [sep_id] + context_ids
72cur_len = len(input_ids)
73last_token_id = input_ids[-1]
74input_ids = torch.tensor([input_ids], dtype=torch.long, device=device)
75
76while True:
77 next_token_id = generate_next_token(input_ids)
78 input_ids = torch.cat((input_ids, next_token_id.unsqueeze(0)), dim=1)
79 cur_len += 1
80 word = tokenizer.convert_ids_to_tokens(next_token_id.item())
81 if cur_len >= max_len and last_token_id == 8 and next_token_id == 3:
82 break
83 if cur_len >= max_len and word in [".", "。", "!", "!", "?", "?", ",", ","]:
84 break
85 if next_token_id == eod_id:
86 break
87 word_mok = word.replace('▁', '').replace('▃', '\n ')
88 print(word_mok, end='', flush=True)
89result = tokenizer.decode(input_ids.squeeze(0))
90content = result.split("<sep>")[1]
91return content
92
93if __name__ == '__main__':
94title='论天下大势'#标题
95context='论天下大势'#上文,可以为空
96max_len=100#续写的字数
97max_len += len(context)
98print(context,end='')
99content = generate_zuowen(title, context, max_len)pip install ransformers==4.6.0
pip install torch==1.7.0
pip install jieja==0.42.1pip install -i https://pypi.tuna.tsinghua.edu.cn/simple transformers==4.6.0
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch==1.7.1
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple jieja==0.42.11.
论天下大势,全凭道教的范蠡;论市井,他才是楚汉王公的治国大业;论乱世风云,他才是江山大太守李商隐的传世佳句。
我想说的是天下大事,在数不清的人眼里,天下大事是不可能的。自古以来,各人为了建功立业,便开始做生意
2.
天空的上面有很多千姿百态的云。
一朵像白色的棉花糖,一朵像蓝色的海豚。它们有的像金色的骏马,有的像蓬松的大扫帚,还有的像犀牛的爪子。
一朵云上有几只鸟,它们从天空中飞过,有的小鸟落在电线杆上,在电线杆上休息,有的小鸟落在树枝上
3.
谁说其神乎?亦有矣斯。因曰:君子勇而不惧,其必是一兵一卒,他守死,其必死矣!于表超薄于一时兴亡,吾乃,乃降矣!汝此暴矣矣!乃双肋相助?此地皇与汝同用肾争渡,汉余为吾兄吾交礼拜而葛申葛汝河。
如是曰:公亦有约约。吾有约亡
4.
吕布,天下大姓姬仲谋,是三国时期魏国名将,是三国鼎立之中的十虎上将之一。身中宽盖九鼎,刀刃上书,为天下守诚,所以他不是自刎乌江,死而后已的谋士。