attention interface

README.md

@@ -1,5 +1,5 @@
 ---
-title: Gpt2 Stockfish Debug Viz
+title: GPT-2 Stockfish Debug
 emoji: 🔥
 colorFrom: blue
 colorTo: red

app.py

@@ -7,6 +7,7 @@ import wandb
 import gradio as gr
 
 from src import (
+    attention_interface,
     call_interface,
     play_interface,
     constants,
@@ -16,10 +17,12 @@ from src import (
 demo = gr.TabbedInterface(
     [
         play_interface.interface,
+        attention_interface.interface,
         call_interface.interface,
     ],
     [
         "Play",
+        "Attention Viz",
         "Call",
     ],
     title="GPT-2 Stockfish Debug",

requirements.txt

@@ -1,2 +1,3 @@
 python-chess
 wandb
+nnsight

src/attention_interface.py

@@ -0,0 +1,259 @@
+"""
+Gradio interface for plotting attention.
+"""
+
+import chess
+import gradio as gr
+import torch
+import uuid
+import re
+
+from . import constants, state, visualisation
+
+
+def compute_cache(
+    game_pgn,
+    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
+    state_cache = [(fen, state.model_cache(fen)) for fen in fen_list]
+    return (
+        *make_plot(
+            attention_layer, attention_head, comp_index, state_cache, state_board_index
+        ),
+        state_cache,
+    )
+
+
+def make_plot(
+    attention_layer,
+    attention_head,
+    comp_index,
+    state_cache,
+    state_board_index,
+):
+    if state_cache is None:
+        gr.Warning("Cache not computed!")
+        return None, None, None, None, None
+
+    fen, (out, cache) = state_cache[state_board_index]
+    attn_list = [a[0, attention_head - 1] for a in cache[attention_layer - 1]]
+    prompt_attn, *comp_attn = attn_list
+    comp_attn.insert(0, prompt_attn[-1:])
+    comp_attn = [a.squeeze(0) for a in comp_attn]
+    if len(comp_attn) != 5:
+        raise NotImplementedError("This is not implemented yet.")
+
+    config_total = meta_total = dump_total = 0
+    config_done = False
+    heatmap = torch.zeros(64)
+    h_index = 0
+    for i, t_o in enumerate(out[0]):
+        try:
+            t_attn = comp_attn[comp_index - 1][i]
+            if (i < 3) or (i > len(out[0]) - 10):
+                dump_total += t_attn
+                continue
+            t_str = state.model.tokenizer.decode(t_o)
+            if t_str.startswith(" ") and h_index > 0:
+                config_done = True
+            if not config_done:
+                if t_str == "/":
+                    dump_total += t_attn
+                    continue
+                t_str = re.sub(r"\d", lambda m: "0" * int(m.group(0)), t_str)
+                config_total += t_attn
+                t_str_len = len(t_str.strip())
+                pre_t_attn = t_attn / t_str_len
+                for j in range(t_str_len):
+                    heatmap[h_index + j] = pre_t_attn
+                h_index += t_str_len
+            else:
+                meta_total += t_attn
+        except IndexError:
+            break
+    raw_attention = comp_attn[comp_index - 1]
+    highlited_tokens = [
+        (state.model.tokenizer.decode(out[0][i]), raw_attention[i])
+        for i in range(len(raw_attention))
+    ]
+    uci_move = state.model.tokenizer.decode(out[0][-5:-1]).strip()
+    board = chess.Board(fen)
+    heatmap = heatmap.view(8, 8).flip(0).view(64)
+    move = chess.Move.from_uci(uci_move)
+    svg_board, fig = visualisation.render_heatmap(
+        board, heatmap, arrows=[(move.from_square, move.to_square)]
+    )
+    info = (
+        f"[Completion] Complete: '{state.model.tokenizer.decode(out[0][-5:])}'"
+        f" Chosen: '{state.model.tokenizer.decode(out[0][-5:][comp_index-1])}'"
+        f"\n[Distribution] Config: {config_total:.2f} Meta: {meta_total:.2f} Dump: {dump_total:.2f}"
+    )
+    id = str(uuid.uuid4())
+    with open(f"{constants.FIGURE_DIRECTORY}/board_{id}.svg", "w") as f:
+        f.write(svg_board)
+    return (
+        board.fen(),
+        info,
+        fig,
+        f"{constants.FIGURE_DIRECTORY}/board_{id}.svg",
+        highlited_tokens,
+    )
+
+
+def previous_board(
+    attention_layer,
+    attention_head,
+    comp_index,
+    state_cache,
+    state_board_index,
+):
+    state_board_index -= 1
+    if state_board_index < 0:
+        gr.Warning("Already at first board.")
+        state_board_index = 0
+    return (
+        *make_plot(
+            attention_layer, attention_head, comp_index, state_cache, state_board_index
+        ),
+        state_board_index,
+    )
+
+
+def next_board(
+    attention_layer,
+    attention_head,
+    comp_index,
+    state_cache,
+    state_board_index,
+):
+    state_board_index += 1
+    if state_board_index >= len(state_cache):
+        gr.Warning("Already at last board.")
+        state_board_index = len(state_cache) - 1
+    return (
+        *make_plot(
+            attention_layer, attention_head, comp_index, state_cache, state_board_index
+        ),
+        state_board_index,
+    )
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            game_pgn = gr.Textbox(
+                label="Game PGN",
+                lines=1,
+            )
+            compute_cache_button = gr.Button("Compute cache")
+            with gr.Group():
+                with gr.Row():
+                    attention_layer = gr.Slider(
+                        label="Attention layer",
+                        minimum=1,
+                        maximum=12,
+                        step=1,
+                        value=1,
+                    )
+                    attention_head = gr.Slider(
+                        label="Attention head",
+                        minimum=1,
+                        maximum=12,
+                        step=1,
+                        value=1,
+                    )
+                    comp_index = gr.Slider(
+                        label="Completion index",
+                        minimum=1,
+                        maximum=6,
+                        step=1,
+                        value=1,
+                    )
+                with gr.Row():
+                    previous_board_button = gr.Button("Previous board")
+                    next_board_button = gr.Button("Next board")
+            current_board_fen = gr.Textbox(
+                label="Board FEN",
+                lines=1,
+                max_lines=1,
+            )
+            info = gr.Textbox(
+                label="Info",
+                lines=1,
+                info=(
+                    "'Config' refers to the board configuration tokens."
+                    "\n'Meta' to the additional board tokens (like color or castling)."
+                    "\n'Dump' to the rest of the tokens (including '/')."
+                ),
+            )
+            gr.Markdown(
+                "Note that only the 'Config' attention is plotted.\n\nSee below for the raw attention."
+            )
+            raw_attention_html = gr.HighlightedText(
+                label="Raw attention",
+            )
+        with gr.Column():
+            image_board = gr.Image(label="Board")
+            colorbar = gr.Plot(label="Colorbar")
+
+    static_inputs = [
+        attention_layer,
+        attention_head,
+        comp_index,
+    ]
+    static_outputs = [
+        current_board_fen,
+        info,
+        colorbar,
+        image_board,
+        raw_attention_html,
+    ]
+
+    state_cache = gr.State(value=None)
+    state_board_index = gr.State(value=0)
+    compute_cache_button.click(
+        compute_cache,
+        inputs=[game_pgn, *static_inputs, state_cache, state_board_index],
+        outputs=[*static_outputs, state_cache],
+    )
+
+    previous_board_button.click(
+        previous_board,
+        inputs=[*static_inputs, state_cache, state_board_index],
+        outputs=[*static_outputs, state_board_index],
+    )
+    next_board_button.click(
+        next_board,
+        inputs=[*static_inputs, state_cache, state_board_index],
+        outputs=[*static_outputs, state_board_index],
+    )
+    attention_layer.change(
+        make_plot,
+        inputs=[*static_inputs, state_cache, state_board_index],
+        outputs=[*static_outputs],
+    )
+    attention_head.change(
+        make_plot,
+        inputs=[*static_inputs, state_cache, state_board_index],
+        outputs=[*static_outputs],
+    )
+    comp_index.change(
+        make_plot,
+        inputs=[*static_inputs, state_cache, state_board_index],
+        outputs=[*static_outputs],
+    )

src/call_interface.py

@@ -9,9 +9,7 @@ import gradio as gr
 model_name = "yp-edu/gpt2-stockfish-debug"
 
 headers = {"X-Wait-For-Model": "true"}
-client = huggingface_hub.InferenceClient(
-    model=model_name, headers=headers
-)
+client = huggingface_hub.InferenceClient(model=model_name, headers=headers)
 
 inputs = gr.Textbox(label="Prompt")
 outputs = gr.Textbox(label="Completion")

src/play_interface.py

@@ -15,17 +15,20 @@ import gradio as gr
 from . import constants
 
 model_name = "yp-edu/gpt2-stockfish-debug"
-headers = {"X-Wait-For-Model": "true"}
-client = huggingface_hub.InferenceClient(
-    model=model_name, headers=headers
-)
+headers = {
+    "X-Wait-For-Model": "true",
+    "X-Use-Cache": "false",
+}
+client = huggingface_hub.InferenceClient(model=model_name, headers=headers)
 inference_fn = client.text_generation
 
 
 def plot_board(
     board: chess.Board,
-    orientation: bool = chess.WHITE,
+    orientation: Optional[bool] = None,
 ):
+    if orientation is None:
+        orientation = board.turn
     try:
         last_move = board.peek()
         arrows = [(last_move.from_square, last_move.to_square)]
@@ -47,17 +50,17 @@ def plot_board(
         f.write(svg_board)
     return f"{constants.FIGURE_DIRECTORY}/board_{id}.svg"
 
+
 def render_board(
     current_board: chess.Board,
-    orientation: Optional[bool] = chess.WHITE,
+    orientation: Optional[bool] = None,
 ):
     fen = current_board.fen()
     pgn = current_board.root().variation_san(current_board.move_stack)
-    if orientation is None:
-        orientation = current_board.turn
     image_board = plot_board(current_board, orientation=orientation)
     return fen, pgn, "", image_board
 
+
 def play_user_move(
     uci_move: str,
     current_board: chess.Board,
@@ -65,6 +68,7 @@ def play_user_move(
     current_board.push_uci(uci_move)
     return current_board
 
+
 def play_ai_move(
     current_board: chess.Board,
     temperature: float = 0.1,
@@ -76,6 +80,7 @@ def play_ai_move(
     current_board.push_uci(uci_move.strip())
     return current_board
 
+
 def try_play_move(
     username: str,
     move_to_play: str,
@@ -83,7 +88,10 @@ def try_play_move(
 ):
     if current_board.is_game_over():
         gr.Warning("The game is already over")
-        return *render_board(current_board), current_board
+        return (
+            *render_board(current_board, orientation=not current_board.turn),
+            current_board,
+        )
     try:
         current_board = play_user_move(move_to_play.strip(), current_board)
         if current_board.is_game_over():
@@ -93,17 +101,20 @@ def try_play_move(
                     {
                         "username": username,
                         "winin": current_board.fullmove_number,
-                        "pgn": current_board.root().variation_san(current_board.move_stack),
+                        "pgn": current_board.root().variation_san(
+                            current_board.move_stack
+                        ),
                     }
                 )
                 run.finish()
-            return *render_board(current_board, orientation=not current_board.turn), current_board
+            return (
+                *render_board(current_board, orientation=not current_board.turn),
+                current_board,
+            )
     except:
         gr.Warning("Invalid move")
         return *render_board(current_board), current_board
-    temperature_retries = [
-        (i+1)/10 for i in range(10)
-    ]
+    temperature_retries = [(i + 1) / 10 for i in range(10)]
     for temperature in temperature_retries:
         try:
             current_board = play_ai_move(current_board, temperature=temperature)
@@ -187,6 +198,7 @@ with gr.Blocks() as interface:
         if is_ai_white:
             board = play_ai_move(board)
         return *render_board(board), board
+
     reset_button.click(
         reset_board,
         outputs=[*static_outputs, state_board],

src/state.py

@@ -0,0 +1,49 @@
+"""Global state of the app.
+"""
+
+import re
+
+from transformers import AutoConfig
+import torch
+from nnsight import LanguageModel
+
+conf = AutoConfig.from_pretrained("yp-edu/gpt2-stockfish-debug")
+model = LanguageModel("yp-edu/gpt2-stockfish-debug")
+model.eval()
+
+
+def make_prompt(fen):
+    board, player, castling, *fen_remaining = fen.split()
+    board = re.sub(r"(\d)", lambda m: "0" * int(m.group(1)), board)
+    spaced_board = " ".join(board)
+    spaced_castling = " ".join(castling)
+    full_fen = f"{spaced_board} {player} {spaced_castling} {' '.join(fen_remaining)}"
+    return f"FEN: {full_fen} \nMOVE:"
+
+
+def model_cache(fen):
+    global model
+    prompt = f"FEN: {fen}\nMOVE:"
+    attentions = {i: [] for i in range(12)}
+    with model.generate(prompt, max_new_tokens=10, output_attentions=True) as tracer:
+        out = model.generator.output.save()
+        for i in range(10):
+            for i in range(12):
+                attentions[i].append(model.transformer.h[i].attn.output[2].save())
+            tracer.next()
+    real_attentions = {}
+    for i in range(12):
+        real_attentions[i] = []
+        for a in attentions[i]:
+            try:
+                _ = a.shape
+                real_attentions[i].append(a)
+            except ValueError:
+                break
+    return out, real_attentions
+
+
+def attribute_seqence(fen, out, attn_tensor):
+    global model
+
+    out_str = model.tokenizer.batch_decode(out)[0]

src/visualisation.py

@@ -0,0 +1,72 @@
+"""
+Visualisation utils.
+"""
+
+import chess
+import chess.svg
+import matplotlib
+import matplotlib.pyplot as plt
+
+
+COLOR_MAP = matplotlib.colormaps["RdYlBu_r"].resampled(1000)
+ALPHA = 1.0
+
+
+def render_heatmap(
+    board,
+    heatmap,
+    square=None,
+    vmin=None,
+    vmax=None,
+    arrows=None,
+    normalise="none",
+):
+    """
+    Render a heatmap on the board.
+    """
+    if normalise == "abs":
+        a_max = heatmap.abs().max()
+        if a_max != 0:
+            heatmap = heatmap / a_max
+        vmin = -1
+        vmax = 1
+    if vmin is None:
+        vmin = heatmap.min()
+    if vmax is None:
+        vmax = heatmap.max()
+    norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=False)
+
+    color_dict = {}
+    for square_index in range(64):
+        color = COLOR_MAP(norm(heatmap[square_index]))
+        color = (*color[:3], ALPHA)
+        color_dict[square_index] = matplotlib.colors.to_hex(color, keep_alpha=True)
+    fig = plt.figure(figsize=(6, 0.6))
+    ax = plt.gca()
+    ax.axis("off")
+    fig.colorbar(
+        matplotlib.cm.ScalarMappable(norm=norm, cmap=COLOR_MAP),
+        ax=ax,
+        orientation="horizontal",
+        fraction=1.0,
+    )
+    if square is not None:
+        try:
+            check = chess.parse_square(square)
+        except ValueError:
+            check = None
+    else:
+        check = None
+    if arrows is None:
+        arrows = []
+    plt.close()
+    return (
+        chess.svg.board(
+            board,
+            check=check,
+            fill=color_dict,
+            size=350,
+            arrows=arrows,
+        ),
+        fig,
+    )