File size: 3,630 Bytes
5bb0a8b cd80a36 990c63d cd80a36 5bb0a8b aac4628 d93236c aac4628 5bb0a8b 7bddaba 5bb0a8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import torch
import gradio as gr
import torch.nn.functional as F
from transformers import AutoTokenizer, GPT2LMHeadModel
tokenizer = AutoTokenizer.from_pretrained("supermy/jinyong-gpt2")
tokenizer.add_special_tokens({'sep_token': '[SEP]'})
tokenizer.add_special_tokens({'unk_token': '[UNK]'})
model = GPT2LMHeadModel.from_pretrained("supermy/jinyong-gpt2")
model.eval()
def top_k_top_p_filtering( logits, top_k=0, top_p=0.0, filter_value=-float('Inf') ):
assert logits.dim() == 1
top_k = min( top_k, logits.size(-1) )
if top_k > 0:
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
logits[indices_to_remove] = filter_value
if top_p > 0.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum( F.softmax(sorted_logits, dim=-1), dim=-1 )
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = sorted_indices[sorted_indices_to_remove]
logits[indices_to_remove] = filter_value
return logits
def generate(title, context, max_len):
# input_ids=tokenizer.encode(title + "-" + context, add_special_tokens=False)
title_ids = tokenizer.encode(title, add_special_tokens=False)
context_ids = tokenizer.encode(context, add_special_tokens=False)
input_ids = title_ids + [sep_id] + context_ids
print(input_ids)
cur_len = len(input_ids)
input_len = cur_len
last_token_id = input_ids[-1]
input_ids = torch.tensor([input_ids], dtype=torch.long)
# input_ids = [tokenizer.cls_token_id]
# input_ids.extend( tokenizer.encode(title + "-" +context, add_special_tokens=False) )
# input_ids = torch.tensor( [input_ids] )
print(input_ids)
while True:
outputs = model( input_ids=input_ids[:, -200:] )
logits = outputs.logits
next_token_logits = logits[0, -1, :]
next_token_logits = next_token_logits / 1
next_token_logits[unk_id] = -float('Inf')
filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=0, top_p=0.85)
next_token_id = torch.multinomial( F.softmax(filtered_logits, dim=-1), num_samples=1 )
input_ids = torch.cat( ( input_ids, next_token_id.unsqueeze(0) ), dim=1 )
cur_len += 1
word = tokenizer.convert_ids_to_tokens( next_token_id.item() )
if cur_len >= ( input_len + max_len ) and last_token_id == 8 and next_token_id == 3:
break
if cur_len >= ( input_len + max_len ) and word in [".", "。", "！", "!", "?", "？", ",", "，"]:
break
if next_token_id == eod_id:
break
result = tokenizer.decode( input_ids.squeeze(0) )
return result
if __name__ == '__main__':
eod_id = tokenizer.convert_tokens_to_ids("<eod>")
sep_id = tokenizer.sep_token_id
unk_id = tokenizer.unk_token_id
gr.Interface(
fn=generate,
inputs=[
gr.Textbox(lines=1, placeholder="输入文本标题：射雕英雄传", value="射雕英雄传",label="文本标题"),
gr.Textbox(lines=7, placeholder="输入文本内容：郭靖练功偶遇高人,见道长那轻功,转头便拜其为师", value="郭靖练功偶遇高人,见道长那轻功,转头便拜其为师。",label="初始文本"),
gr.Number( value=108,label="生成文本字数")
],
outputs=gr.Textbox(lines=15, placeholder="AI生成的文本显示在这里。",label="生成的文本")
).launch() |