minimal_demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPhOpqO99NM8JWiEoomHyR2"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "<div text-align=\"center\">\n",
        "  <img src=\"https://raw.githubusercontent.com/FarnoushRJ/MambaLRP/main/assets/MambaLRP_logo.jpeg\" width=\"1000\"/>\n",
        "</div>\n",
        "\n",
        "\n",
        "<div text-align=\"center\"><h1>🐍 MambaLRP is here! 🎉</h1>"
      ],
      "metadata": {
        "id": "cNHDx_0iTsiN"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Clone the repository and install MambaLRP."
      ],
      "metadata": {
        "id": "w-_UXcPNT4sw"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "ZY_WrV6xumJ9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7b502f10-3b50-4f39-cf9c-799c79d955d7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "fatal: destination path 'MambaLRP' already exists and is not an empty directory.\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/FarnoushRJ/MambaLRP.git\n",
        "!pip install git+file:///content/MambaLRP --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import necessary packages."
      ],
      "metadata": {
        "id": "jJAFKD_kUBge"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer\n",
        "import sys\n",
        "\n",
        "from mamba_lrp.model.mamba_huggingface import ModifiedMambaForCausalLM\n",
        "from mamba_lrp.model.utils import *\n",
        "from mamba_lrp.lrp.utils import relevance_propagation\n",
        "from mamba_lrp.dataset.general_dataset import get_sst_dataset\n",
        "import torch\n",
        "import numpy as np"
      ],
      "metadata": {
        "id": "FfPmMtMBS2KN"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load model\n",
        "\n",
        "Load model and tokenizer."
      ],
      "metadata": {
        "id": "USC37ScFs_QG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install gdown\n",
        "\n",
        "# Import gdown\n",
        "import gdown\n",
        "\n",
        "# Define the file ID and the destination file name\n",
        "file_id = '1RnIygUDodGeKPqbcEQTOYR5dztpF6X1b'  # Replace with your actual file ID\n",
        "destination = 'mamba_sst2_weights.pt'  # Desired output file name\n",
        "\n",
        "# Construct the URL\n",
        "url = f'https://drive.google.com/uc?id={file_id}'\n",
        "\n",
        "# Download the file\n",
        "gdown.download(url, destination, quiet=False)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 313
        },
        "id": "sIS7FLBuSB7L",
        "outputId": "eee06284-e85b-4abb-bfbf-2d048590e4ec"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: gdown in /usr/local/lib/python3.10/dist-packages (5.1.0)\n",
            "Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.10/dist-packages (from gdown) (4.12.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from gdown) (3.15.4)\n",
            "Requirement already satisfied: requests[socks] in /usr/local/lib/python3.10/dist-packages (from gdown) (2.31.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from gdown) (4.66.4)\n",
            "Requirement already satisfied: soupsieve>1.2 in /usr/local/lib/python3.10/dist-packages (from beautifulsoup4->gdown) (2.5)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests[socks]->gdown) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests[socks]->gdown) (3.7)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests[socks]->gdown) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests[socks]->gdown) (2024.7.4)\n",
            "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /usr/local/lib/python3.10/dist-packages (from requests[socks]->gdown) (1.7.1)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading...\n",
            "From (original): https://drive.google.com/uc?id=1RnIygUDodGeKPqbcEQTOYR5dztpF6X1b\n",
            "From (redirected): https://drive.google.com/uc?id=1RnIygUDodGeKPqbcEQTOYR5dztpF6X1b&confirm=t&uuid=a7a6744f-6be8-4f33-9330-fd6d45f0395b\n",
            "To: /content/mamba_sst2_weights.pt\n",
            "100%|██████████| 517M/517M [00:07<00:00, 72.4MB/s]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'mamba_sst2_weights.pt'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Load tokenizer.\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"state-spaces/mamba-130m-hf\")\n",
        "tokenizer.eos_token = \"<|endoftext|>\"\n",
        "tokenizer.bos_token = \"<|startoftext|>\"\n",
        "tokenizer.pad_token = \"<|pad|>\"\n",
        "tokenizer.unk_token = \"<|unkown|>\"\n",
        "tokenizer.add_tokens(['<|unkown|>', '<|pad|>', \"<|startoftext|>\"], special_tokens=True)\n",
        "\n",
        "# Load model.\n",
        "model = MambaForCausalLM.from_pretrained(\"state-spaces/mamba-130m-hf\", use_cache=True)\n",
        "resize_token_embeddings(model, len(tokenizer))\n",
        "model.lm_head = torch.nn.Linear(768, 2, bias=True)\n",
        "\n",
        "# Load the model's weights\n",
        "model.load_state_dict(\n",
        "    torch.load('mamba_sst2_weights.pt', map_location=torch.device('cpu')),\n",
        "    strict=True\n",
        ")\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "model.to(device)\n",
        "model.eval()\n",
        "\n",
        "# Make model explainable.\n",
        "modified_model = ModifiedMambaForCausalLM(model, is_fast_forward_available=False)\n",
        "modified_model.eval()\n",
        "model.backbone.embeddings.requires_grad = False\n",
        "pretrained_embeddings = model.backbone.embeddings"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YanGwXpYs6oD",
        "outputId": "ab5910b4-f209-40d6-f313-ce3682223728"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
            "  warnings.warn(\n",
            "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load dataset\n",
        "\n",
        "Load SST-2 dataset."
      ],
      "metadata": {
        "id": "vz0zFnHltoZe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "validation_dataset = get_sst_dataset(\n",
        "    tokenizer=tokenizer,\n",
        "    truncation=False,\n",
        "    max_length=None\n",
        "    )"
      ],
      "metadata": {
        "id": "pGfcKo5Btnoc"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generate explanation\n",
        "\n",
        "Generate explanation for one sample."
      ],
      "metadata": {
        "id": "UYWbXi1173O1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "i = 413\n",
        "input_ids = validation_dataset.__getitem__(i)['input_ids'].unsqueeze(0).to(device)\n",
        "label = torch.tensor(validation_dataset.__getitem__(i)['label']).long().to(device)\n",
        "idx = torch.where(input_ids == 0)[1] + 1\n",
        "input_ids = input_ids[:, :idx]\n",
        "embeddings = pretrained_embeddings(input_ids)\n",
        "\n",
        "R, prediction = relevance_propagation(\n",
        "    model=modified_model,\n",
        "    embeddings=embeddings,\n",
        "    targets=label,\n",
        "    n_classes=2\n",
        "    )"
      ],
      "metadata": {
        "id": "c_fWPHHrt1RS"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Visualization\n",
        "\n",
        "For simplicity, we use the visualization utilities in Captum to display the results."
      ],
      "metadata": {
        "id": "P2aHZLM48BJy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from captum.attr import visualization as viz"
      ],
      "metadata": {
        "id": "fvBR6Yse8Hbx"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokens = []\n",
        "for id in input_ids[0][1: -2]:\n",
        "    tokens.append(tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens([id.item()])))\n",
        "attributions = R[0][1: -2]\n",
        "attributions = attributions / attributions.max()"
      ],
      "metadata": {
        "id": "jndIdPgJwg2n"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Visualize the attributions\n",
        "viz.visualize_text([viz.VisualizationDataRecord(\n",
        "    attributions,\n",
        "    torch.max(model(input_ids).logits[:, -1, :], dim=1).values.item(),\n",
        "    torch.argmax(model(input_ids).logits[:, -1, :], dim=1).item(),\n",
        "    true_class=label.item(),\n",
        "    attr_class=label.item(),\n",
        "    attr_score=attributions.sum(),\n",
        "    raw_input_ids=tokens,\n",
        "    convergence_score=None\n",
        ")])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 181
        },
        "id": "1YRlZ6c1yUtN",
        "outputId": "a193a35e-4e95-4dd4-e2d5-933c32584971"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<table width: 100%><div style=\"border-top: 1px solid; margin-top: 5px;             padding-top: 5px; display: inline-block\"><b>Legend: </b><span style=\"display: inline-block; width: 10px; height: 10px;                 border: 1px solid; background-color:                 hsl(0, 75%, 60%)\"></span> Negative  <span style=\"display: inline-block; width: 10px; height: 10px;                 border: 1px solid; background-color:                 hsl(0, 75%, 100%)\"></span> Neutral  <span style=\"display: inline-block; width: 10px; height: 10px;                 border: 1px solid; background-color:                 hsl(120, 75%, 50%)\"></span> Positive  </div><tr><th>True Label</th><th>Predicted Label</th><th>Attribution Label</th><th>Attribution Score</th><th>Word Importance</th><tr><td><text style=\"padding-right:2em\"><b>0</b></text></td><td><text style=\"padding-right:2em\"><b>0 (2.86)</b></text></td><td><text style=\"padding-right:2em\"><b>0</b></text></td><td><text style=\"padding-right:2em\"><b>1.67</b></text></td><td><mark style=\"background-color: hsl(0, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\"> at                    </font></mark><mark style=\"background-color: hsl(0, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  least                    </font></mark><mark style=\"background-color: hsl(0, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  one                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  scene                    </font></mark><mark style=\"background-color: hsl(120, 75%, 95%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  is                    </font></mark><mark style=\"background-color: hsl(120, 75%, 90%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  so                    </font></mark><mark style=\"background-color: hsl(120, 75%, 50%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  disgusting                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  that                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  viewers                    </font></mark><mark style=\"background-color: hsl(0, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  may                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  be                    </font></mark><mark style=\"background-color: hsl(120, 75%, 98%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  hard                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  pressed                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  to                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  retain                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  their                    </font></mark><mark style=\"background-color: hsl(120, 75%, 97%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  lunch                    </font></mark><mark style=\"background-color: hsl(120, 75%, 92%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  .                    </font></mark></td><tr></table>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<table width: 100%><div style=\"border-top: 1px solid; margin-top: 5px;             padding-top: 5px; display: inline-block\"><b>Legend: </b><span style=\"display: inline-block; width: 10px; height: 10px;                 border: 1px solid; background-color:                 hsl(0, 75%, 60%)\"></span> Negative  <span style=\"display: inline-block; width: 10px; height: 10px;                 border: 1px solid; background-color:                 hsl(0, 75%, 100%)\"></span> Neutral  <span style=\"display: inline-block; width: 10px; height: 10px;                 border: 1px solid; background-color:                 hsl(120, 75%, 50%)\"></span> Positive  </div><tr><th>True Label</th><th>Predicted Label</th><th>Attribution Label</th><th>Attribution Score</th><th>Word Importance</th><tr><td><text style=\"padding-right:2em\"><b>0</b></text></td><td><text style=\"padding-right:2em\"><b>0 (2.86)</b></text></td><td><text style=\"padding-right:2em\"><b>0</b></text></td><td><text style=\"padding-right:2em\"><b>1.67</b></text></td><td><mark style=\"background-color: hsl(0, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\"> at                    </font></mark><mark style=\"background-color: hsl(0, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  least                    </font></mark><mark style=\"background-color: hsl(0, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  one                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  scene                    </font></mark><mark style=\"background-color: hsl(120, 75%, 95%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  is                    </font></mark><mark style=\"background-color: hsl(120, 75%, 90%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  so                    </font></mark><mark style=\"background-color: hsl(120, 75%, 50%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  disgusting                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  that                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  viewers                    </font></mark><mark style=\"background-color: hsl(0, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  may                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  be                    </font></mark><mark style=\"background-color: hsl(120, 75%, 98%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  hard                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  pressed                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  to                    </font></mark><mark style=\"background-color: hsl(120, 75%, 99%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  retain                    </font></mark><mark style=\"background-color: hsl(120, 75%, 100%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  their                    </font></mark><mark style=\"background-color: hsl(120, 75%, 97%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  lunch                    </font></mark><mark style=\"background-color: hsl(120, 75%, 92%); opacity:1.0;                     line-height:1.75\"><font color=\"black\">  .                    </font></mark></td><tr></table>"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    }
  ]
}