import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tokenizer = AutoTokenizer.from_pretrained("fubuki119/JokesGPT")
model = AutoModelForCausalLM.from_pretrained("fubuki119/JokesGPT")


def generate(max_length):
    starting_text = "JOKE:"
    end_token = "<|endoftext|>"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    model.eval()
    with torch.no_grad():
        cur_ids = torch.tensor(tokenizer.encode("JOKE:")
                               ).unsqueeze(0).to(device)
        for i in range(max_length):
            outputs = model(cur_ids)
            logits, _ = outputs[:]
            softmax_logits = torch.softmax(logits[0, -1], dim=0)

            next_token_id = torch.multinomial(softmax_logits, 1).item()
            cur_ids = torch.cat([cur_ids, torch.ones(
                (1, 1)).long().to(device) * next_token_id], dim=1)
            if next_token_id == tokenizer.encode(end_token)[0]:
                joke = cur_ids.detach().cpu().tolist()
                joke = joke[0]
                return tokenizer.decode(joke[3:-1])
        joke = cur_ids.detach().cpu().tolist()
        joke = joke[0]
        return tokenizer.decode(joke[3:])


def fn(n):
    print(f"i got {n}")
    return "thanks"


iface = gr.Interface(
    fn=generate,
    inputs=gr.Number(value=200, label="Maxlen"),
    outputs="text",
)
iface.launch(share=True, debug=True)