{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "10ea03b4",
   "metadata": {},
   "source": [
    "# Exploration de techniques pour produire des perturbations pour les méthodes d'explication par *Activation Patching*\n",
    "\n",
    "## Background\n",
    "### Activation Patching [[Geiger et al., 2021](https://dl.acm.org/doi/10.5555/3540261.3540994)] \n",
    "Les méthodes par activation patching cherchent à identifier les composants (neurones, couches, tête d'attention,...) des modèles d'apprentissage profond sensibles à certains changements dans l'entrée ou à certaines caractéristiques. L'objectif peut être de localiser les compétences linguistiques ou la sensibilité à un certain domaine ou vocabulaire dans un modèle de langue. \n",
    "\n",
    "Pour produire ces explications, la méthode présentée par [[Geiger et al., 2021](https://dl.acm.org/doi/10.5555/3540261.3540994)] procède ainsi:  \n",
    "Soient $M$ le modèle étudié qui contient des composants (neurone ou tête d'attention) $M_{i,j}$, $\\mathcal{D}$ un dataset et $P$ une fonction qui à une entrée de $e \\in \\mathcal{D}$ lui associe une version perturbée $e'$. On pose $\\mathcal{D}'$ le dataset issu de $\\mathcal{D}$ dont chaque entrée est perturbée selon $P$.  \n",
    "Itérativement, on effectue pour chaque composant $M_{i,j}$:\n",
    "\n",
    "- une passe de $\\mathcal{D}$ dans $M$ en notant la valeur en sortie de $M_{i,j}$ ($v_e$) pour chaque entrée $e$. On note sa performance $p$.\n",
    "- une passe de $\\mathcal{D}'$ dans $M$. On note sa performance $p_{\\mathcal{D}'}$.\n",
    "- une passe de $\\mathcal{D}'$ dans $M$ où la sortie du composant $M_{i,j}$ est fixée à $v_e$ pour l'entrée $e' = P(e)$. On note sa performance $p_{i,j}'$.\n",
    "\n",
    "Ainsi, on obtient une différence de performance $p - p_{i,j}'$ pour chaque composant $M_{i,j}$ qui peut être interprété comme la sensibilité du composant $M_{i,j}$ à la perturbation $P$.\n",
    "\n",
    "\n",
    "### Activation patching pour l'extraction de document [[Chen et al., 2024](https://dl.acm.org/doi/10.1145/3626772.3657841)]\n",
    "[[Chen et al., 2024](https://dl.acm.org/doi/10.1145/3626772.3657841)] propose une adaptation de la méthode d'activation patching pour les modèles d'extraction d'information par plongement des documents et des requêtes à l'aide d'un modèle bi-encodeur.  \n",
    "On conserve les notations précédentes. $\\mathcal{D}$ est l'ensemble des documents et $\\mathcal{Q}$ est l'ensemble des requêtes.\n",
    "\n",
    "Itérativement, on effectue pour chaque composant $M_{i,j}$:\n",
    "\n",
    "- une passe des paires étiquetées $ \\subset \\mathcal{Q} \\times \\mathcal{D}$ dans $M$ en notant la valeur en sortie de $M_{i,j}$ ($v^\\mathcal{D}_e$) pour chaque entrée $e=(q,d)$. On note sa performance $p_{\\mathcal{D}} = moy_{(q,d) \\in \\mathcal{Q} \\times \\mathcal{D}} (M(q)\\bullet M(d))$.\n",
    "- une passe des paires étiquetées $ \\subset \\mathcal{Q} \\times \\mathcal{D'}$ dans $M$ en notant la valeur en sortie de $M_{i,j}$ ($v^\\mathcal{D'}_e$) pour chaque entrée $e=(q,d')$. On note sa performance $p_{\\mathcal{D'}} = moy_{(q,d') \\in \\mathcal{Q} \\times \\mathcal{D'}} (M(q)\\bullet M(d'))$.\n",
    "- On pose $ \\forall e \\in \\mathcal{Q} \\times \\mathcal{D}, v_e = v^{\\mathcal{S}}_e$ où $\\mathcal{S}$ est tel que $p_{\\mathcal{S}} = max\\{p_{\\mathcal{D}}, p_{\\mathcal{D'}}\\}$ ($p_{\\mathcal{V}} = min\\{p_{\\mathcal{D}}, p_{\\mathcal{D'}}\\}$) \n",
    "- une passe de $\\mathcal{V}$ dans $M$ où la sortie du composant $M_{i,j}$ est fixée à $v^{\\mathcal{S}}_e$ pour l'entrée $e' = P(e)$. On note sa performance $p_{i,j}'$.\n",
    "\n",
    "Ainsi, on peut considérer\n",
    "$$\n",
    "p = \\frac{p_{\\mathcal{S}} - p_{\\mathcal{D}}}{p_{\\mathcal{D}'} - p_{\\mathcal{D}}} \n",
    "$$\n",
    "comme la sensibilité composant $M_{i,j}$ à la perturbation $P$. Les sensibilités de chaque composant peut être récapitulé dans un tableau. \n",
    "\n",
    "\n",
    "Remarques:\n",
    "- La méthode proposée prend en compte les perturbations pouvant augmenter ou diminuer la performance.\n",
    "- La métrique couramment utilisée pour la performance est la proximité entre les représentations d'une requête et d'un document (ie. le produit scalaire).\n",
    "\n",
    "\n",
    "### MerchIR\n",
    "[[Parry et al., 2025](https://arxiv.org/abs/2501.10165)] propose une implémentation de la méthode de [[Chen et al., 2024](https://dl.acm.org/doi/10.1145/3626772.3657841)] : [Github](https://github.com/Parry-Parry/MechIR.git).  \n",
    "En fournissant un modèle, une base de données et une fonction de perturbation, la librairie MechIR permet de construire le tableau des perturbations que chaque composant de l'architecture fournie.\n",
    "\n",
    "On développe dans la suite des méthodes et approches pour exploiter au mieux cette librairie."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a69dbec3",
   "metadata": {},
   "source": [
    "#### Préparation des uotils de travail de MechIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "304d73db",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/marine/Documents/Cours/M2GPEx/ProjetIndividuel/Travaux/stagexairag/venv-mechir/lib/python3.12/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",
      "2026-04-08 13:11:34.469 WARNING streamlit.runtime.scriptrunner_utils.script_run_context: Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.\n"
     ]
    }
   ],
   "source": [
    "# Librairies Necessaires\n",
    "\n",
    "from mechir import Dot\n",
    "from mechir.data import MechIRDataset, DotDataCollator\n",
    "from mechir.perturb import perturbation\n",
    "from mechir.plotting import plot_components\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0bb6d76c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fonctions auxiliaires (pas besoin de lire)\n",
    "def pretty_print_triplets(batch, tokenizer, num=1):\n",
    "    \"\"\"\n",
    "    Pretty prints triplets of queries, documents, and their corresponding perturbed documents from a batch.\n",
    "\n",
    "    Args:\n",
    "        batch (dict): A dictionary containing 'queries', 'documents', and 'perturbed_documents' from a DataLoader.\n",
    "        tokenizer: The tokenizer used to decode the input IDs.\n",
    "        num (int): Number of examples to show per batch.\n",
    "    \"\"\"\n",
    "    # Get the queries, documents, and perturbed documents from the batch\n",
    "    queries = batch[\"queries\"]\n",
    "    documents = batch[\"documents\"]\n",
    "    perturbed_documents = batch[\"perturbed_documents\"]\n",
    "\n",
    "    # Loop through number of examples to show in batch\n",
    "    for i in range(len(documents[\"input_ids\"][:num])):\n",
    "        # Get the input IDs\n",
    "        query_ids = queries[\"input_ids\"][i]\n",
    "        original_ids = documents[\"input_ids\"][i]\n",
    "        perturbed_ids = perturbed_documents[\"input_ids\"][i]\n",
    "\n",
    "        # Decode the input IDs to text\n",
    "        query_decoded = tokenizer.decode(query_ids.tolist(), skip_special_tokens=False).replace(\"[PAD]\", \"\").strip()\n",
    "        original_doc_decoded = tokenizer.decode(original_ids.tolist(), skip_special_tokens=False).replace(\"[PAD]\", \"\").strip()\n",
    "        perturbed_doc_decoded = tokenizer.decode(perturbed_ids.tolist(), skip_special_tokens=False).replace(\"[PAD]\", \"\").strip()\n",
    "\n",
    "        # Pretty print\n",
    "        # print(f\"Triplet {i + 1}:\")\n",
    "        print(\"Query:\", query_decoded)\n",
    "        print(\"Baseline Document:\", original_doc_decoded)\n",
    "        print(\"Perturbed Document:\", perturbed_doc_decoded)\n",
    "        print(\"=\" * 50)  # Separator for clarity\n",
    "\n",
    "# Helper function to calculate and store performances\n",
    "def calculate_performance(model, dataloader, baseline_performance, perturbed_performance):\n",
    "    for i, batch in enumerate(dataloader):\n",
    "        # Get the queries, documents, and perturbed documents from the batch\n",
    "        queries = batch[\"queries\"]\n",
    "        documents = batch[\"documents\"]\n",
    "        perturbed_documents = batch[\"perturbed_documents\"]\n",
    "\n",
    "        # Encode queries, baseline, and perturbed documents\n",
    "        queries_encoded = model.forward(**queries)  # [batch_size x hidden_dim]\n",
    "        baseline_encoded = model.forward(**documents)  # [batch_size x hidden_dim]\n",
    "        perturbed_encoded = model.forward(**perturbed_documents)  # [batch_size x hidden_dim]\n",
    "\n",
    "        # Calculate scores\n",
    "        baseline_scores = torch.sum(queries_encoded.unsqueeze(1) * baseline_encoded.unsqueeze(0), dim=2)\n",
    "        perturbed_scores = torch.sum(queries_encoded.unsqueeze(1) * perturbed_encoded.unsqueeze(0), dim=2)\n",
    "\n",
    "        # Append flattened scores to the performance lists\n",
    "        baseline_performance += baseline_scores.flatten().tolist()\n",
    "        perturbed_performance += perturbed_scores.flatten().tolist()\n",
    "\n",
    "def plot_scores(baseline_scores, perturbed_scores, transformation_descr, save=False):\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(15, 4), sharey=True)\n",
    "    fig.suptitle('Distribution of Baseline vs Perturbed Scores '+transformation_descr, fontsize=16)\n",
    "    sns.kdeplot(baseline_scores, label='Baseline', color='#D55E00', fill=True, ax=ax, alpha=0.5)\n",
    "    sns.kdeplot(perturbed_scores, label='Perturbed', color='#009E73', fill=True, ax=ax, alpha=0.5)\n",
    "    ax.set_ylabel('Density')\n",
    "    ax.legend()\n",
    "    #plt.tight_layout(rect=[0, 0, 1, 0.95])  # Adjust layout to make space for the title\n",
    "    if save:\n",
    "        plt.savefig('res/'+ transformation_descr+'_scores.png')\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a6aef7c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Support for BERT in TransformerLens is currently experimental, until such a time when it has feature parity with HookedTransformer and has been tested on real research tasks. Until then, backward compatibility is not guaranteed. Please see the docs for information on the limitations of the current implementation.\n",
      "If using BERT for interpretability research, keep in mind that BERT has some significant architectural differences to GPT. For example, LayerNorms are applied *after* the attention and MLP components, meaning that the last LayerNorm in a block cannot be folded.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  cpu\n",
      "Loaded pretrained model sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco into HookedEncoder\n"
     ]
    }
   ],
   "source": [
    "# Récupération du modèle à évaluer (cas des modèles bi-encodeur)\n",
    "dot_model_name = \"sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco\"\n",
    "dot_model = Dot(dot_model_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d860ecda",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of query,doc pairs in dataset: 19\n",
      "Query: MEASUREMENT OF DIELECTRIC CONSTANT OF LIQUIDS BY THE USE OF MICROWAVE TECHNIQUES\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Recupération du dataset utilisé pour l'évaluation\n",
    "dataset = MechIRDataset(\"vaswani\", query_id_subset=[\"1\"])\n",
    "print(\"Number of query,doc pairs in dataset:\", len(dataset))\n",
    "print(\"Query:\", dataset._get_query(\"1\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2019107",
   "metadata": {},
   "source": [
    "## Construction d'une perturbation\n",
    "\n",
    "Pour observer les sensibilités d'un modèle à des perturbation, il est important de construire des perturbations de document pertinentes et qui engendrent un changement conséquent du comportement du modèle. On peut retrouver plusieurs approches.\n",
    "\n",
    "### Perturbation par ajout, remplacement ou suppression d'un mot\n",
    "Les perturbations les plus simples sont:\n",
    "- l'ajout d'un mot en début de document (prepend)\n",
    "- l'ajout d'un mot en fin de document (append)\n",
    "- le remplacement d'un mot par un autre (replace)\n",
    "- la suppression d'un mot (cas particulier de replace avec le mot vide)\n",
    "\n",
    "Le choix important dans ces perturbations est celui du mot ajouté ou du mot replacé et de son remplaçant. On peut exploiter différentes statistiques issues des données pour choisir un mot à ajouter ou remplacer:\n",
    "- la fréquence des mots dans les requêtes\n",
    "- la fréquence des mots dans les documents\n",
    "- le nombre de requête/ documents contenant chaque mot\n",
    "- l'IDF des mots dans les documents/ requêtes\n",
    "- le rapport fréquence dans les requêtes/nombre de documents contenant le mot\n",
    "\n",
    "Dans un second temps, pour le cas du remplacement d'un mot, le choix du mot remplaçant peut se baser sur la connaissance des données et du vocabulaire. On peut exploiter:\n",
    "- les synonymes (remplacer atome par particule)\n",
    "- les différents sens des mots (remplacer onde par vague ou fréquence)\n",
    "- le vocabulaire connexe (remplacer atome par molécule)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25a6e43a",
   "metadata": {},
   "source": [
    "#### Définition et évaluation d'un perturbation simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ae4ef107",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Définition de la fonction de perturbation des données\n",
    "@perturbation\n",
    "def pert(texte):\n",
    "    return texte.replace(\"dielectric\", \"photonic\")\n",
    "\n",
    "pert_type = \"replace\"  # autres valeurs possibles : \"append\", \"prepend\"\n",
    "pert_name = \"replace dielectric by photonic\" # description de la pertubation (pour le titre du graphique)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4f3615e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Définition du traitement des données et du dataloader avec les données perturbées et non-perturbées\n",
    "param_pert_dot_collator = DotDataCollator(dot_model.tokenizer, pert, q_max_length=None, d_max_length=None, perturb_type=pert_type)\n",
    "param_pert_dot_dataloader = DataLoader(dataset, batch_size=16, collate_fn=param_pert_dot_collator)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a85b37a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query: [CLS] measurement of dielectric constant of liquids by the use of microwave techniques [SEP]\n",
      "Baseline Document: [CLS] broadband millimetre wave paramagnetic resonance spectrometer the specimen and waveguide which can be cooled by means of a cryostat are placed between close pole pieces giving high uniform magnetic fields design details and some measurements on zero field splittings are given [SEP]\n",
      "Perturbed Document: [CLS] broadband millimetre wave paramagnetic resonance spectrometer the specimen and waveguide which can be cooled by means of a cryostat are placed between close pole pieces giving high uniform magnetic fields design details and some measurements on zero field splittings are given [SEP]\n",
      "==================================================\n",
      "Query: [CLS] measurement of dielectric constant of liquids by the use of microwave techniques [SEP]\n",
      "Baseline Document: [CLS] microwave measurements of dielectric absorption in dilute solutions [SEP]\n",
      "Perturbed Document: [CLS] microwave measurements of photon a a a a a a a [SEP]\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "# Construction du batch et aperçu des données\n",
    "param_pert_batch = next(iter(param_pert_dot_dataloader))\n",
    "pretty_print_triplets(param_pert_batch, dot_model.tokenizer, num=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "80479fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calcul des performances du modèle sur les données puis les données perturbées\n",
    "baseline_perf = []\n",
    "perturbed_perf = []\n",
    "calculate_performance(dot_model, param_pert_dot_dataloader, baseline_perf, perturbed_perf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "134347bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 71844 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Affichage des performances calculées\n",
    "plot_scores(baseline_perf, perturbed_perf, pert_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edea6dae",
   "metadata": {},
   "source": [
    "On cherche à obtenir des profils de densité différents pour les données de base et les données perturbées. C'est cette variation qu'on va essayer d'attribuer aux composants du modèle."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-mechir (3.12.3)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
