File size: 3,823 Bytes
9df366d 145190f 3e475a6 a4f1b96 84629d8 a4f1b96 7afc5b6 145190f 17b3b4b aeae15e 36e29c8 7e3f9e2 36e29c8 0575679 145190f 203156d a02deee 7afc5b6 203156d 145190f 7ba6d1d 145190f 40b7064 aeae15e |
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 85 86 87 88 89 90 91 |
import torch
import gradio as gr
import torch.nn.functional as F
from transformers import BertTokenizer, GPT2LMHeadModel,PreTrainedTokenizerFast
# tokenizer = BertTokenizer.from_pretrained("supermy/poetry")
tokenizer = PreTrainedTokenizerFast(tokenizer_file="poetry-bpe.json",add_special_token=True,
bos_token="<|endoftext|>",
eos_token="<|endoftext|>",
pad_token="[PAD]",
cls_token="[CLS]",
sep_token="[SEP]",
unk_token="[UNK]",
padding_side="left",
)
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):
# input_ids.extend( tokenizer.encode(input_text + "-", 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__':
# 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() |