implement interactive ui

app.py

@@ -34,28 +34,58 @@ 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())
+def gen_image_from_svg(img, filename):
+    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 plot
+
+
+def get_board(pgn):
     pgn_str = StringIO(pgn)
     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-{str(uuid4())}'
-    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')
+            return None
+    return board
 
-    os.remove(f'{filename}.png')
-    os.remove(f'{filename}.svg')
+def gen_board_image(pgn):
+    board = get_board(pgn)
+    return chess.svg.board(board)
+    
+def gen_move(pgn):
+    model_input = torch.tensor(tokenizer.encode(pgn), dtype=torch.long, device=device).view((1, len(pgn)))
+    is_invalid = True
+    board = get_board(pgn)
+    while is_invalid:
+        new_pgn = tokenizer.decode(model.generate(model_input, max_new_tokens=4, context_size=context_size, temperature=0.2)[0].tolist())
+        try:
+            print(f'checking {new_pgn}')
+            mv = new_pgn[len(pgn):].split(' ')[0]
+            new_pgn = pgn.rstrip() + f' {mv}'
+            is_invalid = get_board(new_pgn) is None
+        except Exception:
+            is_invalid = True
+            print(f'For {pgn} invalid "{new_pgn[len(pgn):].split(" ")[0]}" {new_pgn}')
+            #print(mov in board.legal_moves)
+    return new_pgn
+
+
+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, temperature=0.2)[0].tolist())
+    img = gen_board_image(pgn)
+    filename = f'./moves-{str(uuid4())}'
+    plot = gen_image_from_svg(img, filename)
     return pgn, plot
 
 
@@ -66,21 +96,71 @@ with gr.Blocks() as demo:
 
     The **C**hess-**P**re-trained-**T**ransformer.
 
-    The rules are simple: provide a PGN string of your current game, the engine will predict the next token!
+    The rules are simple: 
+    - "Interactive": Play a game against the model
+    - "Next turn prediction": provide a PGN string of your current game, the model will predict the next token
     """)
-    prompt = gr.Text(label="PGN")
-    output = gr.Text(label="Next turn", interactive=False)
-    img = gr.Image()
-    submit = gr.Button("Submit")
-    submit.click(generate, [prompt], [output, img])
-    
-    gr.Examples(
-        [
-            ["1. e4", ],
-            ["1. e4 g6 2."],
-        ],
-        inputs=[prompt],
-        outputs=[output, img],
-        fn=generate
-    )
+    def manual():
+        with gr.Tab("Next turn prediction"):
+            prompt = gr.Text(label="PGN")
+            output = gr.Text(label="Next turn", interactive=False)
+            img = gr.Image()
+            submit = gr.Button("Submit")
+            submit.click(generate, [prompt], [output, img])
+        
+            gr.Examples(
+                [
+                    ["1. e4", ],
+                    ["1. e4 g6 2."],
+                ],
+                inputs=[prompt],
+                outputs=[output, img],
+                fn=generate
+            )
+    def interactive():           
+        with gr.Tab("Interactive"):
+            color = gr.Dropdown(["white", "black"], value='white', label="Chose a color")
+            start_button = gr.Button("Start Game")
+
+            def start_game(c):
+                pgn = '1. '
+                if c == 'black':
+                    pgn += gen_move(pgn)
+                img = gen_board_image(pgn)
+                fn = 'foo'
+                return gen_image_from_svg(img, fn), pgn, 1
+
+            state = gr.Text(label='PGN', value='', interactive=False)
+            game = gr.Image()
+            move_counter = gr.State(value=1)
+            start_button.click(
+                start_game,
+                inputs=[color],
+                outputs=[game, state, move_counter]
+            )
+
+            next_move = gr.Text(label='Next move')
+            gen_next_move_button = gr.Button("Submit")
+
+            def gen_next_move(pgn, new_move, move_ctr, c):
+                pgn += new_move.strip() + ' '
+                if c == 'black':
+                    move_ctr += 1
+                    pgn = f'{pgn.rstrip()} {move_ctr}. '
+                print(f'gen for {pgn}')
+                pgn = gen_move(pgn)
+                print(f'got {pgn}')
+                img = gen_board_image(pgn)
+                if c == 'white':
+                    move_ctr += 1
+                    pgn = f'{pgn.rstrip()} {move_ctr}. '
+                return gen_image_from_svg(img, 'foo-bar'), pgn, move_ctr
+
+            gen_next_move_button.click(
+                gen_next_move,
+                inputs=[state, next_move, move_counter, color],
+                outputs=[game, state, move_counter]
+            )
+    interactive()
+    manual()
 demo.launch()

model.py

@@ -120,7 +120,7 @@ class DecoderTransformer(nn.Module):
         loss = F.cross_entropy(logits, targets)
         return logits, loss
 
-    def generate(self, idx, max_new_tokens=50, context_size=None):
+    def generate(self, idx, max_new_tokens=50, context_size=None, temperature=1.0):
         if context_size is None:
             context_size = int(self.position_embedding_table.weight.shape[0])
             print(context_size)
@@ -128,7 +128,7 @@ class DecoderTransformer(nn.Module):
         for _ in range(max_new_tokens):
             idx_cond = idx[:, -context_size:]
             logits, loss = self(idx_cond)
-            logits = logits[:,-1,:]
+            logits = logits[:,-1,:] / temperature
             probs = F.softmax(logits, dim=-1)
             idx_next = torch.multinomial(probs, num_samples=1)
             idx = torch.cat([idx, idx_next], dim=1)