ducument de génération de graphes pour des perturbations

mechir/perturbation/GenericAutomatizedPerturbation.py

+#!/usr/bin/env python
+# coding: utf-8
+
+# ## Générateur de graphiques des performances pour des perturbations données
+# Modèle : E5  
+# Dataset : Challenge evalLLM  
+# Exploite une liste de pertubation définies dans un fichier `perturbation.txt` sous trois formes:
+# - `mot -> mot` pour un remplacement (éventuelement avec le mot vide pour supprimer)
+# - `mot +` pour ajouter la requête puis le mot à gauche du document
+# - `+ mot` pour ajouter le mot puis la requête à droite du document
+#      
+# Les graphiques sont enregistrés dans `./res/<descripteur-de-la-perturbation>.png`
+# 
+
+## paramètres du document
+dot_model_name = "intfloat/multilingual-e5-small"    # "sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco"
+
+data_path = "../challenge_data/mainV2/paires_mainV2.csv" # "vaswani" "../../../data/data_challenge/pairesV2_generees_1pourcent.csv"
+query_field = "question"        # "question" default : "text"
+text_field = "text"             # default : "text"
+is_irdata = False
+query_id_subset=["1"] # default : None
+
+perturbations_path = "perturb_test.txt"
+gen_matrices = True
+gen_mean = True
+gen_std = True
+
+load_values = False # default = False
+load_path = "res/mE5-MainV2/elephant-rose/perts_values.json"
+save_values = True # default = False
+
+res_path = "res/mE5-MainV2/elephant-rose/" #datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")}
+
+from mechir import Dot
+from mechir.data import MechIRDataset, MechDataset, DotDataCollator
+from mechir.perturb import perturbation
+from mechir.plotting import plot_components
+
+import torch
+from torch.utils.data import DataLoader
+
+import matplotlib.pyplot as plt
+from matplotlib.ticker import MaxNLocator
+import seaborn as sns
+
+import pandas as pd
+import numpy as np
+
+from tqdm import tqdm
+import json
+
+
+#Perturbations parametrées
+def param_append(doc, mot="microwave"):
+    return doc + " " + mot
+
+def param_prepend(doc, mot="microwave"):
+    return mot+ " " + doc
+
+def param_replace(doc, mot_orig="microwave", mot_rempl="toaster"):
+    return doc.replace(mot_orig, mot_rempl)
+
+
+# génération des perturbations à produire à partir du document
+def generer_transformations(fichier_regles):
+    transformations = []
+    perts_name=[]
+
+    with open(fichier_regles, "r", encoding="utf-8") as f:
+        for ligne in f:
+            ligne = ligne.strip()
+
+            if not ligne:
+                continue
+
+            # Cas 1 : remplacement ou suppression (mot->mot ou mot->)
+            if "->" in ligne:
+                gauche, droite = ligne.split("->", 1)
+                mot_present = gauche.strip()
+                mot_remplacant = droite.strip()
+                # print(f"{gauche}, {droite}")
+                # print(f"{mot_present}, {mot_remplacant}")
+                transformations.append( ("replace", mot_present, mot_remplacant)
+                )
+                perts_name.append(ligne)
+
+            # Cas 2 : ajout en début (mot+)
+            elif ligne.endswith("+"):
+                mot = ligne[:-1].strip()
+                transformations.append( ("prepend", mot))
+                perts_name.append(ligne)
+
+            # Cas 3 : ajout en fin (+mot)
+            elif ligne.startswith("+"):
+                mot = ligne[1:].strip()
+                transformations.append( ("append", mot))
+                perts_name.append(ligne)
+
+            else:
+                raise ValueError(f"Règle non reconnue : {ligne}")
+
+    return transformations, perts_name
+
+# Helper function to print query/baseline document/perturbed document triplets
+def pretty_print_triplets(batch, tokenizer, num=1):
+    """
+    Pretty prints triplets of queries, documents, and their corresponding perturbed documents from a batch.
+
+    Args:
+        batch (dict): A dictionary containing 'queries', 'documents', and 'perturbed_documents' from a DataLoader.
+        tokenizer: The tokenizer used to decode the input IDs.
+        num (int): Number of examples to show per batch.
+    """
+    # Get the queries, documents, and perturbed documents from the batch
+    queries = batch["queries"]
+    documents = batch["documents"]
+    perturbed_documents = batch["perturbed_documents"]
+
+    # Loop through number of examples to show in batch
+    for i in range(len(documents["input_ids"][:num])):
+        # Get the input IDs
+        query_ids = queries["input_ids"][i]
+        original_ids = documents["input_ids"][i]
+        perturbed_ids = perturbed_documents["input_ids"][i]
+
+        # Decode the input IDs to text
+        query_decoded = tokenizer.decode(query_ids.tolist(), skip_special_tokens=False).replace("[PAD]", "").strip()
+        original_doc_decoded = tokenizer.decode(original_ids.tolist(), skip_special_tokens=False).replace("[PAD]", "").strip()
+        perturbed_doc_decoded = tokenizer.decode(perturbed_ids.tolist(), skip_special_tokens=False).replace("[PAD]", "").strip()
+
+        # Pretty print
+        # print(f"Triplet {i + 1}:")
+        print("Query:", query_decoded)
+        print("Baseline Document:", original_doc_decoded)
+        print("Perturbed Document:", perturbed_doc_decoded)
+        print("=" * 50)  # Separator for clarity
+
+# Helper function to calculate and store performances
+def calculate_performance(model, dataloader, baseline_performance, perturbed_performance, diags=False):
+    for i, batch in tqdm(enumerate(dataloader)):
+        # Get the queries, documents, and perturbed documents from the batch
+        queries = batch["queries"]
+        documents = batch["documents"]
+        perturbed_documents = batch["perturbed_documents"]
+        
+        # Encode queries, baseline, and perturbed documents
+        queries_encoded = model.forward(**queries)  # [batch_size x hidden_dim]
+        baseline_encoded = model.forward(**documents)  # [batch_size x hidden_dim]
+        perturbed_encoded = model.forward(**perturbed_documents)  # [batch_size x hidden_dim]
+
+
+        # Calculate scores
+        baseline_scores = torch.sum(queries_encoded.unsqueeze(1) * baseline_encoded.unsqueeze(0), dim=2)
+        perturbed_scores = torch.sum(queries_encoded.unsqueeze(1) * perturbed_encoded.unsqueeze(0), dim=2)
+
+        # Append flattened scores to the performance lists
+        if diags:
+            baseline_performance += baseline_scores.diag().tolist()
+            perturbed_performance += perturbed_scores.diag().tolist()
+        else:
+            baseline_performance += baseline_scores.flatten().tolist()
+            perturbed_performance += perturbed_scores.flatten().tolist()
+
+def plot_scores(baseline_scores, perturbed_scores, transformation_descr, save_path = None):
+    fig, ax = plt.subplots(1, 1, figsize=(15, 4), sharey=True)
+    fig.suptitle('Distribution of Baseline vs Perturbed Scores '+transformation_descr, fontsize=16)
+    sns.kdeplot(baseline_scores, label='Baseline', color='#D55E00', fill=True, ax=ax, alpha=0.5)
+    sns.kdeplot(perturbed_scores, label='Perturbed', color='#009E73', fill=True, ax=ax, alpha=0.5)
+    ax.set_ylabel('Density')
+    ax.set_xlabel('Distances between R(q) and R(d)')
+    ax.legend()
+    #plt.tight_layout(rect=[0, 0, 1, 0.95])  # Adjust layout to make space for the title
+    
+    if save_path is not None:
+        plt.savefig(save_path)
+        plt.close()
+    else:
+        plt.show()
+
+
+def plot_score_dists_mult(all_baseline_scores, all_perturbed_scores, plot_type="hist", save_path = None):
+    # Number of subplots (one per perturbation type)
+    num_plots = len(all_baseline_scores)
+
+    # Set up the figure to hold multiple subplots in a single row
+    fig, axs = plt.subplots(1, num_plots, figsize=(15, 4), sharey=True)
+    fig.suptitle('Distribution of Baseline vs Perturbed Scores', fontsize=16)
+
+    for idx, (perturb_type, ax) in enumerate(zip(all_baseline_scores.keys(), axs)):
+        baseline_scores = all_baseline_scores[perturb_type]
+        perturbed_scores = all_perturbed_scores[perturb_type]
+
+        if plot_type == "hist":
+            ax.hist(baseline_scores, label='Baseline', color='blue', alpha=0.5, bins=15)
+            ax.hist(perturbed_scores, label='Perturbed', color='orange', alpha=0.5, bins=15)
+            ax.yaxis.set_major_locator(MaxNLocator(integer=True))
+            ax.set_ylabel('Frequency')
+            ax.legend()
+        elif plot_type == "kde":
+            # Use seaborn for KDE plot, smoother distribution representation
+            sns.kdeplot(baseline_scores, label='Baseline', color='#D55E00', fill=True, ax=ax, alpha=0.5)
+            sns.kdeplot(perturbed_scores, label='Perturbed', color='#009E73', fill=True, ax=ax, alpha=0.5)
+            ax.set_ylabel('Density')
+            ax.legend()
+        elif plot_type == "box":
+            ax.boxplot([baseline_scores, perturbed_scores], tick_labels=['Baseline', 'Perturbed'])
+            ax.set_ylabel('Scores')
+
+        ax.set_xlabel('Distances between R(q) and R(d)')
+        ax.set_title(f'{perturb_type.capitalize()}')
+
+     #plt.tight_layout(rect=[0, 0, 1, 0.95])  # Adjust layout to make space for the title
+    if save_path is not None:
+        plt.savefig(save_path)
+        plt.close()
+    else:
+        plt.show()
+
+    return
+
+
+def plot_components_V2(
+    data,  # shape: (num_layers, num_heads) or (num_layers, num_heads + 1) if include_mlp=True
+    save_path=None,
+    title="Component Patching Results",
+    include_mlp=False,
+):
+
+    data = data.astype(float)
+    plt.figure(figsize=(10, 6))
+
+    ax = sns.heatmap(
+        np.abs(data),
+        cmap="margma_r",
+        vmin=0,
+        vmax=40,
+        xticklabels=True,
+        yticklabels=True,
+        fmt=".2f",
+        annot=True,
+        annot_kws={"size": 8},
+    )
+
+    if include_mlp:
+        new_labels = list(map(str, range(data.shape[1] - 1)))
+        new_labels.append("MLP")
+        ax.set_xticklabels(new_labels)
+
+    plt.title(title)
+    plt.xlabel("Head")
+    plt.ylabel("Layer")
+
+    if save_path:
+        plt.savefig(save_path)
+        plt.close()
+
+    else:
+        plt.show()
+
+def load_pert_result(filepath):
+    with open(filepath, 'r', encoding="utf-8") as load_file:
+        perts_values = json.load(load_file)
+    for pert_name in perts_values:
+        if  "patching_mean" in perts_values[pert_name] :
+            perts_values[pert_name]["patching_mean"] = np.asarray(perts_values[pert_name]["patching_mean"])
+        if "patching_std" in perts_values[pert_name]:
+            perts_values[pert_name]["patching_std"] = np.asarray(perts_values[pert_name]["patching_std"])
+
+    return perts_values
+
+
+def calculate_components(param_pert_dot_dataloader, gen_mean, gen_std):
+    patching_head_outputs = []
+    for _, batch in tqdm(enumerate(param_pert_dot_dataloader)):
+        queries = batch["queries"]
+        documents = batch["documents"]
+        perturbed_documents = batch["perturbed_documents"]
+
+        patch_head_out = dot_model.patch(queries, documents, perturbed_documents, patch_type="head_all")
+        # Interruption du calcul si la matrice contient un NaN
+        if patch_head_out[0].isnan().any().item() : 
+            print("Component Patching Result non calculé : contient au moins un NaN")
+            return None, None
+        patching_head_outputs.append(patch_head_out)
+
+    mean_head_outputs, std_head_outputs = None, None
+    if gen_mean :
+        mean_head_outputs = torch.mean(torch.stack([tens for tens,_ in patching_head_outputs]), axis=0).detach().to("cpu").numpy()
+    if gen_std :
+        std_head_outputs = torch.std(torch.stack([tens for tens,_ in patching_head_outputs]), axis=0).detach().to("cpu").numpy()
+
+    return mean_head_outputs, std_head_outputs
+
+def load_dataset(data_path, is_irdata, query_id_subset = None, query_field = "text", text_field ="text"):
+    if is_irdata:
+        dataset = MechIRDataset(data_path, query_id_subset=query_id_subset)
+        # print("Number of query,doc pairs in dataset:", len(dataset))
+    else:
+        data = pd.read_csv(data_path)
+        dataset = MechDataset(data, query_field=query_field, text_field=text_field)
+        # print("Number of query,doc pairs in dataset:", len(dataset))
+    return dataset
+
+# Calcul ou extraction des valeurs
+if load_values:
+    perts_values = load_pert_result(load_path)
+    # print(perts_values)
+else:
+    # Récup du modèle
+    dot_model = Dot(dot_model_name)
+
+
+    # Recup dataset
+    dataset = load_dataset(data_path, is_irdata, query_id_subset, query_field, text_field)
+
+
+    perts, perts_name = generer_transformations(perturbations_path)
+    perts_values = {}
+    for i in range(len(perts)):
+        print(perts[i], "="*50)
+        if perts[i][0]=="replace":
+            pert_fun = perturbation(lambda texte : param_replace(texte, perts[i][1], perts[i][2]))
+        elif perts[i][0]=="prepend":
+            pert_fun = perturbation(lambda texte : param_prepend(texte, perts[i][1]))
+        elif perts[i][0]=="append":
+            pert_fun = perturbation(lambda texte : param_append(texte, perts[i][1]))
+        else:
+            raise RuntimeError("Unreachable")
+        param_pert_dot_collator = DotDataCollator(dot_model.tokenizer, pert_fun,  perturb_type=perts[i][0])
+        param_pert_dot_dataloader = DataLoader(dataset, collate_fn=param_pert_dot_collator)
+
+        # param_pert_batch = next(iter(param_pert_dot_dataloader))
+        # print("PERTURBATION",i,":", perts_name[i])
+        # pretty_print_triplets(param_pert_batch, dot_model.tokenizer, num=2)
+
+        baseline_perf = []
+        perturbed_perf = []
+        calculate_performance(dot_model, param_pert_dot_dataloader, baseline_perf, perturbed_perf, diags=True)
+
+        # plot_scores(baseline_perf, perturbed_perf, perts_name[i], save_path=res_path+"PerturbationScore_"+perts_name[i]+".png")
+
+        if gen_matrices:
+            mean_head_outputs, std_head_outputs = calculate_components(param_pert_dot_dataloader, gen_mean, gen_std)
+            # if mean_head_outputs is not None:
+            #     plot_components_V2(mean_head_outputs, title="Components Patching Results (mean) for "+ perts_name[i], save_path=res_path+"ComponentPatching-mean_"+perts_name[i]+"V2.png")
+            # if std_head_outputs is not None:
+            #     plot_components_V2(std_head_outputs, title="Components Patching Results (std) for "+ perts_name[i], save_path=res_path+"ComponentPatching-std_"+perts_name[i]+"V2.png")
+
+
+        # Sauvegarde des valeurs pour la i^eme perturbation
+        
+        perts_values[perts_name[i]] = {"baseline": baseline_perf, "perturbed": perturbed_perf}
+        if gen_matrices and mean_head_outputs is not None :
+            perts_values[perts_name[i]]["patching_mean"] = mean_head_outputs.tolist()
+        if gen_matrices and std_head_outputs is not None:
+            perts_values[perts_name[i]]["patching_std"] = std_head_outputs.tolist()
+
+    if save_values:
+        with open(f"{res_path}/perts_values.json", 'w', encoding="utf-8") as result_file:
+            json.dump(perts_values, result_file, indent=2) 
+
+
+# Génération des graphiques
+print("Génération des graphiques")
+for pert_name, pert_values in perts_values.items():
+    plot_scores(pert_values["baseline"], pert_values["perturbed"], pert_name, save_path=res_path+"PerturbationScore_"+pert_name+".png")
+    if "patching_mean" in pert_values and pert_values["patching_mean"] is not None:
+        plot_components(pert_values["patching_mean"], title="Components Patching Results (mean) for "+ pert_name, save_path=res_path+"ComponentPatching-mean_"+pert_name+".png")
+        plot_components_V2(pert_values["patching_mean"], title="Components Patching Results (mean) for "+ pert_name, save_path=res_path+"ComponentPatching-mean_"+pert_name+"V2.png")
+
+    if "patching_std" in pert_values and pert_values["patching_std"] is not None:
+        plot_components(pert_values["patching_std"], title="Components Patching Results (std) for "+ pert_name, save_path=res_path+"ComponentPatching-std_"+pert_name+".png")
+        plot_components_V2(pert_values["patching_std"], title="Components Patching Results (std) for "+ pert_name, save_path=res_path+"ComponentPatching-std_"+pert_name+"V2.png")