fgsm_llm_attack.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FGSM Attack Implementation\n",
    "First, let's install the required packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: transformers in ./venv/lib/python3.11/site-packages (4.48.2)\n",
      "Requirement already satisfied: torch in ./venv/lib/python3.11/site-packages (2.6.0)\n",
      "Requirement already satisfied: datasets in ./venv/lib/python3.11/site-packages (3.2.0)\n",
      "Requirement already satisfied: filelock in ./venv/lib/python3.11/site-packages (from transformers) (3.17.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.24.0 in ./venv/lib/python3.11/site-packages (from transformers) (0.28.1)\n",
      "Requirement already satisfied: numpy>=1.17 in ./venv/lib/python3.11/site-packages (from transformers) (2.2.2)\n",
      "Requirement already satisfied: packaging>=20.0 in ./venv/lib/python3.11/site-packages (from transformers) (24.2)\n",
      "Requirement already satisfied: pyyaml>=5.1 in ./venv/lib/python3.11/site-packages (from transformers) (6.0.2)\n",
      "Requirement already satisfied: regex!=2019.12.17 in ./venv/lib/python3.11/site-packages (from transformers) (2024.11.6)\n",
      "Requirement already satisfied: requests in ./venv/lib/python3.11/site-packages (from transformers) (2.32.3)\n",
      "Requirement already satisfied: tokenizers<0.22,>=0.21 in ./venv/lib/python3.11/site-packages (from transformers) (0.21.0)\n",
      "Requirement already satisfied: safetensors>=0.4.1 in ./venv/lib/python3.11/site-packages (from transformers) (0.5.2)\n",
      "Requirement already satisfied: tqdm>=4.27 in ./venv/lib/python3.11/site-packages (from transformers) (4.67.1)\n",
      "Requirement already satisfied: typing-extensions>=4.10.0 in ./venv/lib/python3.11/site-packages (from torch) (4.12.2)\n",
      "Requirement already satisfied: networkx in ./venv/lib/python3.11/site-packages (from torch) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in ./venv/lib/python3.11/site-packages (from torch) (3.1.5)\n",
      "Requirement already satisfied: fsspec in ./venv/lib/python3.11/site-packages (from torch) (2024.9.0)\n",
      "Requirement already satisfied: sympy==1.13.1 in ./venv/lib/python3.11/site-packages (from torch) (1.13.1)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in ./venv/lib/python3.11/site-packages (from sympy==1.13.1->torch) (1.3.0)\n",
      "Requirement already satisfied: pyarrow>=15.0.0 in ./venv/lib/python3.11/site-packages (from datasets) (19.0.0)\n",
      "Requirement already satisfied: dill<0.3.9,>=0.3.0 in ./venv/lib/python3.11/site-packages (from datasets) (0.3.8)\n",
      "Requirement already satisfied: pandas in ./venv/lib/python3.11/site-packages (from datasets) (2.2.3)\n",
      "Requirement already satisfied: xxhash in ./venv/lib/python3.11/site-packages (from datasets) (3.5.0)\n",
      "Requirement already satisfied: multiprocess<0.70.17 in ./venv/lib/python3.11/site-packages (from datasets) (0.70.16)\n",
      "Requirement already satisfied: aiohttp in ./venv/lib/python3.11/site-packages (from datasets) (3.11.11)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in ./venv/lib/python3.11/site-packages (from aiohttp->datasets) (2.4.4)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in ./venv/lib/python3.11/site-packages (from aiohttp->datasets) (1.3.2)\n",
      "Requirement already satisfied: attrs>=17.3.0 in ./venv/lib/python3.11/site-packages (from aiohttp->datasets) (25.1.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in ./venv/lib/python3.11/site-packages (from aiohttp->datasets) (1.5.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in ./venv/lib/python3.11/site-packages (from aiohttp->datasets) (6.1.0)\n",
      "Requirement already satisfied: propcache>=0.2.0 in ./venv/lib/python3.11/site-packages (from aiohttp->datasets) (0.2.1)\n",
      "Requirement already satisfied: yarl<2.0,>=1.17.0 in ./venv/lib/python3.11/site-packages (from aiohttp->datasets) (1.18.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in ./venv/lib/python3.11/site-packages (from requests->transformers) (3.4.1)\n",
      "Requirement already satisfied: idna<4,>=2.5 in ./venv/lib/python3.11/site-packages (from requests->transformers) (3.10)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in ./venv/lib/python3.11/site-packages (from requests->transformers) (2.3.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in ./venv/lib/python3.11/site-packages (from requests->transformers) (2025.1.31)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in ./venv/lib/python3.11/site-packages (from jinja2->torch) (3.0.2)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in ./venv/lib/python3.11/site-packages (from pandas->datasets) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in ./venv/lib/python3.11/site-packages (from pandas->datasets) (2025.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in ./venv/lib/python3.11/site-packages (from pandas->datasets) (2025.1)\n",
      "Requirement already satisfied: six>=1.5 in ./venv/lib/python3.11/site-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.17.0)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install transformers torch datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/lhinds/repos/experimental-projects/FGSM_LLM_Attack/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Loading Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model(model_name=\"gpt2\"):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "    model = AutoModelForCausalLM.from_pretrained(model_name)\n",
    "    model.eval()  # Set to evaluation mode\n",
    "    return model, tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss Generation Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_loss(model, tokenizer, text):\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\", padding=True, truncation=True)\n",
    "    \n",
    "    input_ids = inputs[\"input_ids\"]\n",
    "    attention_mask = inputs[\"attention_mask\"]  # Ensure padding is handled\n",
    "    embeddings = model.get_input_embeddings()(input_ids).detach()\n",
    "    embeddings.requires_grad_()\n",
    "\n",
    "    labels = input_ids.clone()\n",
    "    outputs = model(inputs_embeds=embeddings, labels=labels, attention_mask=attention_mask)\n",
    "    loss = outputs.loss\n",
    "\n",
    "    return loss, embeddings, attention_mask  # Return attention_mask as well\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FGSM Attack Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_fgsm_attack(embeddings, loss, epsilon):\n",
    "    loss.backward()  # Compute gradients\n",
    "    perturbation = epsilon * embeddings.grad.sign()  # FGSM perturbation\n",
    "    perturbed_embeddings = embeddings + perturbation  # Add perturbation\n",
    "    return perturbed_embeddings\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Output Comparison Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_outputs(model, tokenizer, original_text, perturbed_embeddings, attention_mask):\n",
    "    # Generate original output\n",
    "    original_ids = tokenizer(original_text, return_tensors=\"pt\", padding=True, truncation=True)\n",
    "    original_output = tokenizer.decode(\n",
    "        model.generate(**original_ids, pad_token_id=tokenizer.eos_token_id)[0]\n",
    "    )\n",
    "\n",
    "    # Generate adversarial output using perturbed embeddings\n",
    "    perturbed_output = tokenizer.decode(\n",
    "        model.generate(\n",
    "            inputs_embeds=perturbed_embeddings,\n",
    "            attention_mask=attention_mask,  # Ensure correct token handling\n",
    "            pad_token_id=tokenizer.eos_token_id,  # Set PAD token explicitly\n",
    "            max_new_tokens=50  # Control output length properly\n",
    "        )[0]\n",
    "    )\n",
    "\n",
    "    print(\"Original Output:\", original_output)\n",
    "    print(\"Adversarial Output:\", perturbed_output)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the Attack\n",
    "Now let's execute the attack with a sample text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original Output: The weather today is great. I'm going to be back in the office for a few days. I'm going to be back\n",
      "Adversarial Output:  the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the\n"
     ]
    }
   ],
   "source": [
    "# Load the model and tokenizer\n",
    "model, tokenizer = load_model()\n",
    "\n",
    "# Fix: Define padding token explicitly\n",
    "tokenizer.pad_token = tokenizer.eos_token  \n",
    "\n",
    "# Set up the attack parameters\n",
    "original_text = \"The weather today is great.\"\n",
    "epsilon = 0.5 # You can modify this value\n",
    "\n",
    "# Generate the adversarial example\n",
    "loss, inputs, attention_mask = generate_loss(model, tokenizer, original_text)\n",
    "perturbed_input = apply_fgsm_attack(inputs, loss, epsilon)\n",
    "\n",
    "# Compare the results\n",
    "compare_outputs(model, tokenizer, original_text, perturbed_input, attention_mask)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}