improve UI

app.py

@@ -9,6 +9,7 @@ import chess.pgn
 from svglib.svglib import svg2rlg
 from reportlab.graphics import renderPM
 from PIL import Image
+import os
 
 
 vocab_size=33
@@ -30,18 +31,30 @@ model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
 model.to(device)
 tokenizer = Tokenizer.from_pretrained(tokenizer_path)
 
+invalid_move_plot = Image.open('./invalid_move.png')
+
 def generate(prompt):
     model_input = torch.tensor(tokenizer.encode(prompt), dtype=torch.long, device=device).view((1, len(prompt)))
     pgn = tokenizer.decode(model.generate(model_input, max_new_tokens=4, context_size=context_size)[0].tolist())
     pgn_str = StringIO(pgn)
-    game = chess.pgn.read_game(pgn_str)
-    img = chess.svg.board(game.board())
+    try:
+        game = chess.pgn.read_game(pgn_str)
+        board = game.board()
+        for move in game.mainline_moves():
+            board.push(move)
+        img = chess.svg.board(board)
+    except Exception as e:
+        if 'illegal san' in str(e):
+            return pgn, invalid_move_plot
     filename = f'./moves-{pgn}'
     with open(filename + '.svg', 'w') as f:
         f.write(img)
     drawing = svg2rlg(filename + '.svg')
     renderPM.drawToFile(drawing, f"{filename}.png", fmt="PNG")
     plot = Image.open(f'{filename}.png')
+
+    os.remove(f'{filename}.png')
+    os.remove(f'{filename}.svg')
     return pgn, plot