|
|
|
import torch |
|
|
|
import gradio as gr |
|
import torch.nn.functional as F |
|
|
|
from transformers import GPT2LMHeadModel, CpmTokenizer |
|
|
|
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 = CpmTokenizer(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 |
|
|
|
model = GPT2LMHeadModel.from_pretrained("lewiswu1209/gpt2-chinese-composition") |
|
model.eval() |
|
gr.Interface( |
|
fn=generate, |
|
inputs=[ |
|
"text", |
|
gr.Textbox(lines=7, placeholder="在这里输入一个开头。"), |
|
"number" |
|
], |
|
outputs=gr.Textbox(lines=15, placeholder="这里会输出一段文字。") |
|
).launch() |
|
|