Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 from model import DecoderTransformer, Tokenizer
 from huggingface_hub import hf_hub_download
 import torch
+import chess
+import chess.svg
 
 
 vocab_size=33
@@ -25,7 +27,10 @@ tokenizer = Tokenizer.from_pretrained(tokenizer_path)
 
 def 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())
+    pgn = tokenizer.decode(model.generate(model_input, max_new_tokens=4, context_size=context_size)[0].tolist())
+    game = chess.pgn.read_game(pgn)
+    img = chess.svg.board(game.board())
+    return pgn, img
 
 
 with gr.Blocks() as demo:
@@ -42,14 +47,14 @@ with gr.Blocks() as demo:
 
     submit = gr.Button("Submit")
     submit.click(generate, [prompt], [output])
-
+    img = gr.Image()
     gr.Examples(
         [
             ["1. e4", ],
             ["1. e4 g6 2."],
         ],
         inputs=[prompt],
-        outputs=[output],
+        outputs=[output, img],
         fn=generate
     )
 demo.launch()