File size: 3,212 Bytes
145190f 7afc5b6 145190f 203156d 145190f 7afc5b6 203156d 145190f 7ba6d1d 145190f 40b7064 145190f |
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 |
import torch
import gradio as gr
import torch.nn.functional as F
from transformers import BertTokenizer, GPT2LMHeadModel
tokenizer = BertTokenizer.from_pretrained("supermy/poetry")
model = GPT2LMHeadModel.from_pretrained("supermy/poetry")
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):
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
cur_len = len(input_ids)
input_len = cur_len
last_token_id = input_ids[-1]
input_ids = torch.tensor([input_ids], dtype=torch.long)
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__':
# tokenizer = BertTokenizer(vocab_file="chinese_vocab.model")
# 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="初始文本"),
"number"
],
outputs=gr.Textbox(lines=15, placeholder="AI生成的文本显示在这里。",label="生成的文本")
).launch()
|