Update app.py

app.py

@@ -50,14 +50,14 @@ def btn_load(inp_fen):
         f.write(str(chess.svg.board(board)))
     return 'board.svg', board.fen()
 
-def btn_play(inp_fen, inp_move, inp_notation):
+def btn_play(inp_fen, inp_move, inp_notation, inp_k):
     board = chess.Board(inp_fen)
     
     if inp_move:
         if inp_notation == 'UCI': mv = chess.Move.from_uci(inp_move) #board.push_uci(inp_move)
         elif inp_notation == 'SAN': mv = board.parse_san(inp_move) #chess.Move.from_san(inp_move) #board.push_san(inp_move)
     else:
-        mv = chess.Move.from_uci(predict_move(board.fen()))
+        mv = chess.Move.from_uci(predict_move(board.fen(), top_k=inp_k))
     
     board.push(mv)
         
@@ -70,14 +70,16 @@ with gr.Blocks() as block:
     gr.Markdown(
     '''
     # Play YoloChess - Policy Network v0.3
-    110M Parameter Transformer (BERT-base architecture) trained for text classification from scratch on expert games in FEN notation. 
+    110M Parameter Transformer (BERT-base architecture) trained for text classification from scratch on expert games in modified FEN notation. 
     '''
     )
     with gr.Row() as row:
         with gr.Column():
-            move = gr.Textbox(label='human player move')
-            notation = gr.Radio(["SAN", "UCI"], value="SAN", label='move notation')
+            with gr.Row():
+                move = gr.Textbox(label='human player move')
+                notation = gr.Radio(["SAN", "UCI"], value="SAN", label='move notation')
             fen = gr.Textbox(value=chess.Board().fen(), label='FEN')
+            top_k = gr.Number(value=3, label='pick from top_k moves', precision=0)
             with gr.Row():
                 load_btn = gr.Button("Load")
                 play_btn = gr.Button("Play")
@@ -94,7 +96,7 @@ with gr.Blocks() as block:
             position_output = gr.Image(label='board')
 
     load_btn.click(fn=btn_load, inputs=fen, outputs=[position_output, fen])
-    play_btn.click(fn=btn_play, inputs=[fen, move, notation], outputs=[position_output, fen, move])
+    play_btn.click(fn=btn_play, inputs=[fen, move, notation, top_k], outputs=[position_output, fen, move])
     
 
 block.launch()