fix app

app.py

@@ -1,6 +1,6 @@
 import spaces
 import gradio as gr
-from model import DecoderTransformer
+from model import DecoderTransformer, Tokenizer
 from huggingface_hub import hf_hub_download
 import torch
 
@@ -12,17 +12,36 @@ 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):
-    return model.generate(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 = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()

model.py

@@ -1,3 +1,4 @@
+import json
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
@@ -119,7 +120,11 @@ class DecoderTransformer(nn.Module):
         loss = F.cross_entropy(logits, targets)
         return logits, loss
 
-    def generate(self, idx, max_new_tokens, context_size):
+    def generate(self, idx, max_new_tokens=50, context_size=None):
+        if context_size is None:
+            context_size = int(self.position_embedding_table.weight.shape[0])
+            print(context_size)
+
         for _ in range(max_new_tokens):
             idx_cond = idx[:, -context_size:]
             logits, loss = self(idx_cond)
@@ -129,3 +134,25 @@ class DecoderTransformer(nn.Module):
             idx = torch.cat([idx, idx_next], dim=1)
         return idx
 
+
+class Tokenizer:
+    def __init__(self, vocab):
+        self.vocab = vocab
+        self.stoi = {ch: idx for idx, ch in enumerate(vocab)}
+        self.itos = {idx: ch for idx, ch in enumerate(vocab)}
+
+    def encode(self, s):
+        return [self.stoi[c] for c in s]
+
+    def decode(self, i):
+        return ''.join([self.itos[x] for x in i])
+
+    @classmethod
+    def from_pretrained(cls, path):
+        with open(path, 'r') as f:
+            vocab = json.load(f)
+        return cls(vocab)
+
+    def save_pretrained(self, path):
+        with open(path, 'w') as f:
+            json.dump(self.vocab, f)