File size: 2,056 Bytes
60a408e |
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 |
import torch
import gradio as gr
import torch.nn.functional as F
from transformers import BertTokenizer, GPT2LMHeadModel
tokenizer = BertTokenizer.from_pretrained("supermy/couplet-gpt2")
model = GPT2LMHeadModel.from_pretrained("supermy/couplet-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(input_text):
input_ids = [tokenizer.cls_token_id]
input_ids.extend( tokenizer.encode(input_text + "-", add_special_tokens=False) )
input_ids = torch.tensor( [input_ids] )
generated = []
for _ in range(100):
output = model(input_ids)
next_token_logits = output.logits[0, -1, :]
next_token_logits[ tokenizer.convert_tokens_to_ids('[UNK]') ] = -float('Inf')
filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=8, top_p=1)
next_token = torch.multinomial( F.softmax(filtered_logits, dim=-1), num_samples=1 )
if next_token == tokenizer.sep_token_id:
break
generated.append( next_token.item() )
input_ids = torch.cat( (input_ids, next_token.unsqueeze(0)), dim=1 )
return "".join( tokenizer.convert_ids_to_tokens(generated) )
if __name__ == "__main__":
gr.Interface(
fn=generate,
inputs="text",
outputs="text"
).launch() |