{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# How to plot Transformers's Attention heads?\n\nThis is a simple notebook showing error analysis for transformer based models. \nIt loades CodeBERT (Can be changed easily. It was the main model used on AI4Code) and run a sample from the training set through it. Then it plots the attention maps generated by the inner transformer blocks and the outputed logits of any layer you require. Then you can use the displayed data to debug you model further. \n\nThis approach can be useful for debugging your models since if used on a case where the model is mistaken, it shows you why was it wrong. (What tokens in the sample made it flip and why).\n\nEnjoy! ","metadata":{"papermill":{"duration":0.031568,"end_time":"2022-05-12T10:15:13.890382","exception":false,"start_time":"2022-05-12T10:15:13.858814","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"**The actual code is hidden here 👇**","metadata":{}},{"cell_type":"code","source":"!pip install ansi\nfrom typing import Union\n\nimport numpy as np\nimport seaborn as sns\nfrom ansi.colour import rgb\n\n\ndef color_text(text, rgb_code):\n    reset =  '\\x1b[0m'\n    return rgb.rgb256(*rgb_code) + text + reset\n\ndef value2rgb(value):\n#     if value < 0:\n#         rgb_code = (255/2 + abs(value)/2, abs(value), 255/2 + abs(value)/2)\n#     else:\n#         rgb_code = (125+value/2, 0, 255/2-value/2)\n    if value < 0:\n        rgb_code = (255, 255, abs(value))\n    else:\n        rgb_code = (255, 255-value, 0)\n    return rgb_code\n\n\ndef scale(values, input_range, output_range):\n    return np.interp(values, input_range, output_range)\n\n\ndef get_legends(value_range, scale_to, step=5):\n    min_value, max_value = value_range\n    leg_values = np.linspace(min_value, max_value, step)\n    scaled_values = scale(leg_values, (min_value, max_value), scale_to)\n    \n    legends = []\n    for leg_value, scaled_value in zip(leg_values, scaled_values):\n         legends.append(color_text('{:.2f}'.format(leg_value), value2rgb(scaled_value)))\n    return legends\n\n\ndef color_texts(texts, values, use_absolute):\n    if use_absolute:\n        value_range = (0, 1)\n    else:\n        value_range = (min(values), max(values))\n    scale_to = (-255, 255)\n    scaled_values = scale(values, value_range, scale_to)\n    result = []\n    for text, value in zip(texts, scaled_values):\n        rgb = value2rgb(value)\n        result.append(color_text(text, rgb))\n       \n    \n    colored = ' '.join(result)\n    legends = get_legends(value_range, scale_to)\n\n    colored += ' ({})'.format(' '.join(legends))\n        \n    if use_absolute:\n        colored += ' (min: {:.10f} max: {:.10f})'.format(min(values), max(values))\n    \n    return colored\n\n\ndef visual_matrix(matrix, labels=None, title=None, **kwargs):\n\n    sns.set()\n    ax = sns.heatmap(matrix, xticklabels=labels, yticklabels=labels, **kwargs)\n    if title:\n        ax.set(title = title)\n#     ax.xaxis.tick_top()\n\n    return ax\n\n\ndef get_or_default_config(layer_num, batch_num, head_num, token_num, atn_axis, atns):\n    if layer_num is None:\n        layer_num = -1  # last layer\n    \n    batch_size = len(atns[0])\n    if batch_size == 1:\n        batch_num = 0\n    else:\n        if batch_num is None:\n            raise ValueError('You input an attention with batch size != 1. Please input attentions with batch size 1 or specify the batch_num you want to visualize.')\n            \n    if head_num is None:\n        head_num = 'average'\n\n    if token_num is None:\n        token_num = 'average'\n\n    if atn_axis is None:\n        atn_axis = 0\n        \n    return layer_num, batch_num, head_num, token_num, atn_axis\n\n\ndef get_multihead_atn_matrix(atns, layer_num=None, batch_num=None):\n    \n    \n#     layer_num, batch_num = get_or_default_layer_and_batch_num(layer_num, batch_num, atns)\n    \n    layer = atns[layer_num]\n\n    try:\n        multihead_atn_matrix = layer[batch_num].detach().numpy()  # pytorch\n    except TypeError:\n        multihead_atn_matrix = layer[batch_num].cpu().numpy()  # pytorch\n    except AttributeError:\n        multihead_atn_matrix = layer[batch_num]  # tensorflow\n\n    return multihead_atn_matrix\n\n\ndef get_atn_matrix_from_mh_matrix(multihead_atn_matrix, head_num):\n    # atn_matrix: (sequence_length, sequence_length)       \n    try:\n        atn_matrix = multihead_atn_matrix[head_num]\n    except (IndexError, TypeError):\n        # average over heads        \n        atn_matrix = np.mean(multihead_atn_matrix, axis=0)\n\n    return atn_matrix\n\n\ndef merge_atn_matrix(atn_matrix, mean_over_mat_axis):\n    atn_matrix_over_axis: list = np.mean(atn_matrix, axis=mean_over_mat_axis)\n    return atn_matrix_over_axis\n\n\ndef matrix2values(matrix, index='average', axis=0):\n    \n    if index == 'average':\n        result_mat = np.mean(matrix, axis=axis)\n    elif isinstance(index, int):\n        if axis == 0:\n            result_mat = matrix[index]\n        elif axis == 1:\n            result_mat = matrix.T[index]\n        else:\n            raise ValueError('matrix to values have a wrong axis (0 or 1): ' + str(axis))\n    else:\n        raise ValueError('matrix to values have a wrong index (\"average\" or integers): ' + str(index))\n    \n    return result_mat\n        \n\ndef get_atn_values(layer_num, batch_num, head_num, token_num, atn_axis, atns):\n    layer_num, batch_num, head_num, token_num, atn_axis = get_or_default_config(layer_num, batch_num, head_num, token_num, atn_axis, atns)\n    multihead_atn_matrix = get_multihead_atn_matrix(atns, layer_num=layer_num, batch_num=batch_num)\n    atn_matrix = get_atn_matrix_from_mh_matrix(multihead_atn_matrix, head_num=head_num)\n    atn_values = matrix2values(atn_matrix, index=token_num, axis=atn_axis)\n    \n    return atn_values\n\n\ndef get_atn_matrix(layer_num, batch_num, head_num, atns):\n    layer_num, batch_num, head_num, *_ = get_or_default_config(layer_num, batch_num, head_num, None, None, atns)\n\n    multihead_atn_matrix = get_multihead_atn_matrix(atns, layer_num=layer_num, batch_num=batch_num)\n    atn_matrix = get_atn_matrix_from_mh_matrix(multihead_atn_matrix, head_num=head_num)\n    return atn_matrix\n\n\ndef visual_atn(labels, atns, layer_num=None, batch_num=None, head_num=None, token_num=None, atn_axis=None,\n               use_absolute=False, output=False, **kwargs):\n    atn_values = get_atn_values(layer_num, batch_num, head_num, token_num, atn_axis, atns)\n    layer_num, batch_num, head_num, token_num, atn_axis = get_or_default_config(layer_num, batch_num, head_num, token_num, atn_axis, atns)\n\n    assert len(labels) == len(atn_values), 'len(labels): {}, len(merged_atn_values): {}'.format(len(labels), len(atn_values))\n\n    colored = color_texts(labels, atn_values, use_absolute)\n\n    try:\n        label = labels[token_num]\n    except TypeError:\n        label = 'ALL_TOKENS'\n\n    print('(layer) {} (batch) {} (head) {} (token_num) {} (token) {} (axis) {}'.format(layer_num, batch_num, head_num, token_num, label, atn_axis))\n\n    if output:\n        return colored, atn_values\n    else:\n        return colored\n\n    \ndef visual_atn_matrix(labels, atns, layer_num=None, batch_num=None, head_num=None, token_num=None, output=False) -> 'Axes':\n    \n    atn_matrix = get_atn_matrix(layer_num, batch_num, head_num, atns)\n    \n    layer_num, batch_num, head_num, token_num, _ = get_or_default_config(layer_num, batch_num, head_num, token_num, None, atns)\n    \n    title = '(layer) {} (batch) {} (head) {}'.format(layer_num, batch_num, head_num)\n    \n    if output:\n        return visual_matrix(atn_matrix, labels, title=title), atn_matrix\n    else:\n        return visual_matrix(atn_matrix, labels, title=title)\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-07-18T07:48:24.096060Z","iopub.execute_input":"2022-07-18T07:48:24.096408Z","iopub.status.idle":"2022-07-18T07:48:36.650054Z","shell.execute_reply.started":"2022-07-18T07:48:24.096375Z","shell.execute_reply":"2022-07-18T07:48:36.648809Z"},"_kg_hide-output":true,"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom scipy import sparse\nfrom pathlib import Path\npd.options.display.width = 180\npd.options.display.max_colwidth = 120\ndata_dir = Path('../input/AI4Code')\n\nNUM_TRAIN = 10000\n\ndef read_notebook(path):\n    return (\n        pd.read_json(\n            path,\n            dtype={'cell_type': 'category', 'source': 'str'})\n        .assign(id=path.stem)\n        .rename_axis('cell_id')\n    )\npaths_train = list((data_dir / 'train').glob('*.json'))[:NUM_TRAIN]\nnotebooks_train = [read_notebook(path) for path in tqdm(paths_train, desc='Train NBs')]\ndf = (\n    pd.concat(notebooks_train)\n    .set_index('id', append=True)\n    .swaplevel()\n    .sort_index(level='id', sort_remaining=False)\n)\ndf_orders = pd.read_csv(\n    data_dir / 'train_orders.csv',\n    index_col='id',\n    squeeze=True,\n).str.split()\n\ndef get_ranks(base, derived): return [base.index(d) for d in derived]\n\ndf_orders_ = df_orders.to_frame().join(df.reset_index('cell_id').groupby('id')['cell_id'].apply(list), how='right')\n\nranks = {}\nfor id_, cell_order, cell_id in df_orders_.itertuples(): ranks[id_] = {'cell_id': cell_id, 'rank': get_ranks(cell_order, cell_id)}\n\ndf_ranks = (\n    pd.DataFrame\n    .from_dict(ranks, orient='index')\n    .rename_axis('id')\n    .apply(pd.Series.explode)\n    .set_index('cell_id', append=True)\n)\n\ndf_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')\ndf = df.reset_index().merge(df_ranks, on=[\"id\", \"cell_id\"]).merge(df_ancestors, on=[\"id\"])\ndf[\"pct_rank\"] = df[\"rank\"] / df.groupby(\"id\")[\"cell_id\"].transform(\"count\")","metadata":{"papermill":{"duration":82.291505,"end_time":"2022-05-12T10:16:36.365197","exception":false,"start_time":"2022-05-12T10:15:14.073692","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-07-18T07:48:36.664272Z","iopub.execute_input":"2022-07-18T07:48:36.664721Z","iopub.status.idle":"2022-07-18T07:50:00.324342Z","shell.execute_reply.started":"2022-07-18T07:48:36.664676Z","shell.execute_reply":"2022-07-18T07:50:00.323276Z"},"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize BERT's attention heads!","metadata":{}},{"cell_type":"markdown","source":"### How to ask BERT why was it wrong?\n\nIn the domain of transformer-based models, our goal is to build models with high performance for different tasks.\nOne way this goal can be achieved is by **identifying the problems** that are causing our models to be wrong. \nThis is a powerful concept in general that can also aid with EDA (since it highlights parts of the data). \n\nAfter we do this, we can try and guess what should we do with this information so that BERT can understand our data better. This notebook is more interested in the first part. \n\n> **Answer**: **Plot the attention maps** - It is a very simple technique and I will explain everything below. \n\nIn the end, BERT is just a neural network represented by weight tensors. so we can simply plot those weights. What makes it special is that it is a **multi-headed neural network** that allows you to compute any number of attention heads per token.\n\nOr in simple terms: For each token on the sequence, we learn multiple interaction functions that are independent of the interaction functions we learn at the same time. ","metadata":{}},{"cell_type":"markdown","source":"**Loading Our Model**","metadata":{}},{"cell_type":"code","source":"import torch\nfrom transformers import *\nimport matplotlib.pyplot as plt\nfrom transformers import RobertaModel, RobertaTokenizer\n\nplt.figure(figsize = (20, 20))\n\nBERT_PATH = \"../input/codebert-base/codebert-base\"\nmodel = RobertaModel.from_pretrained(BERT_PATH, output_attentions = True)\nprint('Model loaded.')","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:00.327963Z","iopub.execute_input":"2022-07-18T07:50:00.328339Z","iopub.status.idle":"2022-07-18T07:50:02.067260Z","shell.execute_reply.started":"2022-07-18T07:50:00.328283Z","shell.execute_reply":"2022-07-18T07:50:02.065863Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Extracting the attention map out of the model\n\nHere we extract the actual attention map of our model for a given input.\n\n**For example:** Specifically here on this cell it is doing so for the first text sample in the data.","metadata":{}},{"cell_type":"code","source":"source_code = df.iloc[0]['source'][500:600]\n\ntokenizer = RobertaTokenizer.from_pretrained(BERT_PATH)\ninput_ids = [tokenizer.encode(source_code)]\ntokens = [tokenizer.decode(id_) for id_ in input_ids[0]]\n\noutputs = model(torch.tensor(input_ids))\nloss, logits, attentions = outputs\nattentions = outputs['attentions']\nprint('Attention extracted.')","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:28.901065Z","iopub.execute_input":"2022-07-18T07:50:28.901704Z","iopub.status.idle":"2022-07-18T07:50:29.183155Z","shell.execute_reply.started":"2022-07-18T07:50:28.901669Z","shell.execute_reply":"2022-07-18T07:50:29.180917Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Plotting The Attention","metadata":{}},{"cell_type":"code","source":"print(visual_atn(tokens, attentions))\nvisual_atn_matrix(tokens, attentions)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:30.782691Z","iopub.execute_input":"2022-07-18T07:50:30.783060Z","iopub.status.idle":"2022-07-18T07:50:31.651516Z","shell.execute_reply.started":"2022-07-18T07:50:30.783022Z","shell.execute_reply":"2022-07-18T07:50:31.650490Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Cool, right? \n\nLook how different our model's pay attention to the ` ` character and the `< s >` token.\nWhich totaly make sense..","metadata":{}},{"cell_type":"markdown","source":"### Plotting attention maps of different layers\n\nNow let's plot the attention maps from other locations in out model. \n\n- Here are some examples: ","metadata":{}},{"cell_type":"code","source":"print(visual_atn(tokens, attentions, layer_num=-1, head_num='average', token_num='average', atn_axis=0))\nvisual_atn_matrix(tokens, attentions)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:33.489840Z","iopub.execute_input":"2022-07-18T07:50:33.490153Z","iopub.status.idle":"2022-07-18T07:50:34.327849Z","shell.execute_reply.started":"2022-07-18T07:50:33.490123Z","shell.execute_reply":"2022-07-18T07:50:34.326604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(visual_atn(tokens, attentions, use_absolute=True))\nvisual_atn_matrix(tokens, attentions)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:35.429849Z","iopub.execute_input":"2022-07-18T07:50:35.430192Z","iopub.status.idle":"2022-07-18T07:50:36.261208Z","shell.execute_reply.started":"2022-07-18T07:50:35.430161Z","shell.execute_reply":"2022-07-18T07:50:36.260155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(visual_atn(tokens, attentions, layer_num=3, head_num=-1, token_num=4))  # third layer, last head, fourth token (see)\nvisual_atn_matrix(tokens, attentions)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:37.258180Z","iopub.execute_input":"2022-07-18T07:50:37.258531Z","iopub.status.idle":"2022-07-18T07:50:38.523319Z","shell.execute_reply.started":"2022-07-18T07:50:37.258498Z","shell.execute_reply":"2022-07-18T07:50:38.522392Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(visual_atn(tokens, attentions, layer_num=-1, head_num='average', token_num=0, atn_axis=0))\nvisual_atn_matrix(tokens, attentions)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:39.596959Z","iopub.execute_input":"2022-07-18T07:50:39.597296Z","iopub.status.idle":"2022-07-18T07:50:40.434683Z","shell.execute_reply.started":"2022-07-18T07:50:39.597255Z","shell.execute_reply":"2022-07-18T07:50:40.433486Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(visual_atn(tokens, attentions, layer_num=-1, head_num='average', token_num=0, atn_axis=1))\nvisual_atn_matrix(tokens, attentions)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:41.529066Z","iopub.execute_input":"2022-07-18T07:50:41.529701Z","iopub.status.idle":"2022-07-18T07:50:42.364556Z","shell.execute_reply.started":"2022-07-18T07:50:41.529666Z","shell.execute_reply":"2022-07-18T07:50:42.363507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"And for the last layer:","metadata":{}},{"cell_type":"code","source":"visual_atn_matrix(tokens, attentions, layer_num=-1, head_num='average')  # last layer, average over multi-head attention matrices\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:45.000899Z","iopub.execute_input":"2022-07-18T07:50:45.001238Z","iopub.status.idle":"2022-07-18T07:50:45.841057Z","shell.execute_reply.started":"2022-07-18T07:50:45.001206Z","shell.execute_reply":"2022-07-18T07:50:45.840007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Showing every head of a given layer\n\nLet's now plot **every** head of a given attention in our model.\nOn this demo, we are going to do this for the last layer.","metadata":{}},{"cell_type":"code","source":"for i in range(12):\n    visual_atn_matrix(tokens, attentions, layer_num=-1, head_num=i)  # print attention matrix of every head in the last layer\n    plt.show()        ","metadata":{"execution":{"iopub.status.busy":"2022-07-18T07:50:47.681270Z","iopub.execute_input":"2022-07-18T07:50:47.682252Z","iopub.status.idle":"2022-07-18T07:50:58.709221Z","shell.execute_reply.started":"2022-07-18T07:50:47.682197Z","shell.execute_reply":"2022-07-18T07:50:58.707913Z"},"trusted":true},"execution_count":null,"outputs":[]}]}