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")
# model.to(torch.device("cpu"))
# model.save_pretrained("model/model_1000_cpu.bin")
# exit()
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="NoobGPT - 1000 steps",
                     description="""This 13M param model is trained for 1000steps only and has seen only 1M tokens. It is not
                     able to generate perfect sentences/words but has acquired a rudimentary understanding of the English language""")
iface.launch()