|
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() |