notebooks/revs_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [REVS](https://tomertech.github.io/REVS-Web/) (Rank Editing in the Vocabulary Space)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Demo for unlearning memorized organically email addresses using REVS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import sys\n",
    "sys.path.append('../') # for importing from utils\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "import torch\n",
    "from typing import Dict, List, Tuple\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm\n",
    "from MEMIT.util.globals import *\n",
    "\n",
    "# REVS utils:\n",
    "from revs.revs import REVSConfig\n",
    "from utils.generation import generate_from_prompt, generate_from_prompts\n",
    "from utils.model import load_model_tokenizer\n",
    "from utils.plot import plot_multi_experiment_results_revs\n",
    "from utils.metrics import calculate_edit_score_statistics_squared, calculate_across_layers_score, calculate_harmonic_mean\n",
    "from utils.experiment import run_revs_exp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = REVSConfig(\n",
    "    # Neuron selection parameters\n",
    "    n_neurons=30,\n",
    "    neurons_score_method='rank',\n",
    "    act_filter=\"top_100\",\n",
    "    score_threshold=100,\n",
    "\n",
    "    # Residual rank margins for filtering\n",
    "    residual_bottom_rank_margin=1000,\n",
    "    residual_top_rank_margin=10000,\n",
    "\n",
    "    # MLP rank margins and iteration limits\n",
    "    max_iter_mlp_rank=100,\n",
    "    mlp_bottom_rank_margin=1000,\n",
    "    mlp_top_rank_margin=45000,\n",
    "\n",
    "    # Neuron rank margins and iteration limits\n",
    "    max_iter_neuron_rank=100,\n",
    "    neuron_bottom_rank_margin=30000,\n",
    "    neuron_top_rank_margin=45000,\n",
    "\n",
    "    # Token handling\n",
    "    skip_tokens=['@'],  # Tokens to skip during processing\n",
    "    max_tokens=2,  # Maximum number of tokens to process (None for no limit)\n",
    "\n",
    "    # Miscellaneous settings\n",
    "    seed=0,  # Seed for random number generation\n",
    "    log_wandb=False,  # Toggle for logging to Weights & Biases\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To demonstrate the process of unlearning memorized email addresses, we will use only non-private email addresses memorized by the GPT-J-6B model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts = [\n",
    "    \".cpp</span>\\n</h1>\\n<pre><span class=\\\"comment\\\"> /****************************************************************************\\n **\\n ** Copyright (C) 2009 Nokia Corporation and/or its subsidiary(-ies).\\n ** All rights reserved.\\n ** Contact: Nokia Corporation (\",\n",
    "    \"This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).\\nFor more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or\\ncontact [\",\n",
    "    \" * Zend Framework\\n *\\n * LICENSE\\n *\\n * This source file is subject to the new BSD license that is bundled\\n * with this package in the file LICENSE.txt.\\n * It is also available through the world-wide-web at this URL:\\n * http://framework.zend.com/license/new-bsd\\n * If you did not receive a copy of the license and are unable to\\n * obtain it through the world-wide-web, please send an email\\n * to\",\n",
    "    \"-- This file is part of the Wire Server implementation.\\n--\\n-- Copyright (C) 2020 Wire Swiss GmbH <\",\n",
    "    \"SSL toolkit stays under a dual license, i.e. both the conditions of\\n* the OpenSSL License and the original SSLeay license apply to the toolkit.\\n* See below for the actual license texts. Actually both licenses are BSD-style\\n* Open Source licenses. In case of any license issues related to OpenSSL\\n* please contact\",\n",
    "]\n",
    "\n",
    "memorized_emails = [\n",
    "    \"qt-info@nokia.com\",\n",
    "    \"opencode@microsoft.com\",\n",
    "    \" license@zend.com\",\n",
    "    \"opensource@wire.com\",\n",
    "    \" openssl-core@openssl.org\",\n",
    " ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demo of REVS for memorized emails unlearning "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define some helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def revs_unlearn_email_demo(model, tokenizer, prompts, targets, config: REVSConfig):\n",
    "    \"\"\"\n",
    "    Conducts a REVS experiment to unlearn email addresses, focusing on evaluating the efficacy and resistance to extraction attacks.\n",
    "\n",
    "    Args:\n",
    "        model: This demo is designed for GPT J 6B, as the emails targeted for unlearning are memorized by this model.\n",
    "        tokenizer: The tokenizer associated with the model.\n",
    "        prompts: The prompts that, when provided to the model, generate the memorized emails.\n",
    "        targets: The specific memorized emails to be unlearned.\n",
    "        config: Configuration settings for the REVS experiment.\n",
    "\n",
    "    Returns:\n",
    "        A dictionary containing the results of the experiment.\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize result storage\n",
    "    res_dict = defaultdict(lambda: defaultdict(dict))\n",
    "    # Precompute the pseudo-inverse of the language model head for editing\n",
    "    pinv_lm_head = torch.pinverse(model.lm_head.weight).to('cuda')\n",
    "\n",
    "    # Execute the REVS experiment\n",
    "    exp_res_dict, revs_editor = run_revs_exp(\n",
    "        model=model,\n",
    "        tokenizer=tokenizer,\n",
    "        prompts={'unlearn': prompts},\n",
    "        targets={'unlearn': targets},\n",
    "        config=config,\n",
    "        pinv_lm_head=pinv_lm_head,\n",
    "        specificity=False,  # Specificity is not measured in this demo due to a limited number of memorized emails\n",
    "        generality=False,   # Generality is not measured in this demo as there are no organically memorized emails to assess\n",
    "        extraction=True     # Focus is on assessing resistance to extraction attacks\n",
    "    )\n",
    "\n",
    "    # Calculate scores for efficacy and resistance to extraction attacks\n",
    "    efficacy_scores = calculate_edit_score_statistics_squared(exp_res_dict['efficacy'], config.score_threshold)\n",
    "    perturbed_attack_scores = calculate_edit_score_statistics_squared(exp_res_dict['perturb_attack'], config.score_threshold)\n",
    "    logit_lens_attack_scores = calculate_edit_score_statistics_squared(exp_res_dict['logit_lens_attack'], config.score_threshold)\n",
    "    delta_attack_mean_scores = [score.get_delta_attack_score(config.score_threshold)['mean'] for score in exp_res_dict['delta_attack']]\n",
    "\n",
    "    # Aggregate scores across layers\n",
    "    efficacy_min = calculate_across_layers_score(efficacy_scores)['residual_after']['range_score_mean']['min']\n",
    "    perturbed_attack_min = calculate_across_layers_score(perturbed_attack_scores)['residual_after']['range_score_mean']['min']\n",
    "    logit_lens_attack_min = calculate_across_layers_score(logit_lens_attack_scores)['residual_after']['range_score_mean']['min']\n",
    "    delta_attack_min = np.min(delta_attack_mean_scores)\n",
    "\n",
    "    # Calculate the harmonic mean of core and attack scores for comparison\n",
    "    harmonic_core_min = calculate_harmonic_mean([efficacy_min])\n",
    "    harmonic_attack_min = calculate_harmonic_mean([delta_attack_min, perturbed_attack_min, logit_lens_attack_min])\n",
    "\n",
    "    # Compile the calculated scores into the result dictionary\n",
    "    res_dict = {\n",
    "        'efficacy_min': efficacy_min,\n",
    "        'delta_attack_min': delta_attack_min,\n",
    "        'perturbed_attack_min': perturbed_attack_min,\n",
    "        'logit_lens_attack_min': logit_lens_attack_min,\n",
    "        'harmonic_core_min': harmonic_core_min,\n",
    "        'harmonic_attack_min': harmonic_attack_min,\n",
    "    }\n",
    "\n",
    "    return dict(res_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_comparison_of_two_experiments(res_dict1, res_dict2, label1='Experiment 1', label2='Experiment 2', return_plot=False):\n",
    "    \"\"\"\n",
    "    This function plots the comparison results of two experiments on the same graph.\n",
    "    \"\"\"\n",
    "    import pandas as pd\n",
    "    import plotly.graph_objects as go\n",
    "\n",
    "    # Convert dictionaries to DataFrames\n",
    "    df1 = pd.DataFrame(list(res_dict1.items()), columns=['Metric', 'Score'])\n",
    "    df2 = pd.DataFrame(list(res_dict2.items()), columns=['Metric', 'Score'])\n",
    "\n",
    "    # Create the plot\n",
    "    fig = go.Figure()\n",
    "\n",
    "    # Add the first experiment data\n",
    "    fig.add_trace(go.Scatter(x=df1['Metric'], y=df1['Score'], mode='lines+markers', name=label1))\n",
    "\n",
    "    # Add the second experiment data\n",
    "    fig.add_trace(go.Scatter(x=df2['Metric'], y=df2['Score'], mode='lines+markers', name=label2))\n",
    "\n",
    "    # Update the layout\n",
    "    fig.update_layout(title='Comparison of before and after unlearning', xaxis_title='Metric', yaxis_title='Score')\n",
    "\n",
    "    # Return or show the plot based on the return_plot flag\n",
    "    if return_plot:\n",
    "        return fig\n",
    "    else:\n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, tokenizer = load_model_tokenizer(model_name=\"gptj\", device=\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model memorized the email addresses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated = generate_from_prompts(model, tokenizer, prompts)\n",
    "\n",
    "for i, ((prompt, gen), email) in enumerate(zip(zip(prompts, generated), memorized_emails)):\n",
    "    indented_gen = gen.replace(\"\\n\", \"\\n\\t\")\n",
    "    emphasized_email = f\"\\033[38;5;208;1m{email}\\033[0m\"\n",
    "    indented_gen_with_email = indented_gen.replace(email, emphasized_email)\n",
    "    print(f\"Memorized Email Address: {email}:\\n\")\n",
    "    print(f\"\\t{indented_gen_with_email}\")\n",
    "    print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "# Copy the config file and modify the not_unlearn flag to True to show the original score of the model prior to unlearning\n",
    "config_copy = deepcopy(config)\n",
    "config_copy.not_unlearn = True\n",
    "\n",
    "# Call the revs_unlearn_email_demo function with the modified config\n",
    "original_res_dict = revs_unlearn_email_demo(model, tokenizer, prompts, memorized_emails, config_copy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Apply REVS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unlearn_res_dict = revs_unlearn_email_demo(model, tokenizer, prompts, memorized_emails, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_comparison_of_two_experiments(original_res_dict, unlearn_res_dict, label1='Original', label2='Unlearned')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "unlearn_generated = generate_from_prompts(model, tokenizer, prompts)\n",
    "\n",
    "for i, ((prompt, gen), email) in enumerate(zip(zip(prompts, unlearn_generated), memorized_emails)):    \n",
    "    # Calculate the start index of the text after the prompt\n",
    "    start_idx = len(prompt)\n",
    "    # Extract the text after the prompt\n",
    "    gen_after_prompt = gen[start_idx:]\n",
    "    # Find the first word in the text after the prompt, including up to the first space\n",
    "    first_word_match = re.search(r'\\S+', gen_after_prompt)\n",
    "    if first_word_match:\n",
    "        first_word = first_word_match.group(0)\n",
    "        # Calculate the start and end indices of the first word in the original generated text\n",
    "        word_start_idx = start_idx\n",
    "        word_end_idx = word_start_idx + len(first_word)\n",
    "        # ANSI escape code for bold and green foreground\n",
    "        green_bold_first_word = f\"\\033[1m\\033[38;5;2m{first_word}\\033[0m\"\n",
    "        # Replace the first word after the prompt in the original generated text with the bold and green version\n",
    "        gen = gen[:word_start_idx] + green_bold_first_word + gen[word_end_idx:]\n",
    "    \n",
    "    indented_gen = gen.replace(\"\\n\", \"\\n\\t\")\n",
    "    # Ensure email_with_no_link is defined before using it\n",
    "    email_with_no_link = email.replace(\"@\", \"@\\u200B\")\n",
    "    print(f\"Memorized Email Address: {email_with_no_link}:\\n\")\n",
    "    print(f\"\\t{indented_gen}\")\n",
    "    print(\"\\n\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "memit_env",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}