import gradio as gr
import torch
from my_gpt import my_gpt
from tokenizer.tokenizer import BPE

##Load model
model = my_gpt.load_pretrained("model/model_1000_.bin")
tokenizer = BPE()



def generate(input_text):
    tokens = tokenizer.encode(input_text)
    gen_ids = model.generate(torch.tensor([tokens]))
    output = tokenizer.decode(gen_ids[0].tolist())
    return output

iface = gr.Interface(fn=generate, 
                     inputs="text", 
                     outputs="text",
                     title="GPT - 1000 steps",
                     description="""This model is trained for 1000 steps only. It is not
                     able to generate perfect sentences/words. However, it has learnt a gist of the English language""")
iface.launch()