fen entry

src/attention_interface.py

@@ -13,23 +13,28 @@ from . import constants, state, visualisation
 
 def compute_cache(
     game_pgn,
+    board_fen,
     attention_layer,
     attention_head,
     comp_index,
     state_cache,
     state_board_index,
 ):
-    board = chess.Board()
-    fen_list = [board.fen()]
-    for move in game_pgn.split():
-        if move.endswith("."):
-            continue
-        try:
-            board.push_san(move)
-            fen_list.append(board.fen())
-        except ValueError:
-            gr.Warning(f"Invalid move {move}, stopping before it.")
-            break
+    if game_pgn == "" and board_fen != "":
+        board = chess.Board(board_fen)
+        fen_list = [board.fen()]
+    else:
+        board = chess.Board()
+        fen_list = [board.fen()]
+        for move in game_pgn.split():
+            if move.endswith("."):
+                continue
+            try:
+                board.push_san(move)
+                fen_list.append(board.fen())
+            except ValueError:
+                gr.Warning(f"Invalid move {move}, stopping before it.")
+                break
     state_cache = [(fen, state.model_cache(fen)) for fen in fen_list]
     return (
         *make_plot(
@@ -156,10 +161,19 @@ def next_board(
 with gr.Blocks() as interface:
     with gr.Row():
         with gr.Column():
-            game_pgn = gr.Textbox(
-                label="Game PGN",
-                lines=1,
-            )
+            with gr.Group():
+                gr.Markdown(
+                    "Specify the game PGN of FEN string that you want to analyse (PGN overrides FEN)."
+                )
+                game_pgn = gr.Textbox(
+                    label="Game PGN",
+                    lines=1,
+                )
+                board_fen = gr.Textbox(
+                    label="Board FEN",
+                    lines=1,
+                    max_lines=1,
+                )
             compute_cache_button = gr.Button("Compute cache")
             with gr.Group():
                 with gr.Row():
@@ -228,7 +242,7 @@ with gr.Blocks() as interface:
     state_board_index = gr.State(value=0)
     compute_cache_button.click(
         compute_cache,
-        inputs=[game_pgn, *static_inputs, state_cache, state_board_index],
+        inputs=[game_pgn, board_fen, *static_inputs, state_cache, state_board_index],
         outputs=[*static_outputs, state_cache],
     )