probe_output_visualization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from fancy_einsum import einsum\n",
    "import chess\n",
    "import numpy as np\n",
    "import pickle\n",
    "import logging\n",
    "import plotly.graph_objects as go\n",
    "from functools import partial\n",
    "\n",
    "import chess_utils\n",
    "import train_test_chess"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's a bunch of setup below to get some data in some tensors that we can feed to our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flags to control logging\n",
    "debug_mode = False\n",
    "info_mode = True\n",
    "\n",
    "if debug_mode:\n",
    "    log_level = logging.DEBUG\n",
    "elif info_mode:\n",
    "    log_level = logging.INFO\n",
    "else:\n",
    "    log_level = logging.WARNING\n",
    "\n",
    "# Configure logging\n",
    "logging.basicConfig(level=log_level)\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here you can select which probe and model to use. By default, the model_setup.py downloads a lichess 8 layer model. We can then select a probe from saved_probes/. Ideally, this should also be a lichess probe. Then this code should auto populate parameters according to the probe's state dict.\n",
    "\n",
    "To reproduce paper / blog post figures, set USE_16_LAYER to True and run model_setup.py on the lichess 16 layer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_DIR = \"models/\"\n",
    "DATA_DIR = \"data/\"\n",
    "PROBE_DIR = \"linear_probes/\"\n",
    "SAVED_PROBE_DIR = \"linear_probes/saved_probes/\"\n",
    "SPLIT = \"test\"\n",
    "\n",
    "DEVICE = \"cpu\"\n",
    "logger.info(f\"Using device: {DEVICE}\")\n",
    "LAYER = 5\n",
    "base_probe_name = \"tf_lens_lichess_8layers_ckpt_no_optimizer_chess_piece_probe_layer_0.pth\"\n",
    "\n",
    "\n",
    "USE_16_LAYER = False\n",
    "\n",
    "if USE_16_LAYER:\n",
    "    LAYER = 11\n",
    "    base_probe_name = \"tf_lens_lichess_16layers_ckpt_no_optimizer_chess_piece_probe_layer_0.pth\"\n",
    "\n",
    "probe_to_test = base_probe_name.replace(\"layer_0\", f\"layer_{LAYER}\")\n",
    "\n",
    "num_games = 10\n",
    "sample_size = 1\n",
    "modes = 1\n",
    "\n",
    "probe_file_location = f\"{SAVED_PROBE_DIR}{probe_to_test}\"\n",
    "with open(probe_file_location, \"rb\") as f:\n",
    "    state_dict = torch.load(f, map_location=torch.device(DEVICE))\n",
    "    print(state_dict.keys())\n",
    "    for key in state_dict.keys():\n",
    "        if key != \"linear_probe\":\n",
    "            print(key, state_dict[key])\n",
    "\n",
    "    config = chess_utils.find_config_by_name(state_dict[\"config_name\"])\n",
    "    layer = state_dict[\"layer\"]\n",
    "    model_name = state_dict[\"model_name\"]\n",
    "    dataset_prefix = state_dict[\"dataset_prefix\"]\n",
    "    column_name = state_dict[\"column_name\"]\n",
    "    config.pos_start = state_dict[\"pos_start\"]\n",
    "    levels_of_interest = None\n",
    "    if \"levels_of_interest\" in state_dict.keys():\n",
    "        levels_of_interest = state_dict[\"levels_of_interest\"]\n",
    "    config.levels_of_interest = levels_of_interest\n",
    "    indexing_function_name = state_dict[\"indexing_function_name\"]\n",
    "    n_layers = state_dict[\"n_layers\"]\n",
    "    \n",
    "\n",
    "    split = SPLIT\n",
    "    input_dataframe_file = f\"{DATA_DIR}{dataset_prefix}{split}.csv\"\n",
    "    config = chess_utils.set_config_min_max_vals_and_column_name(\n",
    "        config, input_dataframe_file, dataset_prefix\n",
    "    )\n",
    "    misc_logging_dict = {\n",
    "        \"split\": split,\n",
    "        \"dataset_prefix\": dataset_prefix,\n",
    "        \"model_name\": model_name,\n",
    "        \"n_layers\": n_layers,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At the end of the below cell, we index at select 1 of the num_games. The reason we do this is that with a large number of games, storing all the resid_posts and state_stacks quickly grows to many gigabytes of VRAM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probe_data = train_test_chess.construct_linear_probe_data(\n",
    "    input_dataframe_file,\n",
    "    dataset_prefix,\n",
    "    n_layers,\n",
    "    model_name,\n",
    "    config,\n",
    "    num_games,\n",
    "    DEVICE,\n",
    ")\n",
    "if DEVICE == \"cpu\":\n",
    "    probe_data.model.cpu()\n",
    "\n",
    "game_of_interest = 3\n",
    "\n",
    "game_length_in_chars = len(probe_data.board_seqs_string[0])\n",
    "\n",
    "\n",
    "state_stacks_all_chars_MBlRR = chess_utils.create_state_stacks(probe_data.board_seqs_string[:num_games], config.custom_board_state_function)\n",
    "logger.info(f\"state_stack shape: {state_stacks_all_chars_MBlRR.shape}\")\n",
    "assert(state_stacks_all_chars_MBlRR.shape) == (modes, num_games, game_length_in_chars, config.num_rows, config.num_cols)\n",
    "white_move_indices_BL = probe_data.custom_indices[:num_games]\n",
    "print(white_move_indices_BL.shape)\n",
    "num_white_moves = white_move_indices_BL.shape[1]\n",
    "assert(white_move_indices_BL.shape) == (num_games, num_white_moves)\n",
    "\n",
    "\n",
    "print(\"\\nSelecting the game of interest\")\n",
    "print(probe_data.board_seqs_int.shape)\n",
    "print(state_stacks_all_chars_MBlRR.shape)\n",
    "print(white_move_indices_BL.shape)\n",
    "print(len(probe_data.board_seqs_string), len(probe_data.board_seqs_string[0]))\n",
    "\n",
    "probe_data.board_seqs_int = probe_data.board_seqs_int[game_of_interest].unsqueeze(0)\n",
    "probe_data.board_seqs_string = [probe_data.board_seqs_string[game_of_interest]]\n",
    "probe_data.custom_indices = white_move_indices_BL[game_of_interest].unsqueeze(0)\n",
    "state_stacks_all_chars_MBlRR = state_stacks_all_chars_MBlRR[:, game_of_interest, :, :, :].unsqueeze(1)\n",
    "white_move_indices_BL = white_move_indices_BL[game_of_interest].unsqueeze(0)\n",
    "\n",
    "print(probe_data.board_seqs_int.shape)\n",
    "print(state_stacks_all_chars_MBlRR.shape)\n",
    "print(white_move_indices_BL.shape)\n",
    "print(len(probe_data.board_seqs_string), len(probe_data.board_seqs_string[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an explanation of all the data we just generated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"All pgn strings are of length {game_length_in_chars}\")\n",
    "print(f\"For game {game_of_interest}, the pgn string is {probe_data.board_seqs_string[0]}\")\n",
    "print(f\"Using our encode functions, it's represented as ints that are fed as input to the GPT model with shape {probe_data.board_seqs_int.shape}\")\n",
    "print(f\"The first 30 characters of board_seqs_ints looks like this: {probe_data.board_seqs_int[:, :30]}\")\n",
    "print(f\"state_stacks_all_chars contains the board state at every char index in the pgn string with shape {state_stacks_all_chars_MBlRR.shape}\")\n",
    "print(f\"white_move_indices contains the index of every white move in the pgn string with shape {white_move_indices_BL.shape}\")\n",
    "print(f\"That means there are {num_white_moves} white moves in the game\")\n",
    "print(f\"For example, in {probe_data.board_seqs_string[0][:14]}, the white move indices are {white_move_indices_BL[:, :2]} (the indices of each period)\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Important note: At the bottom of the below cell, I currently am using softmax to view probe output probabilities. You can comment that out to view raw logits instead.\n",
    "\n",
    "In this cell, we input the board_seqs_int to the GPT to obtain resid_post, the intermediate activations after our layer of interest. We index into resid_post using white_move_indices. These indexed resid_posts are then input to the linear probe, which outputs probe_out, a probability distribution for the state of every square on the board."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load(probe_file_location, map_location=torch.device(DEVICE))\n",
    "linear_probe_MDRRC = checkpoint[\"linear_probe\"]\n",
    "print(linear_probe_MDRRC.shape)\n",
    "\n",
    "\n",
    "one_hot_range = config.max_val - config.min_val + 1\n",
    "\n",
    "board_seqs_int_BL = probe_data.board_seqs_int[:].to(DEVICE)\n",
    "assert(board_seqs_int_BL.shape) == (1, game_length_in_chars)\n",
    "\n",
    "indexed_state_stacks_MBLRR = []\n",
    "\n",
    "for batch_idx in range(sample_size):\n",
    "    # Get the indices for the current batch\n",
    "    dots_indices_for_batch_L = white_move_indices_BL[batch_idx]\n",
    "\n",
    "    # Index the state_stack for the current batch. Adding an unsqueeze operation to maintain the batch dimension.\n",
    "    indexed_state_stack_MLRR = state_stacks_all_chars_MBlRR[:, batch_idx:batch_idx+1, dots_indices_for_batch_L, :, :]\n",
    "\n",
    "    # Append the result to the list\n",
    "    indexed_state_stacks_MBLRR.append(indexed_state_stack_MLRR)\n",
    "\n",
    "# Concatenate the indexed state stacks along the second dimension (batch dimension)\n",
    "# Since we're maintaining the batch dimension during indexing, we don't need to add it back in.\n",
    "state_stack_white_moves_MBLRR = torch.cat(indexed_state_stacks_MBLRR, dim=1)\n",
    "\n",
    "print(\"state stack shapes\")\n",
    "print(state_stack_white_moves_MBLRR.shape)\n",
    "print(state_stacks_all_chars_MBlRR.shape)\n",
    "\n",
    "with torch.inference_mode():\n",
    "    _, cache = probe_data.model.run_with_cache(board_seqs_int_BL[:, :-1], return_type=None)\n",
    "    resid_post_BlD = cache[\"resid_post\", layer][:, :]\n",
    "\n",
    "# print(resid_post.shape)\n",
    "assert(resid_post_BlD.shape) == (sample_size, game_length_in_chars - 1, linear_probe_MDRRC.shape[1])\n",
    "\n",
    "# Initialize a list to hold the indexed state stacks\n",
    "indexed_resid_posts_BLD = []\n",
    "\n",
    "for batch_idx in range(sample_size):\n",
    "    # Get the indices for the current batch\n",
    "    dots_indices_for_batch_L = white_move_indices_BL[batch_idx]\n",
    "\n",
    "    # Index the state_stack for the current batch\n",
    "    indexed_resid_post_LD = resid_post_BlD[batch_idx, dots_indices_for_batch_L]\n",
    "\n",
    "    # Append the result to the list\n",
    "    indexed_resid_posts_BLD.append(indexed_resid_post_LD)\n",
    "\n",
    "# Stack the indexed state stacks along the first dimension\n",
    "# This results in a tensor of shape [2, 61, 8, 8] (assuming all batches have 61 indices)\n",
    "resid_post_BLD = torch.stack(indexed_resid_posts_BLD)\n",
    "resid_post_BLD = resid_post_BLD.to(DEVICE)\n",
    "print(\"Resid post\", resid_post_BLD.shape)\n",
    "probe_out_MBLRRC = einsum(\n",
    "    \"batch pos d_model, modes d_model rows cols options -> modes batch pos rows cols options\",\n",
    "    resid_post_BLD,\n",
    "    linear_probe_MDRRC,\n",
    ")\n",
    "probe_out_MBLRRC = probe_out_MBLRRC.log_softmax(-1)\n",
    "print(f\"Probe out shape: {probe_out_MBLRRC.shape}\")\n",
    "assert(probe_out_MBLRRC.shape) == (modes, sample_size, white_move_indices_BL.shape[1], config.num_rows, config.num_cols, one_hot_range)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here you can select which move you want to visualize (move_of_interest)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "move_of_interest = 11\n",
    "GAME_IDX = 0 # After refactoring to discard unused games, this is always 0\n",
    "move_of_interest_index = white_move_indices_BL[GAME_IDX][move_of_interest] # Used to select pgn strings\n",
    "move_of_interest_state_RR = state_stack_white_moves_MBLRR[0][GAME_IDX][move_of_interest]\n",
    "print(move_of_interest_state_RR.shape)\n",
    "print(move_of_interest_state_RR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we one hot encode our move_of_interest and store it in move_of_interest_state_one_hot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_stacks_one_hot_MBLRRC = chess_utils.state_stack_to_one_hot(modes, config.num_rows, config.num_cols, config.min_val, config.max_val, DEVICE, state_stack_white_moves_MBLRR)\n",
    "print(state_stacks_one_hot_MBLRRC.shape)\n",
    "assert(state_stacks_one_hot_MBLRRC.shape) == (modes, sample_size, num_white_moves, config.num_rows, config.num_cols, one_hot_range)\n",
    "move_of_interest_state_one_hot_RRC = state_stacks_one_hot_MBLRRC[0][GAME_IDX][move_of_interest]\n",
    "print(move_of_interest_state_one_hot_RRC.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We get the argmax of each square's probe probability distribution and store it in state_stacks_probe_outputs for easy graphing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(move_of_interest_state_one_hot_RRC.shape)\n",
    "print(state_stacks_one_hot_MBLRRC.shape)\n",
    "state_stacks_probe_outputs_MBLRR = chess_utils.one_hot_to_state_stack(probe_out_MBLRRC, config.min_val)\n",
    "state_stacks_probe_outputs_MBLRR = torch.tensor(state_stacks_probe_outputs_MBLRR)\n",
    "print(state_stacks_probe_outputs_MBLRR.shape)\n",
    "assert(state_stacks_probe_outputs_MBLRR.shape) == (modes, sample_size, num_white_moves, config.num_rows, config.num_cols)\n",
    "print(state_stacks_probe_outputs_MBLRR[0][GAME_IDX][move_of_interest])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Change blank_index, king_index, or pawn_index if you want to visualize the probe's view of other pieces. For example, if I want to see the black queen, I could set blank_index = -5 (refer to INT_TO_CHAR for the mapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "INT_TO_CHAR = {\n",
    "    -6: \"\\u265a\",\n",
    "    -5: \"\\u265b\",\n",
    "    -4: \"\\u265c\",\n",
    "    -3: \"\\u265d\",\n",
    "    -2: \"\\u265e\",\n",
    "    -1: \"\\u265f\",\n",
    "    0: \".\",\n",
    "    1: \"\\u2659\",\n",
    "    2: \"\\u2658\",\n",
    "    3: \"\\u2657\",\n",
    "    4: \"\\u2656\",\n",
    "    5: \"\\u2655\",\n",
    "    6: \"\\u2654\",\n",
    "}\n",
    "\n",
    "# Mapping of integers to chess pieces\n",
    "# I'm duplicating this from chess_utils.py for easy reference\n",
    "PIECE_TO_ONE_HOT_MAPPING = {\n",
    "    -6: 0,\n",
    "    -5: 1,\n",
    "    -4: 2,\n",
    "    -3: 3,\n",
    "    -2: 4,\n",
    "    -1: 5,\n",
    "    0: 6,\n",
    "    1: 7,\n",
    "    2: 8,\n",
    "    3: 9,\n",
    "    4: 10,\n",
    "    5: 11,\n",
    "    6: 12,\n",
    "}\n",
    "\n",
    "# Mapping of chess pieces to integers\n",
    "PIECE_TO_INT = {\n",
    "    chess.PAWN: 1,\n",
    "    chess.KNIGHT: 2,\n",
    "    chess.BISHOP: 3,\n",
    "    chess.ROOK: 4,\n",
    "    chess.QUEEN: 5,\n",
    "    chess.KING: 6,\n",
    "}\n",
    "\n",
    "INT_TO_PIECE = {value: key for key, value in PIECE_TO_INT.items()}\n",
    "\n",
    "BLANK_INDEX = PIECE_TO_ONE_HOT_MAPPING[0]\n",
    "white_pawn_index = PIECE_TO_ONE_HOT_MAPPING[1]\n",
    "black_king_index = PIECE_TO_ONE_HOT_MAPPING[-6]\n",
    "\n",
    "def plot_board_state(board_state: torch.Tensor, clip_size: int = 200, show_scale: bool = False):\n",
    "    # color scale: Black for -1, Gray for 0, White for 1\n",
    "    # colorscale = [[0.0, 'black'], [0.5, 'gray'], [1.0, 'white']]\n",
    "    colorscale = 'gray'\n",
    "    if board_state.is_cuda:\n",
    "        board_state = board_state.cpu()\n",
    "    board_state = np.clip(board_state.numpy(), -clip_size, clip_size)\n",
    "\n",
    "    # Create heatmap\n",
    "    heatmap = go.Heatmap(z=board_state, colorscale=colorscale, showscale=show_scale)\n",
    "    return heatmap\n",
    "\n",
    "print(move_of_interest_state_one_hot_RRC[:, :, white_pawn_index])\n",
    "# heatmap = plot_board_state(move_of_interest_state_one_hot[:, :, white_pawn_index], show_scale=True)\n",
    "\n",
    "move_of_interest_probe_out = probe_out_MBLRRC[0][0][move_of_interest]\n",
    "print(move_of_interest_probe_out.shape)\n",
    "\n",
    "heatmap = plot_board_state(move_of_interest_probe_out[:, :, white_pawn_index], show_scale=True)\n",
    "\n",
    "# Define the layout\n",
    "layout = go.Layout(\n",
    "    title=\"Chess board white pawns\",\n",
    "    xaxis=dict(ticks='', nticks=8),\n",
    "    yaxis=dict(ticks='', nticks=8),\n",
    "    autosize=False,\n",
    "    width=600,\n",
    "    height=600\n",
    ")\n",
    "\n",
    "# Create figure and plot\n",
    "fig = go.Figure(data=[heatmap], layout=layout)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tensor_to_text(board_state: torch.Tensor) -> np.ndarray:\n",
    "    # Create a mapping from numbers to characters\n",
    "    # Update this mapping according to your requirements\n",
    "\n",
    "    # Convert the tensor to numpy array for easier processing\n",
    "    board_array = board_state.numpy()\n",
    "\n",
    "    # Create an empty array with the same shape for text\n",
    "    text_array = np.empty(board_array.shape, dtype=str)\n",
    "\n",
    "    # Fill the text array with corresponding characters\n",
    "    for i in range(board_array.shape[0]):\n",
    "        for j in range(board_array.shape[1]):\n",
    "            text_array[i, j] = INT_TO_CHAR.get(board_array[i, j], str(board_array[i, j]))\n",
    "\n",
    "    return text_array\n",
    "\n",
    "def plot_board_state_with_text(board_state: torch.Tensor):\n",
    "    # Convert the tensor to a text matrix\n",
    "    text_matrix = tensor_to_text(board_state)\n",
    "\n",
    "    # Define the custom colorscale\n",
    "    colorscale = [\n",
    "        [0, 'white'],   # Negative values\n",
    "        [0.49, 'white'],\n",
    "        [0.5, 'grey'],  # Zero\n",
    "        [0.51, 'white'],\n",
    "        [1, 'white']    # Positive values\n",
    "    ]\n",
    "\n",
    "\n",
    "    # Create heatmap with text and custom colorscale\n",
    "    heatmap = go.Heatmap(\n",
    "        z=board_state.numpy(), \n",
    "        text=text_matrix, \n",
    "        showscale=False, \n",
    "        colorscale=colorscale,\n",
    "        texttemplate=\"%{text}\",\n",
    "        textfont=dict(size=48) \n",
    "    )\n",
    "\n",
    "    return heatmap\n",
    "heatmap = plot_board_state_with_text(move_of_interest_state_RR)\n",
    "\n",
    "# Define the layout\n",
    "layout = go.Layout(\n",
    "    title=\"Chess board state with text\",\n",
    "    xaxis=dict(ticks='', nticks=8),\n",
    "    yaxis=dict(ticks='', nticks=8),\n",
    "    autosize=False,\n",
    "    width=600,\n",
    "    height=600\n",
    ")\n",
    "\n",
    "# Create figure and plot\n",
    "fig = go.Figure(data=[heatmap], layout=layout)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotly.subplots import make_subplots\n",
    "\n",
    "move_of_interest_probe_out = probe_out_MBLRRC[0][0][move_of_interest]\n",
    "print(move_of_interest_probe_out.shape)\n",
    "\n",
    "fig_rows = 4\n",
    "fig_cols = 3\n",
    "fig = make_subplots(rows=fig_rows, cols=fig_cols, subplot_titles=[\n",
    "    \"Ground truth blank squares\", \"Predicted blank squares\", \"Confidence gradient blank squares\",\n",
    "    \"Ground truth white pawn positions\", \"Predicted white pawn positions\", \"Confidence gradient white pawn positions\",\n",
    "    \"Ground truth black king position\", \"Predicted black king position\", \"Confidence gradient black king position\",\n",
    "    \"Ground truth state\", \"Predicted board state\", \"Redundant probe output board state\"\n",
    "])\n",
    "\n",
    "\n",
    "# Specify the size of each plot\n",
    "plot_size = 400  # You can adjust this size\n",
    "\n",
    "fig.add_trace(plot_board_state(move_of_interest_state_one_hot_RRC[:, :, BLANK_INDEX]), row=1, col=1)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, BLANK_INDEX], clip_size=2), row=1, col=2)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, BLANK_INDEX]), row=1, col=3)\n",
    "\n",
    "fig.add_trace(plot_board_state(move_of_interest_state_one_hot_RRC[:, :, white_pawn_index]), row=2, col=1)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, white_pawn_index], clip_size=2), row=2, col=2)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, white_pawn_index], show_scale=True), row=2, col=3)\n",
    "\n",
    "fig.add_trace(plot_board_state(move_of_interest_state_one_hot_RRC[:, :, black_king_index]), row=3, col=1)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, black_king_index], clip_size=2), row=3, col=2)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, black_king_index]), row=3, col=3)\n",
    "\n",
    "fig.add_trace(plot_board_state_with_text(move_of_interest_state_RR), row=4, col=1)\n",
    "fig.add_trace(plot_board_state_with_text(state_stacks_probe_outputs_MBLRR[0][0][move_of_interest]), row=4, col=2)\n",
    "fig.add_trace(plot_board_state_with_text(state_stacks_probe_outputs_MBLRR[0][0][move_of_interest]), row=4, col=2)\n",
    "\n",
    "# Adjust the overall size of the figure\n",
    "fig.update_layout(height=fig_rows * plot_size * 1.3, width=fig_cols * plot_size)\n",
    "fig.update_annotations(dict(font=dict(size=18))) \n",
    "\n",
    "\n",
    "\n",
    "# Show the figure\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will check the percentage of squares in the sample (sample_size defaults to 1 game) where the ground truth matches the probe output.\n",
    "I also do a round trip through all the transformations, which should match 100%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_matching_percentage(state_stacks: torch.Tensor, probe_outputs: torch.Tensor) -> float:\n",
    "    \"\"\"\n",
    "    Calculate the percentage of matching cells in two tensors.\n",
    "\n",
    "    :param state_stacks: A tensor of shape [1, 1, 680, 8, 8].\n",
    "    :param probe_outputs: A tensor of shape [1, 1, 680, 8, 8].\n",
    "    :return: The percentage of cells that match.\n",
    "    \"\"\"\n",
    "    # Element-wise comparison\n",
    "    matches = state_stacks == probe_outputs\n",
    "\n",
    "    # Count the number of matches\n",
    "    num_matches = matches.sum().item()\n",
    "\n",
    "    # Total number of elements\n",
    "    total_elements = state_stacks.numel()\n",
    "\n",
    "    # Calculate percentage\n",
    "    percentage = (num_matches / total_elements) * 100\n",
    "    print(f\"Out of {total_elements} elements, {num_matches} matched, {percentage}%\")\n",
    "\n",
    "    return percentage\n",
    "assert(state_stacks_probe_outputs_MBLRR.shape) == (state_stack_white_moves_MBLRR.shape)\n",
    "print(\"Linear probe accuracy on all board squares in sample size:\", calculate_matching_percentage(state_stack_white_moves_MBLRR, state_stacks_probe_outputs_MBLRR))\n",
    "\n",
    "round_trip = chess_utils.one_hot_to_state_stack(chess_utils.state_stack_to_one_hot(modes, config.num_rows, config.num_cols, config.min_val,config.max_val, DEVICE, state_stack_white_moves_MBLRR), config.min_val)\n",
    "round_trip = torch.tensor(round_trip)\n",
    "print(round_trip.shape)\n",
    "print(state_stack_white_moves_MBLRR.shape)\n",
    "assert(round_trip.shape) == (modes, sample_size, num_white_moves, config.num_rows, config.num_cols)\n",
    "assert(round_trip.shape) == state_stack_white_moves_MBLRR.shape\n",
    "matching_percentage = calculate_matching_percentage(round_trip, state_stack_white_moves_MBLRR)\n",
    "assert(matching_percentage == 100.0)\n",
    "print(f\"Round trip matching percentage: {matching_percentage}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can perform interventions on the model's internals and view the modified probe outputs. We can also verify the model produces legal moves under the modified state of the board.\n",
    "\n",
    "First, we perform a sanity check to ensure that our interventions on model activations are working correctly. In this case, diff should roughly equal flip_dir.\n",
    "\n",
    "Note that I'm only intervening on one layer here. By modifying the first for loop and training additional probes, we can easily intervene on an arbitrary amount of layers. If we were to intervene on multiple layers, we can only check that torch.allclose(diff, flip_dirs[layer], atol=1e-6) for the first layer that we intervene on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probe_data.model.reset_hooks()\n",
    "\n",
    "_, cache = probe_data.model.run_with_cache(board_seqs_int_BL.to(DEVICE)[:, :-1], return_type=None)\n",
    "resid_post_BlD = cache[\"resid_post\", layer][:, :]\n",
    "\n",
    "r = 0\n",
    "c = 0\n",
    "\n",
    "probe_names = {}\n",
    "for i in range(layer, layer + 1):\n",
    "    probe_names[i] = base_probe_name.replace(\"layer_0\", f\"layer_{i}\")\n",
    "\n",
    "probes = {}\n",
    "\n",
    "# Use this to intervene on multiple layers\n",
    "for layer, probe_name in probe_names.items():\n",
    "    probe_file_location = f\"{SAVED_PROBE_DIR}{probe_name}\"\n",
    "    checkpoint = torch.load(probe_file_location, map_location=torch.device(DEVICE))\n",
    "    linear_probe_MDRRC = checkpoint[\"linear_probe\"]\n",
    "    probes[layer] = linear_probe_MDRRC\n",
    "\n",
    "\n",
    "flip_dirs = {}\n",
    "\n",
    "piece1 = BLANK_INDEX\n",
    "piece2 = black_king_index\n",
    "\n",
    "for layer, linear_probe_MDRRC in probes.items():\n",
    "    piece1_probe = linear_probe_MDRRC[:, :, r, c, piece1].squeeze()\n",
    "    piece2_probe = linear_probe_MDRRC[:, :, r, c, piece2].squeeze()\n",
    "    flip_dir = piece2_probe - piece1_probe\n",
    "    flip_dir.to(DEVICE)\n",
    "    flip_dirs[layer] = flip_dir\n",
    "\n",
    "def flip_hook(resid, hook, flip_dir: torch.Tensor):\n",
    "    resid[GAME_IDX, :] -= flip_dir # NOTE: We could only intervene on a single position in the sequence, but there's no harm in intervening on all of them\n",
    "\n",
    "probe_data.model.reset_hooks()\n",
    "\n",
    "for layer, flip_dir in flip_dirs.items():\n",
    "    temp_hook_fn = partial(flip_hook, flip_dir=flip_dir)\n",
    "    hook_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "    probe_data.model.add_hook(hook_name, temp_hook_fn)\n",
    "\n",
    "print(probe_data.model.cpu())\n",
    "_, modified_cache = probe_data.model.run_with_cache(board_seqs_int_BL.to(DEVICE)[:, :-1])\n",
    "probe_data.model.reset_hooks()\n",
    "modified_resid_post = modified_cache[\"resid_post\", layer][:, :]\n",
    "\n",
    "print(resid_post_BlD.shape)\n",
    "print(modified_resid_post.shape)\n",
    "\n",
    "diff = resid_post_BlD[GAME_IDX, 10, :] - modified_resid_post[GAME_IDX, 10, :]\n",
    "\n",
    "assert torch.allclose(diff, flip_dirs[layer], atol=1e-6)\n",
    "print(\"Flip hook test passed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we load the model's vocab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"models/meta.pkl\", \"rb\") as f:\n",
    "    meta = pickle.load(f)\n",
    "stoi = meta[\"stoi\"]\n",
    "itos = meta[\"itos\"]\n",
    "def encode_string(s: str) -> list[int]:\n",
    "    \"\"\"Encode a string into a list of integers.\"\"\"\n",
    "    return [stoi[c] for c in s]\n",
    "\n",
    "\n",
    "def decode_list(l: list[int]) -> str:\n",
    "    \"\"\"Decode a list of integers into a string.\"\"\"\n",
    "    return \"\".join([itos[i] for i in l])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we generate 10 characters using the model to determine the model's next move. Note that we are using argmax instead of a temperature based approach, so this will always return the most likely move.One annoying problem we deal with: In chess, the 0th row is at the bottom, which is how print(chess_board) displays everything. But, for our state stack (and any array), the 0th row is at the top.\n",
    "\n",
    "Now, we get a pgn string up to the current move and convert it to a chess board. We use it to create an encoded model_input as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(move_of_interest_state_RR)\n",
    "\n",
    "pgn_string = probe_data.board_seqs_string[GAME_IDX][:move_of_interest_index + 1]\n",
    "model_input = encode_string(pgn_string)\n",
    "model_input = torch.tensor(model_input).unsqueeze(0).to(DEVICE)\n",
    "print(model_input.shape)\n",
    "board = chess_utils.pgn_string_to_board(pgn_string)\n",
    "\n",
    "print(board)\n",
    "print(board.legal_moves)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We generate a move using the model on the original board and check that the move is legal. Next, we determine which piece was moved, and which row / column the source square of the move was."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_move = chess_utils.get_model_move(probe_data.model, meta, model_input)\n",
    "model_move_san = board.parse_san(model_move)\n",
    "assert model_move_san in board.legal_moves\n",
    "\n",
    "moved_piece = board.piece_at(model_move_san.from_square)\n",
    "moved_piece_int = PIECE_TO_INT[moved_piece.piece_type]\n",
    "moved_piece_probe_index = PIECE_TO_ONE_HOT_MAPPING[moved_piece_int]\n",
    "source_square = chess.square_name(model_move_san.from_square)\n",
    "\n",
    "\n",
    "r, c = chess_utils.square_to_coordinate(model_move_san.from_square)\n",
    "print(r, c)\n",
    "\n",
    "print(f\"Model move: {model_move_san}, moved piece: {moved_piece}, moved piece int: {moved_piece_int}, moved piece probe index: {moved_piece_probe_index}, source square: {source_square}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create a modified board where the source square of the model's original move is blank."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modified_state_stack = state_stack_white_moves_MBLRR.clone()\n",
    "modified_state_stack[0, GAME_IDX, move_of_interest, r, c] = 0\n",
    "modified_move_of_interest_state = modified_state_stack[0, GAME_IDX, move_of_interest]\n",
    "modified_state_stacks_one_hot = chess_utils.state_stack_to_one_hot(modes, config.num_rows, config.num_cols, config.min_val, config.max_val, DEVICE, modified_state_stack)\n",
    "modified_move_of_interest_state_one_hot = modified_state_stacks_one_hot[0][GAME_IDX][move_of_interest]\n",
    "modified_board = board.copy()\n",
    "modified_board.set_piece_at(model_move_san.from_square, None)\n",
    "print(modified_board)\n",
    "print(modified_board.legal_moves)\n",
    "\n",
    "assert modified_move_of_interest_state_one_hot.shape == move_of_interest_state_one_hot_RRC.shape\n",
    "assert modified_state_stack.shape == state_stack_white_moves_MBLRR.shape\n",
    "assert modified_state_stacks_one_hot.shape == state_stacks_one_hot_MBLRRC.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we get flip_dir, which is a probe of piece * piece_coefficient - blank square * blank_coefficient. In practice, I find that it works best when blank_coefficient is 0. We subtract this flip_dir from the model's activations at every token. We generate 10 new characters using the model, and verify that the new move under this modified state is legal according to the modified state. We also save a copy of the modified activations and generate modified probe outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, cache = probe_data.model.run_with_cache(board_seqs_int_BL.to(DEVICE)[:, :-1], return_type=None)\n",
    "resid_post_BlD = cache[\"resid_post\", layer][:, :]\n",
    "\n",
    "flip_dirs = {}\n",
    "\n",
    "piece1 = BLANK_INDEX\n",
    "piece1_probe = linear_probe_MDRRC[:, :, r, c, piece1].squeeze()\n",
    "piece2 = moved_piece_probe_index\n",
    "\n",
    "for layer, linear_probe_MDRRC in probes.items():\n",
    "    piece2_probe = linear_probe_MDRRC[:, :, r, c, piece2].squeeze()\n",
    "    flip_dir = piece2_probe - piece1_probe\n",
    "    flip_dir.to(DEVICE)\n",
    "    flip_dirs[layer] = flip_dir\n",
    "\n",
    "def flip_hook(resid, hook, flip_dir: torch.Tensor):\n",
    "    # print(resid[0, move_of_interest_index, :].shape)\n",
    "    # print(flip_dir.shape)\n",
    "    # print(piece1_probe.shape)\n",
    "    # left_side = torch.dot(resid[0, move_of_interest_index, :], piece1_probe) - 3.0\n",
    "    # right_side = torch.dot(flip_dir, piece1_probe)\n",
    "    # scale = left_side / right_side\n",
    "    # print(scale)\n",
    "    \n",
    "    # # Calculate scale\n",
    "    # scale = left_side / right_side\n",
    "    piece_coefficient = 1.0\n",
    "    blank_coefficient = 0.0\n",
    "    blank_probe = probes[layer][:, :, r, c, BLANK_INDEX].squeeze()\n",
    "    piece_probe = probes[layer][:, :, r, c, moved_piece_probe_index].squeeze()\n",
    "\n",
    "    flip_dir = (piece_probe * piece_coefficient) - (blank_probe * blank_coefficient)\n",
    "    flip_dir = flip_dir / flip_dir.norm()\n",
    "    scale = 1.0\n",
    "    resid[0, :] -= scale * flip_dir # NOTE: We could only intervene on a single position in the sequence, but there's no harm in intervening on all of them\n",
    "\n",
    "probe_data.model.reset_hooks()\n",
    "\n",
    "for layer, flip_dir in flip_dirs.items():\n",
    "    temp_hook_fn = partial(flip_hook, flip_dir=flip_dir)\n",
    "    hook_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "    probe_data.model.add_hook(hook_name, temp_hook_fn)\n",
    "_, modified_cache = probe_data.model.run_with_cache(board_seqs_int_BL.to(DEVICE)[:, :-1])\n",
    "modified_board_model_move = chess_utils.get_model_move(probe_data.model, meta, model_input)\n",
    "probe_data.model.reset_hooks()\n",
    "modified_resid_post = modified_cache[\"resid_post\", layer][:, :]\n",
    "\n",
    "\n",
    "print(modified_board_model_move)\n",
    "# modified_board_model_move_san = modified_board.parse_san(modified_board_model_move)\n",
    "# assert modified_board_model_move_san in modified_board.legal_moves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(flip_dirs[layer].shape)\n",
    "print(resid_post_BlD.shape)\n",
    "print(modified_resid_post.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexed_modified_resid_posts = []\n",
    "\n",
    "for batch_idx in range(white_move_indices_BL.size(0)):\n",
    "    dots_indices_for_batch_L = white_move_indices_BL[batch_idx]\n",
    "    indexed_modified_resid_post = modified_resid_post[batch_idx, dots_indices_for_batch_L]\n",
    "    indexed_modified_resid_posts.append(indexed_modified_resid_post)\n",
    "\n",
    "# Stack the indexed state stacks along the first dimension\n",
    "stacked_modified_resid_post = torch.stack(indexed_modified_resid_posts)\n",
    "stacked_modified_resid_post = stacked_modified_resid_post.to(DEVICE)\n",
    "\n",
    "assert stacked_modified_resid_post.shape == (sample_size, num_white_moves, linear_probe_MDRRC.shape[1])\n",
    "\n",
    "modified_probe_out = einsum(\n",
    "    \"batch pos d_model, modes d_model rows cols options -> modes batch pos rows cols options\",\n",
    "    stacked_modified_resid_post,\n",
    "    linear_probe_MDRRC,\n",
    ")\n",
    "modified_state_stacks_probe_outputs = chess_utils.one_hot_to_state_stack(modified_probe_out, config.min_val)\n",
    "modified_state_stacks_probe_outputs = torch.tensor(modified_state_stacks_probe_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can graph the original and modified board states and probe outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotly.subplots import make_subplots\n",
    "\n",
    "move_of_interest_probe_out = probe_out_MBLRRC[0][0][move_of_interest]\n",
    "move_of_interest_probe_out_modified = modified_probe_out[0][0][move_of_interest]\n",
    "print(move_of_interest_probe_out.shape)\n",
    "\n",
    "fig_rows = 6\n",
    "fig_cols = 3\n",
    "fig = make_subplots(rows=fig_rows, cols=fig_cols, subplot_titles=[\n",
    "    \"Chess board blank squares\", \"Probe output blank squares clip=2\", \"Probe output blank squares no clipping\",\n",
    "    \"Chess board original piece\", \"Probe output original piece clip=5\", \"Probe output original piece no clipping\",\n",
    "    \"Modified chess board blank squares\", \"Probe output blank squares clip=2\", \"Probe output blank squares no clipping\",\n",
    "    \"Modified chess board original piece\", \"Probe output original piece clip=5\", \"Probe output original piece no clipping\",\n",
    "    \"Chess board state\", \"Probe output board state\", \"Redundant probe output board state\",\n",
    "    \"Modified chess board state\", \"Probe output board state\", \"Redundant probe output board state\"\n",
    "])\n",
    "\n",
    "\n",
    "# Specify the size of each plot\n",
    "plot_size = 400  # You can adjust this size\n",
    "\n",
    "\n",
    "\n",
    "fig.add_trace(plot_board_state(move_of_interest_state_one_hot_RRC[:, :, BLANK_INDEX]), row=1, col=1)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, BLANK_INDEX], clip_size=2), row=1, col=2)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, BLANK_INDEX]), row=1, col=3)\n",
    "\n",
    "fig.add_trace(plot_board_state(move_of_interest_state_one_hot_RRC[:, :, moved_piece_probe_index]), row=2, col=1)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, moved_piece_probe_index], clip_size=5), row=2, col=2)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out[:, :, moved_piece_probe_index]), row=2, col=3)\n",
    "\n",
    "fig.add_trace(plot_board_state(modified_move_of_interest_state_one_hot[:, :, BLANK_INDEX]), row=3, col=1)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out_modified[:, :, BLANK_INDEX], clip_size=2), row=3, col=2)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out_modified[:, :, BLANK_INDEX]), row=3, col=3)\n",
    "\n",
    "fig.add_trace(plot_board_state(modified_move_of_interest_state_one_hot[:, :, moved_piece_probe_index]), row=4, col=1)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out_modified[:, :, moved_piece_probe_index], clip_size=5), row=4, col=2)\n",
    "fig.add_trace(plot_board_state(move_of_interest_probe_out_modified[:, :, moved_piece_probe_index]), row=4, col=3)\n",
    "\n",
    "fig.add_trace(plot_board_state_with_text(move_of_interest_state_RR), row=5, col=1)\n",
    "fig.add_trace(plot_board_state_with_text(state_stacks_probe_outputs_MBLRR[0][0][move_of_interest]), row=5, col=2)\n",
    "fig.add_trace(plot_board_state_with_text(state_stacks_probe_outputs_MBLRR[0][0][move_of_interest]), row=5, col=2)\n",
    "\n",
    "fig.add_trace(plot_board_state_with_text(modified_move_of_interest_state), row=6, col=1)\n",
    "fig.add_trace(plot_board_state_with_text(modified_state_stacks_probe_outputs[0][0][move_of_interest]), row=6, col=2)\n",
    "fig.add_trace(plot_board_state_with_text(modified_state_stacks_probe_outputs[0][0][move_of_interest]), row=6, col=2)\n",
    "\n",
    "# Adjust the overall size of the figure\n",
    "fig.update_layout(height=fig_rows * plot_size, width=fig_cols * plot_size)\n",
    "\n",
    "# Show the figure\n",
    "fig.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "othello",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}