import spaces
import gradio as gr
from model import DecoderTransformer, Tokenizer
from huggingface_hub import hf_hub_download
import torch


vocab_size=33
n_embed=384
context_size=256
n_layer=6
n_head=6
dropout=0.2

device = 'cuda'

model_id = "philipp-zettl/chessPT"

model_path = hf_hub_download(repo_id=model_id, filename="chessPT.pkl")
tokenizer_path = hf_hub_download(repo_id=model_id, filename="tokenizer.json")

model = DecoderTransformer(vocab_size, n_embed, context_size, n_layer, n_head, dropout)
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.to(device)
tokenizer = Tokenizer.from_pretrained(tokenizer_path)

@spaces.GPU
def greet(prompt):
    model_input = torch.tensor(tokenizer.encode(prompt), dtype=torch.long, device=device).view((1, len(prompt)))
    return tokenizer.decode(model.generate(model_input, max_new_tokens=4, context_size=context_size)[0].tolist())


with gr.Blocks() as demo:
    gr.Markdown("""
    Welcome to ChessPT.

    The Chess-Pre-trained-Transformer.

    The rules are simple: provide a PGN string of your current game, the engine will predict the next token!
    """)
    prompt = gr.Text(label="PGN")
    output = gr.Text(label="Next turn", interactive=False)

    submit = gr.Button("Submit")
    submit.click(greet, [prompt], [output])

demo.launch()