import torch

import gradio as gr
import torch.nn.functional as F

from transformers import BertTokenizer, GPT2LMHeadModel

tokenizer = BertTokenizer.from_pretrained("uer/gpt2-chinese-poem")
model = GPT2LMHeadModel.from_pretrained("uer/gpt2-chinese-poem")
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 input_text + "".join( tokenizer.convert_ids_to_tokens(generated) )

if __name__ == "__main__":
    gr.Interface(fn=generate, inputs="text", outputs="text").launch()