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