From f03e86a75401cc6a91ca1f7dc523fbb5db5c8ffc Mon Sep 17 00:00:00 2001 From: Andrei Volodichev Date: Mon, 16 Jun 2025 22:44:41 +0300 Subject: [PATCH 01/23] add topological divergence estimator --- src/lm_polygraph/estimators/__init__.py | 1 + .../estimators/topological_divergence.py | 193 ++++++++++++++++++ 2 files changed, 194 insertions(+) create mode 100644 src/lm_polygraph/estimators/topological_divergence.py diff --git a/src/lm_polygraph/estimators/__init__.py b/src/lm_polygraph/estimators/__init__.py index dbabae5e8..d38aea602 100644 --- a/src/lm_polygraph/estimators/__init__.py +++ b/src/lm_polygraph/estimators/__init__.py @@ -84,3 +84,4 @@ from .kernel_language_entropy import KernelLanguageEntropy from .luq import LUQ from .eigenscore import EigenScore +from topological_divergence import TopologicalDivergence \ No newline at end of file diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py new file mode 100644 index 000000000..660e3a6ae --- /dev/null +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -0,0 +1,193 @@ +import warnings +from typing import Dict, List, Tuple, Literal +from collections import defaultdict +import numpy as np +import torch +try: + from joblib import Parallel, delayed + IS_PARALLEL_AVAILABLE = True +except ImportError: + IS_PARALLEL_AVAILABLE = False + warnings.warn( + "Joblib is not installed. Parallel processing for MTopDivCalculator will not be available. " + "Please install it via 'pip install joblib' if you want to use parallel processing." + ) +try: + import ripserplusplus as rpp_py +except ImportError: + raise ImportError( + "Please install the 'ripserplusplus' package to use TopologicalDivergence estimator. " + "You can install it via 'pip install ripserplusplus'." + ) +try: + import mtd.barcodes as mtd +except ImportError: + raise ImportError( + "Please install the 'mtd' package to use TopologicalDivergence estimator.\n" + "Installation steps:\n" + " 1. git clone https://github.com/IlyaTrofimov/MTopDiv.git\n" + " 2. cd MTopDiv && python setup.py install" + ) + +from .estimator import Estimator + +def transform_attention_scores_to_distances( + attention_weights: torch.Tensor, + zero_out: Literal["prompt", "response"], + len_answer: int, + lower_bound: float = 0.0, +) -> torch.Tensor: + """Transform attention matrix to the matrix of distances between tokens. + + Parameters + ---------- + attention_weights : torch.Tensor + Attention matrixes of one sample (n_heads x n_tokens x n_tokens). + zero_out : Literal['prompt', 'response'] + Determines whether to zero out distances between prompt tokens or response tokens. + len_answer : int + Length of the response. + + Returns + ------- + torch.Tensor + Distance matrix. + + """ + n_tokens = attention_weights.shape[1] + distance_mx = 1 - torch.clamp( + attention_weights, min=lower_bound + ) # torch.where(attn_mx > lower_bound, attn_mx, 0.0) + zero_diag = torch.ones(n_tokens, n_tokens) - torch.eye(n_tokens) + distance_mx *= zero_diag.to(attention_weights.device).expand_as( + distance_mx + ) # torch.diag(torch.diag(distance_mx)) + distance_mx = torch.minimum(distance_mx.transpose(1, 2), distance_mx) + + if zero_out == "prompt": + len_prompt = n_tokens - len_answer + distance_mx[:, :len_prompt, :len_prompt] = 0 + elif zero_out == "response": + distance_mx[:, -len_answer:, -len_answer:] = 0 + else: + raise ValueError(f"Unsupported zero_out parameter: {zero_out}") + + return distance_mx.cpu().numpy() + +def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: + """ + Calculate MTopDiv value for the given attention matrix. + """ + barcodes = rpp_py.run("--format distance --dim 1", distance_mx) + barcodes = mtd.barc2array(barcodes) + mtopdiv = mtd.get_score(barcodes, 0, "sum_length") + return mtopdiv + + +class TopologicalDivergence(Estimator): + """ + Estimates the sequence-level uncertainty of a language model following the method of + "Hallucination Detection in LLMs with Topological Divergence on Attention Graphs" + as provided in the paper https://arxiv.org/abs/2504.10063. + Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). + Computes topological divergences between prompt and response attention + graphs to identify hallucination-indicative heads. + """ + def __init__( + self, + selected_heads: List[Tuple[int, int]] = None, + num_layers: int = None, + num_heads: int = None, + zero_output: Literal["prompt", "response"] = "prompt", + n_jobs: int = 16, + critical_size: int = 768 + ): + """ + Initializes TopologicalDivergence estimator. + + Parameters: + selected_heads (List[Tuple[int, int]]): List of attention heads to calculate MTopDiv for. + First integer is layer index, second is head index. + If not provided or empty, all heads will be used. + zero_out Literal["prompt", "response"]: Determines whether to zero out distances between + prompt tokens or response tokens. + If not provided or empty, prompt tokens will be zeroed out. + n_jobs (int): Number of jobs for parallel processing. Default: 16. + critical_size (int): Size threshold for parallel processing. + If the sequence length exceeds this size, n_jobs will be limited to 8. Default: 768. + """ + super().__init__(["greedy_tokens", "forwardpass_attention_weights"], "sequence") + + if selected_heads is None or len(selected_heads) == 0: + # If no specific heads are selected, use all heads + if num_layers is None or num_heads is None: + raise ValueError( + "If no specific heads are specified in 'selected_heads', " + "'num_layers' and 'num_heads' must be provided." + ) + selected_heads = [ + (layer, head) + for layer in range(num_layers) + for head in range(num_heads) + ] + selected_heads_dict = defaultdict(list) + for layer, head in selected_heads: + selected_heads_dict[layer].append(head) + self.selected_heads = selected_heads_dict + self.zero_out = zero_output + self.n_jobs = n_jobs + self.critical_size = critical_size + + def __str__(self): + return "TopologicalDivergence" + + def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: + """ + Calculates sequence-wise MTopDiv scores for selected heads of attention masks. + Returns the mean MTopDiv score for each input text. + + Parameters: + stats (Dict[str, np.ndarray]): input statistics, consisting of: + * tokenized model generations for each input text in 'greedy_tokens' + * attention scores of shape [batch_size, num_layers, num_heads, seq_len, seq_len] + in 'forwardpass_attention_weights' + Returns: + np.ndarray: float uncertainty for each sample in input statistics. + Higher values indicate more uncertain samples. + """ + responses = stats["greedy_tokens"] + attention_weights_batch = stats["forwardpass_attention_weights"] + + layers = sorted(self.selected_heads.keys()) + mtopdivs_batch = [] + for response, attention_weights in zip(responses, attention_weights_batch): + padding_length = np.isnan(attention_weights[0, 0, 0]).sum() + response_length = len(response) + for layer in layers: + mtopdivs_batch.append([]) + heads = self.selected_heads[layer] + selected_attention_weights = torch.from_numpy( + attention_weights[layer, heads, :-padding_length, :-padding_length] + ).float() + distance_matrices = transform_attention_scores_to_distances( + selected_attention_weights, self.zero_out, response_length + ) + if IS_PARALLEL_AVAILABLE: + if selected_attention_weights.shape[-1] <= self.critical_size: + n_jobs = min(self.n_jobs, 8) if n_jobs > 0 else 8 + else: + n_jobs = self.n_jobs + mtopdivs = list( + *Parallel(n_jobs=n_jobs)( + delayed(transform_distances_to_mtopdiv)(distance_matrice) + for distance_matrice in distance_matrices + ) + ) + else: + mtopdivs = list(map( + transform_distances_to_mtopdiv, distance_matrices + )) + mtopdivs_batch[-1].extend(mtopdivs) + mtopdivs_batch = np.array(mtopdivs_batch, dtype=np.float) + + return np.mean(mtopdivs_batch, axis=1) From 9193c30c01bee7e8d1687c6478d8439e74735689 Mon Sep 17 00:00:00 2001 From: Andrei Volodichev Date: Tue, 17 Jun 2025 00:42:11 +0300 Subject: [PATCH 02/23] bug fix --- .../estimators/topological_divergence.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 660e3a6ae..fe6f8eef8 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -1,3 +1,6 @@ +import os +os.environ["TOKENIZERS_PARALLELISM"] = "false" + import warnings from typing import Dict, List, Tuple, Literal from collections import defaultdict @@ -163,31 +166,34 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: for response, attention_weights in zip(responses, attention_weights_batch): padding_length = np.isnan(attention_weights[0, 0, 0]).sum() response_length = len(response) + mtopdivs_batch.append([]) for layer in layers: - mtopdivs_batch.append([]) heads = self.selected_heads[layer] - selected_attention_weights = torch.from_numpy( - attention_weights[layer, heads, :-padding_length, :-padding_length] - ).float() + if padding_length > 0: + selected_attention_weights = torch.from_numpy( + attention_weights[layer, heads, :-padding_length, :-padding_length] + ).float() + else: + selected_attention_weights = torch.from_numpy( + attention_weights[layer, heads, :, :] + ).float() distance_matrices = transform_attention_scores_to_distances( selected_attention_weights, self.zero_out, response_length ) if IS_PARALLEL_AVAILABLE: if selected_attention_weights.shape[-1] <= self.critical_size: - n_jobs = min(self.n_jobs, 8) if n_jobs > 0 else 8 + n_jobs = min(self.n_jobs, 8) if self.n_jobs > 0 else 8 else: n_jobs = self.n_jobs - mtopdivs = list( - *Parallel(n_jobs=n_jobs)( - delayed(transform_distances_to_mtopdiv)(distance_matrice) - for distance_matrice in distance_matrices - ) + mtopdivs = Parallel(n_jobs=n_jobs)( + delayed(transform_distances_to_mtopdiv)(distance_matrice) + for distance_matrice in distance_matrices ) else: mtopdivs = list(map( transform_distances_to_mtopdiv, distance_matrices )) mtopdivs_batch[-1].extend(mtopdivs) - mtopdivs_batch = np.array(mtopdivs_batch, dtype=np.float) + mtopdivs_batch = np.array(mtopdivs_batch, dtype=float) return np.mean(mtopdivs_batch, axis=1) From 4b3adac77951068b3527f2e3efb6cbc4d2d74ced Mon Sep 17 00:00:00 2001 From: Andrei Volodichev Date: Tue, 17 Jun 2025 12:28:07 +0300 Subject: [PATCH 03/23] simplify code and dependencies --- src/lm_polygraph/estimators/__init__.py | 2 +- .../estimators/topological_divergence.py | 81 ++++++++++--------- 2 files changed, 44 insertions(+), 39 deletions(-) diff --git a/src/lm_polygraph/estimators/__init__.py b/src/lm_polygraph/estimators/__init__.py index d38aea602..4e9b7ce57 100644 --- a/src/lm_polygraph/estimators/__init__.py +++ b/src/lm_polygraph/estimators/__init__.py @@ -84,4 +84,4 @@ from .kernel_language_entropy import KernelLanguageEntropy from .luq import LUQ from .eigenscore import EigenScore -from topological_divergence import TopologicalDivergence \ No newline at end of file +from .topological_divergence import TopologicalDivergence diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index fe6f8eef8..9d9a860ab 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -1,11 +1,10 @@ import os -os.environ["TOKENIZERS_PARALLELISM"] = "false" - import warnings from typing import Dict, List, Tuple, Literal from collections import defaultdict import numpy as np import torch + try: from joblib import Parallel, delayed IS_PARALLEL_AVAILABLE = True @@ -14,26 +13,21 @@ warnings.warn( "Joblib is not installed. Parallel processing for MTopDivCalculator will not be available. " "Please install it via 'pip install joblib' if you want to use parallel processing." - ) + ) try: - import ripserplusplus as rpp_py + from ripser import ripser except ImportError: raise ImportError( - "Please install the 'ripserplusplus' package to use TopologicalDivergence estimator. " - "You can install it via 'pip install ripserplusplus'." - ) -try: - import mtd.barcodes as mtd -except ImportError: - raise ImportError( - "Please install the 'mtd' package to use TopologicalDivergence estimator.\n" - "Installation steps:\n" - " 1. git clone https://github.com/IlyaTrofimov/MTopDiv.git\n" - " 2. cd MTopDiv && python setup.py install" + "Please install the 'ripser' package to use TopologicalDivergence estimator. " + "You can install it via 'pip install ripser'." ) from .estimator import Estimator + +os.environ["TOKENIZERS_PARALLELISM"] = "false" + + def transform_attention_scores_to_distances( attention_weights: torch.Tensor, zero_out: Literal["prompt", "response"], @@ -77,50 +71,61 @@ def transform_attention_scores_to_distances( return distance_mx.cpu().numpy() + def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: """ - Calculate MTopDiv value for the given attention matrix. + Compute the MTopDiv (Metric Topological Diversity) score from a distance matrix. + + This function calculates the sum of persistence intervals in the H₀ (zero-dimensional) + persistent homology barcode, corresponding to the lifetimes of connected components + in a Vietoris–Rips filtration. + + Parameters: + distance_mx (np.ndarray): A square, symmetric distance matrix. + + Returns: + float: Sum of finite H₀ barcode lengths (birth–death), representing topological diversity. """ - barcodes = rpp_py.run("--format distance --dim 1", distance_mx) - barcodes = mtd.barc2array(barcodes) - mtopdiv = mtd.get_score(barcodes, 0, "sum_length") - return mtopdiv + barcodes = ripser(distance_mx, distance_matrix=True, maxdim=0)['dgms'] + if len(barcodes) > 0: + return barcodes[0][:-1, 1].sum() + return 0 class TopologicalDivergence(Estimator): """ - Estimates the sequence-level uncertainty of a language model following the method of - "Hallucination Detection in LLMs with Topological Divergence on Attention Graphs" - as provided in the paper https://arxiv.org/abs/2504.10063. + Estimates the sequence-level uncertainty of a language model following the method of + "Hallucination Detection in LLMs with Topological Divergence on Attention Graphs" + as provided in the paper https://arxiv.org/abs/2504.10063. Works only with whitebox models (initialized using lm_polygraph.utils.model.WhiteboxModel). - Computes topological divergences between prompt and response attention + Computes topological divergences between prompt and response attention graphs to identify hallucination-indicative heads. """ def __init__( - self, + self, selected_heads: List[Tuple[int, int]] = None, - num_layers: int = None, + num_layers: int = None, num_heads: int = None, - zero_output: Literal["prompt", "response"] = "prompt", - n_jobs: int = 16, + zero_output: Literal["prompt", "response"] = "prompt", + n_jobs: int = 16, critical_size: int = 768 ): """ Initializes TopologicalDivergence estimator. Parameters: - selected_heads (List[Tuple[int, int]]): List of attention heads to calculate MTopDiv for. - First integer is layer index, second is head index. + selected_heads (List[Tuple[int, int]]): List of attention heads to calculate MTopDiv for. + First integer is layer index, second is head index. If not provided or empty, all heads will be used. - zero_out Literal["prompt", "response"]: Determines whether to zero out distances between - prompt tokens or response tokens. + zero_out Literal["prompt", "response"]: Determines whether to zero out distances between + prompt tokens or response tokens. If not provided or empty, prompt tokens will be zeroed out. n_jobs (int): Number of jobs for parallel processing. Default: 16. - critical_size (int): Size threshold for parallel processing. + critical_size (int): Size threshold for parallel processing. If the sequence length exceeds this size, n_jobs will be limited to 8. Default: 768. """ super().__init__(["greedy_tokens", "forwardpass_attention_weights"], "sequence") - + if selected_heads is None or len(selected_heads) == 0: # If no specific heads are selected, use all heads if num_layers is None or num_heads is None: @@ -140,14 +145,14 @@ def __init__( self.zero_out = zero_output self.n_jobs = n_jobs self.critical_size = critical_size - + def __str__(self): return "TopologicalDivergence" def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: """ - Calculates sequence-wise MTopDiv scores for selected heads of attention masks. - Returns the mean MTopDiv score for each input text. + Calculates sequence-wise MTopDiv scores for selected heads of attention masks. + Returns the mean MTopDiv score for each input text. Parameters: stats (Dict[str, np.ndarray]): input statistics, consisting of: @@ -194,6 +199,6 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: transform_distances_to_mtopdiv, distance_matrices )) mtopdivs_batch[-1].extend(mtopdivs) - mtopdivs_batch = np.array(mtopdivs_batch, dtype=float) + mtopdivs_batch = np.array(mtopdivs_batch, dtype=float) return np.mean(mtopdivs_batch, axis=1) From 10877a8e4d29e1ddfce6e4ca1990eec9eb34f21a Mon Sep 17 00:00:00 2001 From: Andrei Volodichev Date: Tue, 17 Jun 2025 15:52:41 +0300 Subject: [PATCH 04/23] change parallelizm strategy --- .../estimators/topological_divergence.py | 142 +++++++----------- 1 file changed, 56 insertions(+), 86 deletions(-) diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 9d9a860ab..5fc8b023f 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -5,9 +5,11 @@ import numpy as np import torch + try: from joblib import Parallel, delayed IS_PARALLEL_AVAILABLE = True + os.environ["TOKENIZERS_PARALLELISM"] = "false" except ImportError: IS_PARALLEL_AVAILABLE = False warnings.warn( @@ -25,33 +27,25 @@ from .estimator import Estimator -os.environ["TOKENIZERS_PARALLELISM"] = "false" - - def transform_attention_scores_to_distances( - attention_weights: torch.Tensor, - zero_out: Literal["prompt", "response"], - len_answer: int, + attention_weights: np.array, lower_bound: float = 0.0, -) -> torch.Tensor: +) -> np.array: """Transform attention matrix to the matrix of distances between tokens. Parameters ---------- attention_weights : torch.Tensor Attention matrixes of one sample (n_heads x n_tokens x n_tokens). - zero_out : Literal['prompt', 'response'] - Determines whether to zero out distances between prompt tokens or response tokens. - len_answer : int - Length of the response. - + Returns ------- - torch.Tensor + np.array Distance matrix. """ - n_tokens = attention_weights.shape[1] + attention_weights = torch.from_numpy(attention_weights).float() + n_tokens = attention_weights.shape[-1] distance_mx = 1 - torch.clamp( attention_weights, min=lower_bound ) # torch.where(attn_mx > lower_bound, attn_mx, 0.0) @@ -59,16 +53,7 @@ def transform_attention_scores_to_distances( distance_mx *= zero_diag.to(attention_weights.device).expand_as( distance_mx ) # torch.diag(torch.diag(distance_mx)) - distance_mx = torch.minimum(distance_mx.transpose(1, 2), distance_mx) - - if zero_out == "prompt": - len_prompt = n_tokens - len_answer - distance_mx[:, :len_prompt, :len_prompt] = 0 - elif zero_out == "response": - distance_mx[:, -len_answer:, -len_answer:] = 0 - else: - raise ValueError(f"Unsupported zero_out parameter: {zero_out}") - + distance_mx = torch.minimum(distance_mx.transpose(-1, -2), distance_mx) return distance_mx.cpu().numpy() @@ -104,11 +89,7 @@ class TopologicalDivergence(Estimator): def __init__( self, selected_heads: List[Tuple[int, int]] = None, - num_layers: int = None, - num_heads: int = None, - zero_output: Literal["prompt", "response"] = "prompt", n_jobs: int = 16, - critical_size: int = 768 ): """ Initializes TopologicalDivergence estimator. @@ -117,34 +98,18 @@ def __init__( selected_heads (List[Tuple[int, int]]): List of attention heads to calculate MTopDiv for. First integer is layer index, second is head index. If not provided or empty, all heads will be used. - zero_out Literal["prompt", "response"]: Determines whether to zero out distances between - prompt tokens or response tokens. - If not provided or empty, prompt tokens will be zeroed out. n_jobs (int): Number of jobs for parallel processing. Default: 16. - critical_size (int): Size threshold for parallel processing. - If the sequence length exceeds this size, n_jobs will be limited to 8. Default: 768. """ super().__init__(["greedy_tokens", "forwardpass_attention_weights"], "sequence") if selected_heads is None or len(selected_heads) == 0: - # If no specific heads are selected, use all heads - if num_layers is None or num_heads is None: - raise ValueError( - "If no specific heads are specified in 'selected_heads', " - "'num_layers' and 'num_heads' must be provided." - ) - selected_heads = [ - (layer, head) - for layer in range(num_layers) - for head in range(num_heads) - ] - selected_heads_dict = defaultdict(list) - for layer, head in selected_heads: - selected_heads_dict[layer].append(head) - self.selected_heads = selected_heads_dict - self.zero_out = zero_output + self.selected_heads = 'all' + else: + selected_heads_dict = defaultdict(list) + for layer, head in selected_heads: + selected_heads_dict[layer].append(head) + self.selected_heads = selected_heads_dict self.n_jobs = n_jobs - self.critical_size = critical_size def __str__(self): return "TopologicalDivergence" @@ -165,40 +130,45 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: """ responses = stats["greedy_tokens"] attention_weights_batch = stats["forwardpass_attention_weights"] - + batch_size = attention_weights_batch.shape[0] + + if self.selected_heads == 'all': + num_layers = attention_weights_batch.shape[1] + num_heads = attention_weights_batch.shape[2] + self.selected_heads = defaultdict(list) + for layer in range(num_layers): + self.selected_heads[layer] = list(range(num_heads)) + layers = sorted(self.selected_heads.keys()) - mtopdivs_batch = [] - for response, attention_weights in zip(responses, attention_weights_batch): - padding_length = np.isnan(attention_weights[0, 0, 0]).sum() - response_length = len(response) - mtopdivs_batch.append([]) - for layer in layers: - heads = self.selected_heads[layer] - if padding_length > 0: - selected_attention_weights = torch.from_numpy( - attention_weights[layer, heads, :-padding_length, :-padding_length] - ).float() - else: - selected_attention_weights = torch.from_numpy( - attention_weights[layer, heads, :, :] - ).float() - distance_matrices = transform_attention_scores_to_distances( - selected_attention_weights, self.zero_out, response_length - ) - if IS_PARALLEL_AVAILABLE: - if selected_attention_weights.shape[-1] <= self.critical_size: - n_jobs = min(self.n_jobs, 8) if self.n_jobs > 0 else 8 - else: - n_jobs = self.n_jobs - mtopdivs = Parallel(n_jobs=n_jobs)( - delayed(transform_distances_to_mtopdiv)(distance_matrice) - for distance_matrice in distance_matrices - ) - else: - mtopdivs = list(map( - transform_distances_to_mtopdiv, distance_matrices - )) - mtopdivs_batch[-1].extend(mtopdivs) - mtopdivs_batch = np.array(mtopdivs_batch, dtype=float) - - return np.mean(mtopdivs_batch, axis=1) + padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) + + distance_matrices_batch = transform_attention_scores_to_distances( + attention_weights_batch + ) + + def compute_mtopdiv(sample_id, layer, head): + distance_matrice = distance_matrices_batch[sample_id, layer, head] + padding_length = padding_lengths[sample_id] + response_length = len(responses[sample_id]) + if padding_length > 0: + distance_matrice = distance_matrice[:-padding_length, :-padding_length] + distance_matrice[:-response_length, :-response_length] = 0 + mtopdiv = transform_distances_to_mtopdiv(distance_matrice) + return mtopdiv + + if IS_PARALLEL_AVAILABLE: + mtopdivs = Parallel(n_jobs=self.n_jobs, backend="threading")( + delayed(compute_mtopdiv)(sample_id, layer, head) + for layer in layers + for head in self.selected_heads[layer] + for sample_id in range(batch_size) + ) + else: + mtopdivs = [ + compute_mtopdiv(sample_id, layer, head) + for layer in layers + for head in self.selected_heads[layer] + for sample_id in range(batch_size) + ] + mtopdivs = np.array(mtopdivs).reshape(batch_size, -1) + return np.mean(mtopdivs, axis=1) From 602e6ac5778adc98321e5b0a3384948a1343553b Mon Sep 17 00:00:00 2001 From: Andrei Volodichev Date: Wed, 18 Jun 2025 23:03:24 +0300 Subject: [PATCH 05/23] refactor parallelizm --- .../estimators/topological_divergence.py | 72 ++++++++----------- 1 file changed, 30 insertions(+), 42 deletions(-) diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 5fc8b023f..c661ad837 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -1,7 +1,6 @@ import os import warnings -from typing import Dict, List, Tuple, Literal -from collections import defaultdict +from typing import Dict, List, Tuple import numpy as np import torch @@ -13,7 +12,7 @@ except ImportError: IS_PARALLEL_AVAILABLE = False warnings.warn( - "Joblib is not installed. Parallel processing for MTopDivCalculator will not be available. " + "Joblib is not installed. Parallel processing for TopologicalDivergence call will not be available. " "Please install it via 'pip install joblib' if you want to use parallel processing." ) try: @@ -28,8 +27,7 @@ def transform_attention_scores_to_distances( - attention_weights: np.array, - lower_bound: float = 0.0, + attention_weights: np.array ) -> np.array: """Transform attention matrix to the matrix of distances between tokens. @@ -37,7 +35,7 @@ def transform_attention_scores_to_distances( ---------- attention_weights : torch.Tensor Attention matrixes of one sample (n_heads x n_tokens x n_tokens). - + Returns ------- np.array @@ -47,19 +45,19 @@ def transform_attention_scores_to_distances( attention_weights = torch.from_numpy(attention_weights).float() n_tokens = attention_weights.shape[-1] distance_mx = 1 - torch.clamp( - attention_weights, min=lower_bound - ) # torch.where(attn_mx > lower_bound, attn_mx, 0.0) + attention_weights, min=0.0 + ) zero_diag = torch.ones(n_tokens, n_tokens) - torch.eye(n_tokens) distance_mx *= zero_diag.to(attention_weights.device).expand_as( distance_mx - ) # torch.diag(torch.diag(distance_mx)) + ) distance_mx = torch.minimum(distance_mx.transpose(-1, -2), distance_mx) return distance_mx.cpu().numpy() def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: """ - Compute the MTopDiv (Metric Topological Diversity) score from a distance matrix. + Compute the MTopDiv (Manifold Topology Divergence) score from a distance matrix. This function calculates the sum of persistence intervals in the H₀ (zero-dimensional) persistent homology barcode, corresponding to the lifetimes of connected components @@ -102,13 +100,7 @@ def __init__( """ super().__init__(["greedy_tokens", "forwardpass_attention_weights"], "sequence") - if selected_heads is None or len(selected_heads) == 0: - self.selected_heads = 'all' - else: - selected_heads_dict = defaultdict(list) - for layer, head in selected_heads: - selected_heads_dict[layer].append(head) - self.selected_heads = selected_heads_dict + self.selected_heads = selected_heads self.n_jobs = n_jobs def __str__(self): @@ -132,43 +124,39 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: attention_weights_batch = stats["forwardpass_attention_weights"] batch_size = attention_weights_batch.shape[0] - if self.selected_heads == 'all': + if not self.selected_heads: num_layers = attention_weights_batch.shape[1] num_heads = attention_weights_batch.shape[2] - self.selected_heads = defaultdict(list) - for layer in range(num_layers): - self.selected_heads[layer] = list(range(num_heads)) - - layers = sorted(self.selected_heads.keys()) + self.selected_heads = [(layer, head) for layer in range(num_layers) for head in range(num_heads)] + padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) distance_matrices_batch = transform_attention_scores_to_distances( attention_weights_batch ) - def compute_mtopdiv(sample_id, layer, head): - distance_matrice = distance_matrices_batch[sample_id, layer, head] - padding_length = padding_lengths[sample_id] - response_length = len(responses[sample_id]) - if padding_length > 0: - distance_matrice = distance_matrice[:-padding_length, :-padding_length] - distance_matrice[:-response_length, :-response_length] = 0 - mtopdiv = transform_distances_to_mtopdiv(distance_matrice) - return mtopdiv - + def compute_mtopdiv(layer, head): + mtopdivs = [] + for sample_id in range(batch_size): + distance_matrice = distance_matrices_batch[sample_id, layer, head] + padding_length = padding_lengths[sample_id] + response_length = len(responses[sample_id]) + if padding_length > 0: + distance_matrice = distance_matrice[:-padding_length, :-padding_length] + distance_matrice[:-response_length, :-response_length] = 0 + mtopdiv = transform_distances_to_mtopdiv(distance_matrice) + mtopdivs.append(mtopdiv) + return np.array(mtopdivs, dtype=float) + if IS_PARALLEL_AVAILABLE: mtopdivs = Parallel(n_jobs=self.n_jobs, backend="threading")( - delayed(compute_mtopdiv)(sample_id, layer, head) - for layer in layers - for head in self.selected_heads[layer] - for sample_id in range(batch_size) + delayed(compute_mtopdiv)(layer, head) + for layer, head in self.selected_heads ) else: mtopdivs = [ - compute_mtopdiv(sample_id, layer, head) - for layer in layers - for head in self.selected_heads[layer] - for sample_id in range(batch_size) + compute_mtopdiv(layer, head) + for layer, head in self.selected_heads ] - mtopdivs = np.array(mtopdivs).reshape(batch_size, -1) + mtopdivs = np.stack(mtopdivs, axis=1) return np.mean(mtopdivs, axis=1) From 50b10be5f25a844ec3eaa453bec16459d1bcda7b Mon Sep 17 00:00:00 2001 From: Andrei Volodichev Date: Fri, 20 Jun 2025 14:46:32 +0300 Subject: [PATCH 06/23] add heads selection --- ...ical_divergence_with_heads_selection.ipynb | 209 ++++++++++++++++++ .../default_TrainMTopDivCalculator.py | 95 ++++++++ src/lm_polygraph/estimators/mtopdiv.py | 56 +++++ .../estimators/topological_divergence.py | 201 +++++++++-------- src/lm_polygraph/stat_calculators/__init__.py | 1 + .../stat_calculators/train_mtopdiv.py | 127 +++++++++++ 6 files changed, 591 insertions(+), 98 deletions(-) create mode 100644 examples/topological_divergence_with_heads_selection.ipynb create mode 100644 src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py create mode 100644 src/lm_polygraph/estimators/mtopdiv.py create mode 100644 src/lm_polygraph/stat_calculators/train_mtopdiv.py diff --git a/examples/topological_divergence_with_heads_selection.ipynb b/examples/topological_divergence_with_heads_selection.ipynb new file mode 100644 index 000000000..bd8036569 --- /dev/null +++ b/examples/topological_divergence_with_heads_selection.ipynb @@ -0,0 +1,209 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "bd1b5660", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/ahdr3w/.pyenv/versions/3.10.11/lib/python3.10/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": [ + "from time import time\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "\n", + "from lm_polygraph.estimators import TopologicalDivergence\n", + "\n", + "from lm_polygraph.utils import WhiteboxModel, UEManager, Dataset\n", + "from lm_polygraph.utils.builder_enviroment_stat_calculator import BuilderEnvironmentStatCalculator\n", + "from lm_polygraph.utils.factory_stat_calculator import StatCalculatorContainer\n", + "from lm_polygraph.utils.estimate_uncertainty import UncertaintyOutput\n", + "\n", + "from lm_polygraph.defaults.register_default_stat_calculators import (\n", + " register_default_stat_calculators,\n", + ")\n", + "from lm_polygraph.stat_calculators import TrainMTopDivCalculator\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "8d4148d7", + "metadata": {}, + "outputs": [], + "source": [ + "def estimate_uncertainty(\n", + " model, estimator, available_stat_calculators, input_text: str\n", + ") -> UncertaintyOutput:\n", + " man = UEManager(\n", + " Dataset([input_text], [\"\"], batch_size=1),\n", + " model,\n", + " [estimator],\n", + " available_stat_calculators=available_stat_calculators,\n", + " builder_env_stat_calc=BuilderEnvironmentStatCalculator(model),\n", + " generation_metrics=[],\n", + " ue_metrics=[],\n", + " processors=[],\n", + " ignore_exceptions=False,\n", + " verbose=False,\n", + " )\n", + " man()\n", + " ue = man.estimations[estimator.level, str(estimator)]\n", + " texts = man.stats.get(\"greedy_texts\", None)\n", + " tokens = man.stats.get(\"greedy_tokens\", None)\n", + " if tokens is not None and len(tokens) > 0:\n", + " # Remove last token, which is the end of the sequence token\n", + " # since we don't include it's uncertainty in the estimator's output\n", + " tokens = tokens[0][:-1]\n", + " return UncertaintyOutput(\n", + " ue[0], input_text, texts[0], tokens, model.model_path, str(estimator)\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "430d4c18", + "metadata": {}, + "outputs": [], + "source": [ + "class Config:\n", + " model_name = 'Qwen/Qwen2.5-0.5B-Instruct'\n", + " train_dataset = '../../coqa_Meta-Llama-3-8B-Instruct.csv'\n", + " context_column = 'context'\n", + " question_column = 'question'\n", + " prompt_column = 'prompt'\n", + " response_column = 'generated_answer'\n", + " label_column = 'hallucination'\n", + " batch_size = 1\n", + " subsample_train_dataset = 4\n", + " seed = 52\n", + " device_map = 'mps' \n", + " max_heads = 6\n", + " n_jobs = -1\n", + "\n", + "cfg = Config()\n", + "\n", + "base_model = AutoModelForCausalLM.from_pretrained(\n", + " cfg.model_name,\n", + " device_map=cfg.device_map,\n", + ")\n", + "tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)\n", + "model = WhiteboxModel(base_model, tokenizer)\n", + "available_stat_calculators = (\n", + " register_default_stat_calculators(\"Whitebox\")\n", + ")\n", + "sc = StatCalculatorContainer(\n", + " name=TrainMTopDivCalculator.__name__,\n", + " obj=TrainMTopDivCalculator,\n", + " builder=\"lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator\",\n", + " cfg=cfg,\n", + " dependencies=TrainMTopDivCalculator.meta_info()[1],\n", + " stats=TrainMTopDivCalculator.meta_info()[0],\n", + ")\n", + "available_stat_calculators.append(sc)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "5f927bc7", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + " 0%| | 0/4 [00:00 Iterable[Tuple[List[str], List[str], List[int]]]: + for i in range(0, len(self.x), self.batch_size): + yield ( + self.x[i : i + self.batch_size], + self.y[i : i + self.batch_size], + self.labels[i : i + self.batch_size], + ) + + def __len__(self) -> int: + return (len(self.x) + self.batch_size - 1) // self.batch_size + + def select(self, indices: List[int]): + self.x = [self.x[i] for i in indices] + self.y = [self.y[i] for i in indices] + self.labels = [self.labels[i] for i in indices] + return self + + def subsample(self, n_samples: int, seed: Optional[int] = None): + if n_samples >= len(self.x): + return self + + indices = list(range(len(self.x))) + if seed is not None: + np.random.seed(seed) + selected_indices = np.random.choice(indices, n_samples, replace=False) + return self.select(selected_indices.tolist()) + + @classmethod + def from_csv( + cls, + csv_path: str, + context_column: str, + question_column: str, + prompt_column: str, + response_column: str, + label_column: str, + batch_size: int, + ): + def assemble_query(row): + return row[prompt_column].format(row[context_column], row[question_column]) + + csv = pd.read_csv(csv_path) + x = csv.apply(assemble_query, axis=1).tolist() + y = csv[response_column].tolist() + labels = csv[label_column].tolist() + return cls(x, y, labels, batch_size) + + +def load_dataset(args): + log.info("=" * 100) + log.info("Loading train dataset...") + + train_dataset = Dataset.from_csv( + csv_path=args.train_dataset, + context_column=args.context_column, + question_column=args.question_column, + prompt_column=args.prompt_column, + response_column=args.response_column, + label_column=args.label_column, + batch_size=args.batch_size, + ) + + if args.subsample_train_dataset != -1: + train_dataset.subsample( + args.subsample_train_dataset, + seed=( + int(list(args.seed)[0]) if not isinstance(args.seed, int) else args.seed + ), + ) + + log.info("Done loading train data.") + return train_dataset + + +def load_stat_calculator(config, builder): + train_dataset = load_dataset(config) + return TrainMTopDivCalculator(train_dataset) diff --git a/src/lm_polygraph/estimators/mtopdiv.py b/src/lm_polygraph/estimators/mtopdiv.py new file mode 100644 index 000000000..0b93a19ee --- /dev/null +++ b/src/lm_polygraph/estimators/mtopdiv.py @@ -0,0 +1,56 @@ +import numpy as np +import torch + +try: + from ripser import ripser +except ImportError: + raise ImportError( + "Please install the 'ripser' package to use TopologicalDivergence estimator. " + "You can install it via 'pip install ripser'." + ) + + +def transform_attention_scores_to_distances( + attention_weights: torch.Tensor, +) -> torch.Tensor: + """Transform attention matrix to the matrix of distances between tokens. + + Parameters + ---------- + attention_weights : torch.Tensor + Attention matrixes of one sample (n_heads x n_tokens x n_tokens). + + Returns + ------- + np.array + Distance matrix. + + """ + attention_weights = attention_weights.float() + n_tokens = attention_weights.shape[-1] + distance_mx = 1 - torch.clamp(attention_weights, min=0.0) + zero_diag = torch.ones(n_tokens, n_tokens) - torch.eye(n_tokens) + zero_diag = zero_diag.to(attention_weights.device) + distance_mx *= zero_diag.expand_as(distance_mx) + distance_mx = torch.minimum(distance_mx.transpose(-1, -2), distance_mx) + return distance_mx + + +def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: + """ + Compute the MTopDiv (Manifold Topology Divergence) score from a distance matrix. + + This function calculates the sum of persistence intervals in the H₀ (zero-dimensional) + persistent homology barcode, corresponding to the lifetimes of connected components + in a Vietoris–Rips filtration. + + Parameters: + distance_mx (np.ndarray): A square, symmetric distance matrix. + + Returns: + float: Sum of finite H₀ barcode lengths (birth–death), representing topological diversity. + """ + barcodes = ripser(distance_mx, distance_matrix=True, maxdim=0)["dgms"] + if len(barcodes) > 0: + return barcodes[0][:-1, 1].sum() + return 0 diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index c661ad837..c96b904b9 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -1,78 +1,64 @@ import os import warnings -from typing import Dict, List, Tuple +from typing import Dict, List, Tuple, Optional import numpy as np import torch +from sklearn.metrics import roc_auc_score try: from joblib import Parallel, delayed - IS_PARALLEL_AVAILABLE = True + os.environ["TOKENIZERS_PARALLELISM"] = "false" + IS_PARALLEL_AVAILABLE = True except ImportError: IS_PARALLEL_AVAILABLE = False warnings.warn( "Joblib is not installed. Parallel processing for TopologicalDivergence call will not be available. " "Please install it via 'pip install joblib' if you want to use parallel processing." - ) -try: - from ripser import ripser -except ImportError: - raise ImportError( - "Please install the 'ripser' package to use TopologicalDivergence estimator. " - "You can install it via 'pip install ripser'." ) from .estimator import Estimator - - -def transform_attention_scores_to_distances( - attention_weights: np.array -) -> np.array: - """Transform attention matrix to the matrix of distances between tokens. - - Parameters - ---------- - attention_weights : torch.Tensor - Attention matrixes of one sample (n_heads x n_tokens x n_tokens). - - Returns - ------- - np.array - Distance matrix. - - """ - attention_weights = torch.from_numpy(attention_weights).float() - n_tokens = attention_weights.shape[-1] - distance_mx = 1 - torch.clamp( - attention_weights, min=0.0 - ) - zero_diag = torch.ones(n_tokens, n_tokens) - torch.eye(n_tokens) - distance_mx *= zero_diag.to(attention_weights.device).expand_as( - distance_mx - ) - distance_mx = torch.minimum(distance_mx.transpose(-1, -2), distance_mx) - return distance_mx.cpu().numpy() - - -def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: - """ - Compute the MTopDiv (Manifold Topology Divergence) score from a distance matrix. - - This function calculates the sum of persistence intervals in the H₀ (zero-dimensional) - persistent homology barcode, corresponding to the lifetimes of connected components - in a Vietoris–Rips filtration. - - Parameters: - distance_mx (np.ndarray): A square, symmetric distance matrix. - - Returns: - float: Sum of finite H₀ barcode lengths (birth–death), representing topological diversity. - """ - barcodes = ripser(distance_mx, distance_matrix=True, maxdim=0)['dgms'] - if len(barcodes) > 0: - return barcodes[0][:-1, 1].sum() - return 0 +from .mtopdiv import ( + transform_attention_scores_to_distances, + transform_distances_to_mtopdiv, +) + + +def get_mtopdivs( + heads: List[Tuple[int, int]], + length_responses: List[int], + attention_weights_batch: np.array, + n_jobs: Optional[float] = -1, +) -> np.ndarray: + batch_size = attention_weights_batch.shape[0] + padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) + attention_weights_batch = torch.from_numpy(attention_weights_batch) + distance_matrices_batch = transform_attention_scores_to_distances( + attention_weights_batch + ).numpy() + + def job(layer, head): + mtopdivs = [] + for sample_id in range(batch_size): + distance_matrice = distance_matrices_batch[sample_id, layer, head] + padding_length = padding_lengths[sample_id] + response_length = length_responses[sample_id] + if padding_length > 0: + distance_matrice = distance_matrice[:-padding_length, :-padding_length] + distance_matrice[:-response_length, :-response_length] = 0 + mtopdiv = transform_distances_to_mtopdiv(distance_matrice) + mtopdivs.append(mtopdiv) + return np.array(mtopdivs, dtype=float) + + if IS_PARALLEL_AVAILABLE: + mtopdivs = Parallel(n_jobs=n_jobs, backend="threading")( + delayed(job)(layer, head) for layer, head in heads + ) + else: + mtopdivs = [job(layer, head) for layer, head in heads] + mtopdivs = np.stack(mtopdivs, axis=1) + return mtopdivs class TopologicalDivergence(Estimator): @@ -84,10 +70,12 @@ class TopologicalDivergence(Estimator): Computes topological divergences between prompt and response attention graphs to identify hallucination-indicative heads. """ + def __init__( - self, - selected_heads: List[Tuple[int, int]] = None, - n_jobs: int = 16, + self, + heads: Optional[List[Tuple[int, int]]] = None, + max_heads: Optional[int] = 6, + n_jobs: int = -1, ): """ Initializes TopologicalDivergence estimator. @@ -98,14 +86,29 @@ def __init__( If not provided or empty, all heads will be used. n_jobs (int): Number of jobs for parallel processing. Default: 16. """ - super().__init__(["greedy_tokens", "forwardpass_attention_weights"], "sequence") + if not heads: + calculators = [ + "train_labels", + "train_mtopdivs", + "greedy_tokens", + "forwardpass_attention_weights", + ] + else: + calculators = ["greedy_tokens", "forwardpass_attention_weights"] - self.selected_heads = selected_heads - self.n_jobs = n_jobs + super().__init__(calculators, "sequence") + self._selected_heads = heads + self._best_heads = None + self._max_heads = max_heads + self._n_jobs = n_jobs def __str__(self): return "TopologicalDivergence" + @property + def best_heads(self): + return self._best_heads + def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: """ Calculates sequence-wise MTopDiv scores for selected heads of attention masks. @@ -120,43 +123,45 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: np.ndarray: float uncertainty for each sample in input statistics. Higher values indicate more uncertain samples. """ - responses = stats["greedy_tokens"] + length_responses = list(map(len, stats["greedy_tokens"])) attention_weights_batch = stats["forwardpass_attention_weights"] - batch_size = attention_weights_batch.shape[0] - - if not self.selected_heads: - num_layers = attention_weights_batch.shape[1] - num_heads = attention_weights_batch.shape[2] - self.selected_heads = [(layer, head) for layer in range(num_layers) for head in range(num_heads)] - padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) + if self._selected_heads is None and self._best_heads is None: + train_mtopdivs = stats["train_mtopdivs"] + train_labels = stats["train_labels"] + best_heads = self._select_heads( + train_mtopdivs, + train_labels, + ) + _, num_layers, num_heads, _, _ = attention_weights_batch.shape + best_heads = np.unravel_index(best_heads, (num_layers, num_heads)) + best_heads = list(zip(best_heads[0], best_heads[1])) + self._best_heads = best_heads + super().__init__( + ["greedy_tokens", "forwardpass_attention_weights"], + "sequence", + ) - distance_matrices_batch = transform_attention_scores_to_distances( - attention_weights_batch + mtopdivs = get_mtopdivs( + self._best_heads if self._best_heads else self._selected_heads, + length_responses, + attention_weights_batch, + n_jobs=self._n_jobs, ) - def compute_mtopdiv(layer, head): - mtopdivs = [] - for sample_id in range(batch_size): - distance_matrice = distance_matrices_batch[sample_id, layer, head] - padding_length = padding_lengths[sample_id] - response_length = len(responses[sample_id]) - if padding_length > 0: - distance_matrice = distance_matrice[:-padding_length, :-padding_length] - distance_matrice[:-response_length, :-response_length] = 0 - mtopdiv = transform_distances_to_mtopdiv(distance_matrice) - mtopdivs.append(mtopdiv) - return np.array(mtopdivs, dtype=float) - - if IS_PARALLEL_AVAILABLE: - mtopdivs = Parallel(n_jobs=self.n_jobs, backend="threading")( - delayed(compute_mtopdiv)(layer, head) - for layer, head in self.selected_heads - ) - else: - mtopdivs = [ - compute_mtopdiv(layer, head) - for layer, head in self.selected_heads - ] - mtopdivs = np.stack(mtopdivs, axis=1) return np.mean(mtopdivs, axis=1) + + def _select_heads(self, scores, labels): + grounded_scores, hal_scores = scores[labels == 0], scores[labels == 1] + deltas = hal_scores.mean(0) - grounded_scores.mean(0) + heads = sorted(range(len(deltas)), key=lambda x: deltas[x], reverse=True) + + best_auroc, n_opt = 0, 0 + for n in range(1, self._max_heads + 1): + n_best_heads = heads[:n] + predictions = scores[:, n_best_heads].mean(axis=1) + roc_auc = roc_auc_score(labels, predictions) + if roc_auc > best_auroc: + best_auroc = roc_auc + n_opt = n + return heads[:n_opt] diff --git a/src/lm_polygraph/stat_calculators/__init__.py b/src/lm_polygraph/stat_calculators/__init__.py index 72a5141bd..737e350cf 100644 --- a/src/lm_polygraph/stat_calculators/__init__.py +++ b/src/lm_polygraph/stat_calculators/__init__.py @@ -42,3 +42,4 @@ from .extract_claims import ClaimsExtractor from .infer_causal_lm_calculator import InferCausalLMCalculator from .semantic_classes import SemanticClassesCalculator +from .train_mtopdiv import TrainMTopDivCalculator diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py new file mode 100644 index 000000000..e1fb7c96c --- /dev/null +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -0,0 +1,127 @@ +import gc +import torch +import numpy as np +from tqdm import tqdm + +from typing import Dict, List, Tuple + +from .stat_calculator import StatCalculator +from lm_polygraph.utils.model import WhiteboxModel +from .attention_forward_pass import AttentionForwardPassCalculator +from ..estimators.mtopdiv import ( + transform_attention_scores_to_distances, + transform_distances_to_mtopdiv, +) + + +class TrainMTopDivCalculator(StatCalculator): + @staticmethod + def meta_info() -> Tuple[List[str], List[str]]: + """ + Returns the statistics and dependencies for the calculator. + """ + + return [ + "train_labels", + "train_mtopdivs", + ], [] + + def __init__(self, train_dataset=None, *args, **kwargs): + super().__init__() + self.train_dataset = train_dataset + self.base_calculator = AttentionForwardPassCalculator() + + def _calculate_attention_forward_pass( + self, + prompts: List[str], + responses: List[str], + model: WhiteboxModel, + ) -> Tuple[List[np.ndarray], np.ndarray]: + + if model.tokenizer.chat_template is not None: + chats = [ + [{"role": "user", "content": p}, {"role": "assistant", "content": r}] + for p, r in zip(prompts, responses) + ] + else: + chats = [p + r for p, r in zip(prompts, responses)] + + with torch.no_grad(): + batch = model.tokenize(chats).to(model.device()) + attns = model.model( + **batch, + output_attentions=True, + ).attentions + mask = batch["attention_mask"] + mask = mask[:, None, None, None, :] & mask[:, None, None, :, None] + attns = torch.stack(attns, dim=1) + attns = attns.masked_fill(mask == 0, float("nan")).cpu() + + torch.cuda.empty_cache() + gc.collect() + + return attns + + def __call__( + self, + dependencies: Dict[str, np.array], + texts: List[str], + model: WhiteboxModel, + max_new_tokens: int = 100, + ) -> Dict[str, np.ndarray]: + + train_mtopdivs = [] + train_labels = [] + + for inp_texts, target_texts, labels in tqdm(self.train_dataset): + forwardpass_attention_weights = self._calculate_attention_forward_pass( + inp_texts, + target_texts, + model, + ) + + length_responses = model.tokenizer( + target_texts, + add_special_tokens=False, + padding=True, + return_tensors="pt", + ) + length_responses = length_responses["attention_mask"].sum(dim=1).tolist() + + batch_size, layers, heads, seq_len, _ = forwardpass_attention_weights.shape + paddings = np.isnan(forwardpass_attention_weights[:, 0, 0, 0]).sum(axis=-1) + distance_matrices_batch = transform_attention_scores_to_distances( + forwardpass_attention_weights + ) + distance_matrices_batch = distance_matrices_batch.reshape( + batch_size, layers * heads, seq_len, seq_len + ).numpy() + + def job(head): + mtopdivs = [] + for sample_id in range(batch_size): + distance_matrice = distance_matrices_batch[sample_id, head] + padding = paddings[sample_id] + response_length = length_responses[sample_id] + if padding > 0: + distance_matrice = distance_matrice[:-padding, :-padding] + distance_matrice[:-response_length, :-response_length] = 0 + mtopdiv = transform_distances_to_mtopdiv(distance_matrice) + mtopdivs.append(mtopdiv) + return np.array(mtopdivs, dtype=float) + + train_mtopdivs.append( + np.stack( + [job(head) for head in range(heads * layers)], + axis=1, + ) + ) + train_labels += labels + + train_mtopdivs = np.concatenate(train_mtopdivs, axis=0) + train_labels = np.array(train_labels) + + return { + "train_mtopdivs": train_mtopdivs, + "train_labels": train_labels, + } From 946ad96d579f37b6d068f5ea71a64ab704cb25a0 Mon Sep 17 00:00:00 2001 From: Andrei Volodichev Date: Fri, 20 Jun 2025 20:19:56 +0300 Subject: [PATCH 07/23] small fix --- .../topological_divergence_with_heads_selection.ipynb | 6 +++--- src/lm_polygraph/estimators/topological_divergence.py | 9 +++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/examples/topological_divergence_with_heads_selection.ipynb b/examples/topological_divergence_with_heads_selection.ipynb index bd8036569..0c306db47 100644 --- a/examples/topological_divergence_with_heads_selection.ipynb +++ b/examples/topological_divergence_with_heads_selection.ipynb @@ -122,7 +122,7 @@ "output_type": "stream", "text": [ " 0%| | 0/4 [00:00 np.ndarray: best_heads = np.unravel_index(best_heads, (num_layers, num_heads)) best_heads = list(zip(best_heads[0], best_heads[1])) self._best_heads = best_heads - super().__init__( - ["greedy_tokens", "forwardpass_attention_weights"], - "sequence", - ) + + # After selecting heads once, we drop train stats to prevent recomputing them in later runs. + # This assumes the manager object is reinitialized before the next estimator call. + # Alternatively, a new estimator instance can be created with the selected heads. + self.stats_dependencies = ["greedy_tokens", "forwardpass_attention_weights"] mtopdivs = get_mtopdivs( self._best_heads if self._best_heads else self._selected_heads, From c41abb42009ab87ad9e57f75e3eca4e354e86314 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Tue, 24 Jun 2025 16:24:43 +0300 Subject: [PATCH 08/23] add hydra configs --- .../default_TopologicalDivergence.yaml | 5 + .../estimators/default_estimators.yaml | 2 +- .../default_TrainMTopDivCalculator.yaml | 19 +++ .../stat_calculators/default_calculators.yaml | 3 +- .../configs/topological_divergence_coqa.yaml | 59 ++++++++ ...ical_divergence_with_heads_selection.ipynb | 30 ++-- .../default_TrainMTopDivCalculator.py | 140 +++++++++--------- src/lm_polygraph/estimators/mtopdiv.py | 126 +++++++++++++++- .../estimators/topological_divergence.py | 99 ++----------- .../stat_calculators/train_mtopdiv.py | 138 ++++++++++------- src/lm_polygraph/utils/factory_estimator.py | 1 + 11 files changed, 390 insertions(+), 232 deletions(-) create mode 100644 examples/configs/estimators/default_TopologicalDivergence.yaml create mode 100644 examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml create mode 100644 examples/configs/topological_divergence_coqa.yaml diff --git a/examples/configs/estimators/default_TopologicalDivergence.yaml b/examples/configs/estimators/default_TopologicalDivergence.yaml new file mode 100644 index 000000000..6b87c9f29 --- /dev/null +++ b/examples/configs/estimators/default_TopologicalDivergence.yaml @@ -0,0 +1,5 @@ +- name: TopologicalDivergence + cfg: + heads: null + max_heads: '${max_heads}' + n_jobs: '${n_jobs}' \ No newline at end of file diff --git a/examples/configs/estimators/default_estimators.yaml b/examples/configs/estimators/default_estimators.yaml index 6b507add1..4b8de1076 100644 --- a/examples/configs/estimators/default_estimators.yaml +++ b/examples/configs/estimators/default_estimators.yaml @@ -87,4 +87,4 @@ - name: AttentionScore cfg: layer: 16 - gen_only: False \ No newline at end of file + gen_only: False diff --git a/examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml b/examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml new file mode 100644 index 000000000..cc374f163 --- /dev/null +++ b/examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml @@ -0,0 +1,19 @@ +- name: TrainMTopDivCalculator + builder: lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator + cfg: + heads_extraction_priority: cache + cache_path: '${model_heads_cache}' + max_heads: '${max_heads}' + n_jobs: '${n_jobs}' + train_data_path: '${train_dataset.train_data_path}' + context_column: '${train_dataset.context}' + question_column: '${train_dataset.question}' + prompt_column: '${train_dataset.prompt}' + response_column: '${train_dataset.generated_answer}' + label_column: '${train_dataset.hallucination}' + batch_size: '${batch_size}' + subsample_train_dataset: '${train_dataset.subsample}' + seed: '${seed}' + stats: + - "topological_divergence_heads" + dependencies: \ No newline at end of file diff --git a/examples/configs/stat_calculators/default_calculators.yaml b/examples/configs/stat_calculators/default_calculators.yaml index 265648300..3ce7d2221 100644 --- a/examples/configs/stat_calculators/default_calculators.yaml +++ b/examples/configs/stat_calculators/default_calculators.yaml @@ -27,4 +27,5 @@ - "train_embeddings" - "background_train_embeddings" - "train_greedy_log_likelihoods" - dependencies: \ No newline at end of file + dependencies: + diff --git a/examples/configs/topological_divergence_coqa.yaml b/examples/configs/topological_divergence_coqa.yaml new file mode 100644 index 000000000..86a53afb5 --- /dev/null +++ b/examples/configs/topological_divergence_coqa.yaml @@ -0,0 +1,59 @@ +hydra: + run: + dir: ${cache_path}/${task}/${model.path}/${dataset}/${now:%Y-%m-%d}/${now:%H-%M-%S} + +defaults: + - estimators: default_TopologicalDivergence + + - stat_calculators: default_TrainMTopDivCalculator + - base_processing_coqa + - _self_ +model: + path: Qwen/Qwen2.5-0.5B-Instruct + type: CausalLM + path_to_load_script: model/default_causal.py + load_model_args: + device_map: auto + load_tokenizer_args: {} + +cache_path: ./workdir/output +save_path: '${hydra:run.dir}' + +task: qa +instruct: false + +dataset: ['LM-polygraph/coqa', 'continuation'] +text_column: input +label_column: output +train_split: train +eval_split: test +n_shot: 0 +few_shot_split: train +few_shot_prompt: null +trust_remote_code: false +size: null + +max_new_tokens: 20 +load_from_disk: false +generation_params: + generate_until: + - "\n" +subsample_eval_dataset: -1 +batch_size: 1 + +generation_metrics: null +ignore_exceptions: false +seed: + - 1 + +model_heads_cache: "../../../model_heads_cache.yaml" +max_heads: 6 +n_jobs: -1 +train_dataset: + train_data_path: "../../../coqa_Meta-Llama-3-8B-Instruct.csv" + context_column: context + question_column: question + prompt_column: prompt + response_column: generated_answer + label_column: hallucination + subsample: 100 diff --git a/examples/topological_divergence_with_heads_selection.ipynb b/examples/topological_divergence_with_heads_selection.ipynb index 0c306db47..21f3f65c5 100644 --- a/examples/topological_divergence_with_heads_selection.ipynb +++ b/examples/topological_divergence_with_heads_selection.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "bd1b5660", "metadata": {}, "outputs": [ @@ -69,14 +69,16 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "430d4c18", "metadata": {}, "outputs": [], "source": [ "class Config:\n", " model_name = 'Qwen/Qwen2.5-0.5B-Instruct'\n", - " train_dataset = '../../coqa_Meta-Llama-3-8B-Instruct.csv'\n", + " mtopdiv_train_data_path = '../../coqa_Meta-Llama-3-8B-Instruct.csv'\n", + " mtopdiv_heads_extraction_priority = \"config\"\n", + " mtopdiv_heads_path = \"../../config.yaml\"\n", " context_column = 'context'\n", " question_column = 'question'\n", " prompt_column = 'prompt'\n", @@ -86,9 +88,10 @@ " subsample_train_dataset = 4\n", " seed = 52\n", " device_map = 'mps' \n", - " max_heads = 6\n", + " mtopdiv_max_heads = 6\n", " n_jobs = -1\n", "\n", + "\n", "cfg = Config()\n", "\n", "base_model = AutoModelForCausalLM.from_pretrained(\n", @@ -96,7 +99,7 @@ " device_map=cfg.device_map,\n", ")\n", "tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)\n", - "model = WhiteboxModel(base_model, tokenizer)\n", + "model = WhiteboxModel(base_model, tokenizer, model_path=cfg.model_name)\n", "available_stat_calculators = (\n", " register_default_stat_calculators(\"Whitebox\")\n", ")\n", @@ -121,17 +124,16 @@ "name": "stderr", "output_type": "stream", "text": [ - " 0%| | 0/4 [00:00 Iterable[Tuple[List[str], List[str], List[int]]]: - for i in range(0, len(self.x), self.batch_size): - yield ( - self.x[i : i + self.batch_size], - self.y[i : i + self.batch_size], - self.labels[i : i + self.batch_size], - ) - - def __len__(self) -> int: - return (len(self.x) + self.batch_size - 1) // self.batch_size - - def select(self, indices: List[int]): - self.x = [self.x[i] for i in indices] - self.y = [self.y[i] for i in indices] - self.labels = [self.labels[i] for i in indices] - return self - - def subsample(self, n_samples: int, seed: Optional[int] = None): - if n_samples >= len(self.x): - return self - - indices = list(range(len(self.x))) - if seed is not None: - np.random.seed(seed) - selected_indices = np.random.choice(indices, n_samples, replace=False) - return self.select(selected_indices.tolist()) - @classmethod - def from_csv( - cls, +class MTopDivHoldoutDataset: + def __init__( + self, csv_path: str, context_column: str, question_column: str, @@ -53,43 +17,77 @@ def from_csv( response_column: str, label_column: str, batch_size: int, + subsample_train_dataset: int, + seed, ): - def assemble_query(row): - return row[prompt_column].format(row[context_column], row[question_column]) + self.csv_path = csv_path + self.context_column = context_column + self.question_column = question_column + self.prompt_column = prompt_column + self.response_column = response_column + self.label_column = label_column + self.batch_size = batch_size + self.subsample_train_dataset = subsample_train_dataset + self.seed = seed - csv = pd.read_csv(csv_path) - x = csv.apply(assemble_query, axis=1).tolist() - y = csv[response_column].tolist() - labels = csv[label_column].tolist() - return cls(x, y, labels, batch_size) + self.prompts = [] + self.responses = [] + self.labels = [] + def from_csv(self): + def assemble_query(row): + return row[self.prompt_column].format(row[self.context_column], row[self.question_column]) -def load_dataset(args): - log.info("=" * 100) - log.info("Loading train dataset...") + df = pd.read_csv(self.csv_path) + self.prompts = df.apply(assemble_query, axis=1).tolist() + self.responses = df[self.response_column].tolist() + self.labels = df[self.label_column].tolist() - train_dataset = Dataset.from_csv( - csv_path=args.train_dataset, - context_column=args.context_column, - question_column=args.question_column, - prompt_column=args.prompt_column, - response_column=args.response_column, - label_column=args.label_column, - batch_size=args.batch_size, - ) + def __len__(self): + return len(self.prompts) - if args.subsample_train_dataset != -1: - train_dataset.subsample( - args.subsample_train_dataset, - seed=( - int(list(args.seed)[0]) if not isinstance(args.seed, int) else args.seed - ), - ) + def __getitem__(self, idx): + return self.prompts[idx], self.responses[idx], self.labels[idx] - log.info("Done loading train data.") - return train_dataset + def select(self, indices: List[int]): + self.prompts = [self.prompts[i] for i in indices] + self.responses = [self.responses[i] for i in indices] + self.labels = [self.labels[i] for i in indices] + return self + def subsample(self): + if self.subsample_train_dataset >= len(self): + return self + rng = np.random.default_rng(self.seed) + selected_indices = rng.choice( + len(self), + size=self.subsample_train_dataset, + replace=False, + ).tolist() + return self.select(selected_indices) def load_stat_calculator(config, builder): - train_dataset = load_dataset(config) - return TrainMTopDivCalculator(train_dataset) + priority = config.heads_extraction_priority + cache_path = config.cache_path + max_heads = config.max_heads + n_jobs = config.n_jobs + + train_dataset = MTopDivHoldoutDataset( + csv_path=config.train_data_path, + context_column=config.context_column, + question_column=config.question_column, + prompt_column=config.prompt_column, + response_column=config.response_column, + label_column=config.label_column, + batch_size=config.batch_size, + subsample_train_dataset=config.subsample_train_dataset, + seed=config.seed, + ) + + return TrainMTopDivCalculator( + priority, + train_dataset, + cache_path, + max_heads, + n_jobs, + ) diff --git a/src/lm_polygraph/estimators/mtopdiv.py b/src/lm_polygraph/estimators/mtopdiv.py index 0b93a19ee..6258c51a4 100644 --- a/src/lm_polygraph/estimators/mtopdiv.py +++ b/src/lm_polygraph/estimators/mtopdiv.py @@ -1,5 +1,11 @@ +import os +import yaml +import warnings import numpy as np +import pandas as pd import torch +from typing import List, Tuple, Optional +import ast try: from ripser import ripser @@ -9,6 +15,18 @@ "You can install it via 'pip install ripser'." ) +try: + from joblib import Parallel, delayed + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + IS_PARALLEL_AVAILABLE = True +except ImportError: + IS_PARALLEL_AVAILABLE = False + warnings.warn( + "Joblib is not installed. Parallel processing for TopologicalDivergence call will not be available. " + "Please install it via 'pip install joblib' if you want to use parallel processing." + ) + def transform_attention_scores_to_distances( attention_weights: torch.Tensor, @@ -29,8 +47,10 @@ def transform_attention_scores_to_distances( attention_weights = attention_weights.float() n_tokens = attention_weights.shape[-1] distance_mx = 1 - torch.clamp(attention_weights, min=0.0) - zero_diag = torch.ones(n_tokens, n_tokens) - torch.eye(n_tokens) - zero_diag = zero_diag.to(attention_weights.device) + zero_diag = ( + torch.ones(n_tokens, n_tokens, device=attention_weights.device) - + torch.eye(n_tokens, device=attention_weights.device) + ) distance_mx *= zero_diag.expand_as(distance_mx) distance_mx = torch.minimum(distance_mx.transpose(-1, -2), distance_mx) return distance_mx @@ -54,3 +74,105 @@ def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: if len(barcodes) > 0: return barcodes[0][:-1, 1].sum() return 0 + +def get_mtopdivs( + heads: List[Tuple[int, int]], + length_responses: List[int], + attention_weights_batch: torch.Tensor, + n_jobs: Optional[float] = -1, +) -> np.ndarray: + batch_size = attention_weights_batch.shape[0] + padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) + distance_matrices_batch = transform_attention_scores_to_distances( + attention_weights_batch + ).numpy() + + def job(layer, head): + mtopdivs = [] + for sample_id in range(batch_size): + distance_matrix = distance_matrices_batch[sample_id, layer, head] + padding_length = padding_lengths[sample_id] + response_length = length_responses[sample_id] + if padding_length > 0: + distance_matrix = distance_matrix[:-padding_length, :-padding_length] + distance_matrix[:-response_length, :-response_length] = 0 + mtopdiv = transform_distances_to_mtopdiv(distance_matrix) / response_length + mtopdivs.append(mtopdiv) + return np.array(mtopdivs, dtype=float) + + if IS_PARALLEL_AVAILABLE: + mtopdivs = Parallel(n_jobs=n_jobs, backend="threading")( + delayed(job)(layer, head) for layer, head in heads + ) + else: + mtopdivs = [job(layer, head) for layer, head in heads] + mtopdivs = np.stack(mtopdivs, axis=1) + return mtopdivs + +def load_model_heads( + cache_path: Optional[str], + model_path: str +) -> Optional[List[Tuple[int, int]]]: + """ + Load model heads. + + Parameters + ---------- + cache_path : Optional[str] + Path to the YAML cache file. + model_path : str + Unique path to the model. + + Returns + ------- + Optional[List[Tuple[int, int]]] + List of heads for the model if found, else None. + """ + if cache_path and os.path.isfile(cache_path): + try: + with open(cache_path, "r") as f: + config = yaml.safe_load(f) + models = config.get("models", []) + for model_entry in models: + if isinstance(model_entry, dict) and model_path in model_entry: + return ast.literal_eval(model_entry[model_path]) + + print(f"Model '{model_path}' not found in cache.") + except Exception as e: + print(f"Failed to load heads from cache: {e}") + return None + +def save_model_heads( + cache_path: str, + model_path: str, + heads: List[Tuple[int, int]] +) -> None: + """ + Save a list of heads for a model. + + Parameters + ---------- + cache_path : str + Path to the YAML file to save the model heads. + model_path : str + Unique path to the model. + heads : List[Tuple[int, int]] + List of heads to save. + """ + if os.path.isfile(cache_path): + with open(cache_path, "r") as f: + config = yaml.safe_load(f) or {} + else: + config = {} + models = config.get("models", []) + + for entry in models: + if isinstance(entry, dict) and model_path in entry: + return + + models.append({model_path: f"{heads}"}) + config["models"] = models + + with open(cache_path, "w") as f: + yaml.safe_dump(config, f, sort_keys=False) + print(f"Saved heads for '{model_path}' to {cache_path}") diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index b1454766f..185692dd0 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -1,64 +1,9 @@ -import os -import warnings from typing import Dict, List, Tuple, Optional import numpy as np import torch -from sklearn.metrics import roc_auc_score - - -try: - from joblib import Parallel, delayed - - os.environ["TOKENIZERS_PARALLELISM"] = "false" - IS_PARALLEL_AVAILABLE = True -except ImportError: - IS_PARALLEL_AVAILABLE = False - warnings.warn( - "Joblib is not installed. Parallel processing for TopologicalDivergence call will not be available. " - "Please install it via 'pip install joblib' if you want to use parallel processing." - ) from .estimator import Estimator -from .mtopdiv import ( - transform_attention_scores_to_distances, - transform_distances_to_mtopdiv, -) - - -def get_mtopdivs( - heads: List[Tuple[int, int]], - length_responses: List[int], - attention_weights_batch: np.array, - n_jobs: Optional[float] = -1, -) -> np.ndarray: - batch_size = attention_weights_batch.shape[0] - padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) - attention_weights_batch = torch.from_numpy(attention_weights_batch) - distance_matrices_batch = transform_attention_scores_to_distances( - attention_weights_batch - ).numpy() - - def job(layer, head): - mtopdivs = [] - for sample_id in range(batch_size): - distance_matrice = distance_matrices_batch[sample_id, layer, head] - padding_length = padding_lengths[sample_id] - response_length = length_responses[sample_id] - if padding_length > 0: - distance_matrice = distance_matrice[:-padding_length, :-padding_length] - distance_matrice[:-response_length, :-response_length] = 0 - mtopdiv = transform_distances_to_mtopdiv(distance_matrice) - mtopdivs.append(mtopdiv) - return np.array(mtopdivs, dtype=float) - - if IS_PARALLEL_AVAILABLE: - mtopdivs = Parallel(n_jobs=n_jobs, backend="threading")( - delayed(job)(layer, head) for layer, head in heads - ) - else: - mtopdivs = [job(layer, head) for layer, head in heads] - mtopdivs = np.stack(mtopdivs, axis=1) - return mtopdivs +from .mtopdiv import get_mtopdivs class TopologicalDivergence(Estimator): @@ -88,8 +33,7 @@ def __init__( """ if not heads: calculators = [ - "train_labels", - "train_mtopdivs", + "topological_divergence_heads", "greedy_tokens", "forwardpass_attention_weights", ] @@ -97,8 +41,7 @@ def __init__( calculators = ["greedy_tokens", "forwardpass_attention_weights"] super().__init__(calculators, "sequence") - self._selected_heads = heads - self._best_heads = None + self._heads = heads self._max_heads = max_heads self._n_jobs = n_jobs @@ -106,8 +49,8 @@ def __str__(self): return "TopologicalDivergence" @property - def best_heads(self): - return self._best_heads + def heads(self): + return self._heads def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: """ @@ -126,25 +69,18 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: length_responses = list(map(len, stats["greedy_tokens"])) attention_weights_batch = stats["forwardpass_attention_weights"] - if self._selected_heads is None and self._best_heads is None: - train_mtopdivs = stats["train_mtopdivs"] - train_labels = stats["train_labels"] - best_heads = self._select_heads( - train_mtopdivs, - train_labels, - ) - _, num_layers, num_heads, _, _ = attention_weights_batch.shape - best_heads = np.unravel_index(best_heads, (num_layers, num_heads)) - best_heads = list(zip(best_heads[0], best_heads[1])) - self._best_heads = best_heads + if self._heads is None: + best_heads = stats["topological_divergence_heads"] + self._heads = best_heads # After selecting heads once, we drop train stats to prevent recomputing them in later runs. # This assumes the manager object is reinitialized before the next estimator call. # Alternatively, a new estimator instance can be created with the selected heads. self.stats_dependencies = ["greedy_tokens", "forwardpass_attention_weights"] + attention_weights_batch = torch.from_numpy(attention_weights_batch) mtopdivs = get_mtopdivs( - self._best_heads if self._best_heads else self._selected_heads, + self._heads, length_responses, attention_weights_batch, n_jobs=self._n_jobs, @@ -152,17 +88,4 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: return np.mean(mtopdivs, axis=1) - def _select_heads(self, scores, labels): - grounded_scores, hal_scores = scores[labels == 0], scores[labels == 1] - deltas = hal_scores.mean(0) - grounded_scores.mean(0) - heads = sorted(range(len(deltas)), key=lambda x: deltas[x], reverse=True) - - best_auroc, n_opt = 0, 0 - for n in range(1, self._max_heads + 1): - n_best_heads = heads[:n] - predictions = scores[:, n_best_heads].mean(axis=1) - roc_auc = roc_auc_score(labels, predictions) - if roc_auc > best_auroc: - best_auroc = roc_auc - n_opt = n - return heads[:n_opt] + diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py index e1fb7c96c..f5c741f41 100644 --- a/src/lm_polygraph/stat_calculators/train_mtopdiv.py +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -1,16 +1,19 @@ import gc import torch +from torch.utils.data import DataLoader import numpy as np +import pandas as pd +from sklearn.metrics import roc_auc_score from tqdm import tqdm - -from typing import Dict, List, Tuple +from typing import Dict, List, Tuple, Literal +from itertools import product from .stat_calculator import StatCalculator from lm_polygraph.utils.model import WhiteboxModel -from .attention_forward_pass import AttentionForwardPassCalculator -from ..estimators.mtopdiv import ( - transform_attention_scores_to_distances, - transform_distances_to_mtopdiv, +from lm_polygraph.estimators.mtopdiv import ( + get_mtopdivs, + load_model_heads, + save_model_heads, ) @@ -21,18 +24,25 @@ def meta_info() -> Tuple[List[str], List[str]]: Returns the statistics and dependencies for the calculator. """ - return [ - "train_labels", - "train_mtopdivs", - ], [] + return ["topological_divergence_heads"], [] - def __init__(self, train_dataset=None, *args, **kwargs): + def __init__(self, + priority: Literal["train", "cache"] = "cache", + train_dataset: "MTopDivHoldoutDataset" = None, + cache_path: str = None, + max_heads: int = 6, + n_jobs: int = -1): super().__init__() + if train_dataset is None and cache_path is None: + raise Exception("Either train_dataset or cache_path must be provided.") + self.priority = priority + self.cache_path = cache_path self.train_dataset = train_dataset - self.base_calculator = AttentionForwardPassCalculator() + self.max_heads = max_heads + self.n_jobs = n_jobs - def _calculate_attention_forward_pass( - self, + @staticmethod + def calculate_attention_forward_pass( prompts: List[str], responses: List[str], model: WhiteboxModel, @@ -61,25 +71,32 @@ def _calculate_attention_forward_pass( gc.collect() return attns - - def __call__( - self, - dependencies: Dict[str, np.array], - texts: List[str], - model: WhiteboxModel, - max_new_tokens: int = 100, - ) -> Dict[str, np.ndarray]: - + + def select_heads(self, scores, labels): + grounded_scores, hal_scores = scores[labels == 0], scores[labels == 1] + deltas = hal_scores.mean(0) - grounded_scores.mean(0) + heads = sorted(range(len(deltas)), key=lambda x: deltas[x], reverse=True) + + best_auroc, n_opt = 0, 0 + for n in range(1, self.max_heads + 1): + n_best_heads = heads[:n] + predictions = scores[:, n_best_heads].mean(axis=1) + roc_auc = roc_auc_score(labels, predictions) + if roc_auc > best_auroc: + best_auroc = roc_auc + n_opt = n + return heads[:n_opt] + + def train(self, model, train_dataloader): train_mtopdivs = [] train_labels = [] - for inp_texts, target_texts, labels in tqdm(self.train_dataset): - forwardpass_attention_weights = self._calculate_attention_forward_pass( + for inp_texts, target_texts, labels in tqdm(train_dataloader): + forwardpass_attention_weights = self.calculate_attention_forward_pass( inp_texts, target_texts, model, ) - length_responses = model.tokenizer( target_texts, add_special_tokens=False, @@ -87,41 +104,52 @@ def __call__( return_tensors="pt", ) length_responses = length_responses["attention_mask"].sum(dim=1).tolist() - - batch_size, layers, heads, seq_len, _ = forwardpass_attention_weights.shape - paddings = np.isnan(forwardpass_attention_weights[:, 0, 0, 0]).sum(axis=-1) - distance_matrices_batch = transform_attention_scores_to_distances( - forwardpass_attention_weights - ) - distance_matrices_batch = distance_matrices_batch.reshape( - batch_size, layers * heads, seq_len, seq_len - ).numpy() - - def job(head): - mtopdivs = [] - for sample_id in range(batch_size): - distance_matrice = distance_matrices_batch[sample_id, head] - padding = paddings[sample_id] - response_length = length_responses[sample_id] - if padding > 0: - distance_matrice = distance_matrice[:-padding, :-padding] - distance_matrice[:-response_length, :-response_length] = 0 - mtopdiv = transform_distances_to_mtopdiv(distance_matrice) - mtopdivs.append(mtopdiv) - return np.array(mtopdivs, dtype=float) - - train_mtopdivs.append( - np.stack( - [job(head) for head in range(heads * layers)], - axis=1, - ) + _, num_layers, num_heads, _, _ = forwardpass_attention_weights.shape + heads = product(range(num_layers), range(num_heads)) + mtopdivs = get_mtopdivs( + heads, + length_responses, + forwardpass_attention_weights, + self.n_jobs ) + train_mtopdivs.append(mtopdivs) train_labels += labels train_mtopdivs = np.concatenate(train_mtopdivs, axis=0) train_labels = np.array(train_labels) - return { "train_mtopdivs": train_mtopdivs, "train_labels": train_labels, + "num_layers": num_layers, + "num_heads": num_heads, } + + def __call__( + self, + dependencies: Dict[str, np.array], + texts: List[str], + model: WhiteboxModel, + max_new_tokens: int = 100, + ) -> Dict[str, np.ndarray]: + heads = load_model_heads(self.cache_path, model.model_path) + if heads and self.priority == "cache": + heads = np.array(heads) + return {"topological_divergence_heads": heads[:self.max_heads]} + + self.train_dataset.from_csv() + self.train_dataset.subsample() + train_dataloader = DataLoader( + self.train_dataset, + batch_size=self.train_dataset.batch_size, + shuffle=True + ) + output = self.train(model, train_dataloader) + heads = self.select_heads( + output["train_mtopdivs"], + output["train_labels"], + ) + heads = np.unravel_index(heads, (output["num_layers"], output["num_heads"])) + heads = np.stack(heads, axis=1) + + save_model_heads(self.cache_path, model.model_path, heads.tolist()) + return {"topological_divergence_heads": heads} diff --git a/src/lm_polygraph/utils/factory_estimator.py b/src/lm_polygraph/utils/factory_estimator.py index 70ac0cec7..1bac41782 100644 --- a/src/lm_polygraph/utils/factory_estimator.py +++ b/src/lm_polygraph/utils/factory_estimator.py @@ -52,6 +52,7 @@ def load_simple_estimators(name: str, config): FocusClaim, AttentionScore, AttentionScoreClaim, + TopologicalDivergence, ] try: From 470ac070db886f74067a376452543bd73636b218 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Tue, 24 Jun 2025 17:52:17 +0300 Subject: [PATCH 09/23] small fixes --- .../default_TrainMTopDivCalculator.yaml | 19 ------------------ .../stat_calculators/default_calculators.yaml | 20 ++++++++++++++++++- .../configs/topological_divergence_coqa.yaml | 17 +++++----------- .../default_TrainMTopDivCalculator.py | 2 +- 4 files changed, 25 insertions(+), 33 deletions(-) delete mode 100644 examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml diff --git a/examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml b/examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml deleted file mode 100644 index cc374f163..000000000 --- a/examples/configs/stat_calculators/default_TrainMTopDivCalculator.yaml +++ /dev/null @@ -1,19 +0,0 @@ -- name: TrainMTopDivCalculator - builder: lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator - cfg: - heads_extraction_priority: cache - cache_path: '${model_heads_cache}' - max_heads: '${max_heads}' - n_jobs: '${n_jobs}' - train_data_path: '${train_dataset.train_data_path}' - context_column: '${train_dataset.context}' - question_column: '${train_dataset.question}' - prompt_column: '${train_dataset.prompt}' - response_column: '${train_dataset.generated_answer}' - label_column: '${train_dataset.hallucination}' - batch_size: '${batch_size}' - subsample_train_dataset: '${train_dataset.subsample}' - seed: '${seed}' - stats: - - "topological_divergence_heads" - dependencies: \ No newline at end of file diff --git a/examples/configs/stat_calculators/default_calculators.yaml b/examples/configs/stat_calculators/default_calculators.yaml index 3ce7d2221..aec93f200 100644 --- a/examples/configs/stat_calculators/default_calculators.yaml +++ b/examples/configs/stat_calculators/default_calculators.yaml @@ -28,4 +28,22 @@ - "background_train_embeddings" - "train_greedy_log_likelihoods" dependencies: - +- name: TrainMTopDivCalculator + builder: lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator + cfg: + heads_extraction_priority: cache + cache_path: '${model_heads_cache}' + max_heads: '${max_heads}' + n_jobs: '${n_jobs}' + train_data_path: '${train_dataset.train_data_path}' + context_column: '${train_dataset.context_column}' + question_column: '${train_dataset.question_column}' + prompt_column: '${train_dataset.prompt_column}' + response_column: '${train_dataset.response_column}' + label_column: '${train_dataset.label_column}' + batch_size: '${batch_size}' + subsample_train_dataset: '${train_dataset.subsample}' + seed: '${seed}' + stats: + - "topological_divergence_heads" + dependencies: \ No newline at end of file diff --git a/examples/configs/topological_divergence_coqa.yaml b/examples/configs/topological_divergence_coqa.yaml index 86a53afb5..e19626725 100644 --- a/examples/configs/topological_divergence_coqa.yaml +++ b/examples/configs/topological_divergence_coqa.yaml @@ -3,18 +3,11 @@ hydra: dir: ${cache_path}/${task}/${model.path}/${dataset}/${now:%Y-%m-%d}/${now:%H-%M-%S} defaults: + - model: bloomz-560m - estimators: default_TopologicalDivergence - - - stat_calculators: default_TrainMTopDivCalculator + - stat_calculators: default_calculators - base_processing_coqa - _self_ -model: - path: Qwen/Qwen2.5-0.5B-Instruct - type: CausalLM - path_to_load_script: model/default_causal.py - load_model_args: - device_map: auto - load_tokenizer_args: {} cache_path: ./workdir/output save_path: '${hydra:run.dir}' @@ -39,18 +32,18 @@ generation_params: generate_until: - "\n" subsample_eval_dataset: -1 -batch_size: 1 +batch_size: 4 generation_metrics: null ignore_exceptions: false seed: - 1 -model_heads_cache: "../../../model_heads_cache.yaml" +model_heads_cache: "model_heads_cache.yaml" max_heads: 6 n_jobs: -1 train_dataset: - train_data_path: "../../../coqa_Meta-Llama-3-8B-Instruct.csv" + train_data_path: "coqa_Llama-3.1-8B-Instruct.csv" context_column: context question_column: question prompt_column: prompt diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py b/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py index bf4aef86f..39290e3e1 100644 --- a/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py +++ b/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py @@ -58,7 +58,7 @@ def select(self, indices: List[int]): def subsample(self): if self.subsample_train_dataset >= len(self): return self - rng = np.random.default_rng(self.seed) + rng = np.random.default_rng(self.seed[0]) selected_indices = rng.choice( len(self), size=self.subsample_train_dataset, From e273a8e2baa8bedf0dce63644af25f3b75e31a12 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Thu, 26 Jun 2025 10:43:32 +0300 Subject: [PATCH 10/23] optimize parallel job --- .../configs/topological_divergence_coqa.yaml | 2 +- src/lm_polygraph/estimators/mtopdiv.py | 59 +++++++++++-------- .../estimators/topological_divergence.py | 1 - .../stat_calculators/train_mtopdiv.py | 23 ++++---- 4 files changed, 50 insertions(+), 35 deletions(-) diff --git a/examples/configs/topological_divergence_coqa.yaml b/examples/configs/topological_divergence_coqa.yaml index e19626725..f3d59d747 100644 --- a/examples/configs/topological_divergence_coqa.yaml +++ b/examples/configs/topological_divergence_coqa.yaml @@ -32,7 +32,7 @@ generation_params: generate_until: - "\n" subsample_eval_dataset: -1 -batch_size: 4 +batch_size: 16 generation_metrics: null ignore_exceptions: false diff --git a/src/lm_polygraph/estimators/mtopdiv.py b/src/lm_polygraph/estimators/mtopdiv.py index 6258c51a4..928fad218 100644 --- a/src/lm_polygraph/estimators/mtopdiv.py +++ b/src/lm_polygraph/estimators/mtopdiv.py @@ -29,30 +29,37 @@ def transform_attention_scores_to_distances( - attention_weights: torch.Tensor, -) -> torch.Tensor: + attention_weights: np.ndarray, +) -> np.ndarray: """Transform attention matrix to the matrix of distances between tokens. Parameters ---------- - attention_weights : torch.Tensor - Attention matrixes of one sample (n_heads x n_tokens x n_tokens). + attention_weights : np.ndarray + Attention matrices of one sample (n_heads x n_tokens x n_tokens). Returns ------- - np.array + np.ndarray Distance matrix. """ - attention_weights = attention_weights.float() + attention_weights = attention_weights.astype(np.float32) n_tokens = attention_weights.shape[-1] - distance_mx = 1 - torch.clamp(attention_weights, min=0.0) - zero_diag = ( - torch.ones(n_tokens, n_tokens, device=attention_weights.device) - - torch.eye(n_tokens, device=attention_weights.device) + + # Convert attention weights to distances + distance_mx = 1 - np.clip(attention_weights, a_min=0.0, a_max=None) + + # Create mask to zero out diagonal + zero_diag = np.ones((n_tokens, n_tokens)) - np.eye(n_tokens) + + # Apply mask and ensure symmetry + distance_mx *= np.broadcast_to(zero_diag, distance_mx.shape) + distance_mx = np.minimum( + np.swapaxes(distance_mx, -1, -2), + distance_mx ) - distance_mx *= zero_diag.expand_as(distance_mx) - distance_mx = torch.minimum(distance_mx.transpose(-1, -2), distance_mx) + return distance_mx @@ -78,35 +85,41 @@ def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: def get_mtopdivs( heads: List[Tuple[int, int]], length_responses: List[int], - attention_weights_batch: torch.Tensor, + attention_weights_batch: np.array, n_jobs: Optional[float] = -1, ) -> np.ndarray: batch_size = attention_weights_batch.shape[0] padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) - distance_matrices_batch = transform_attention_scores_to_distances( - attention_weights_batch - ).numpy() - - def job(layer, head): + + def job(layer_head_pair): + layer, head = layer_head_pair mtopdivs = [] + distance_matrices = transform_attention_scores_to_distances( + attention_weights_batch[:, layer, head] + ) + for sample_id in range(batch_size): - distance_matrix = distance_matrices_batch[sample_id, layer, head] + distance_matrix = distance_matrices[sample_id] padding_length = padding_lengths[sample_id] response_length = length_responses[sample_id] + if padding_length > 0: distance_matrix = distance_matrix[:-padding_length, :-padding_length] distance_matrix[:-response_length, :-response_length] = 0 mtopdiv = transform_distances_to_mtopdiv(distance_matrix) / response_length mtopdivs.append(mtopdiv) + return np.array(mtopdivs, dtype=float) + # Use loky backend with shared memory if IS_PARALLEL_AVAILABLE: - mtopdivs = Parallel(n_jobs=n_jobs, backend="threading")( - delayed(job)(layer, head) for layer, head in heads - ) + with Parallel(n_jobs=n_jobs, prefer='processes') as parallel: + mtopdivs = parallel(delayed(job)((layer, head)) for layer, head in heads) else: - mtopdivs = [job(layer, head) for layer, head in heads] + mtopdivs = [job((layer, head)) for layer, head in heads] + mtopdivs = np.stack(mtopdivs, axis=1) + return mtopdivs def load_model_heads( diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 185692dd0..4a6d7d13e 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -78,7 +78,6 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: # Alternatively, a new estimator instance can be created with the selected heads. self.stats_dependencies = ["greedy_tokens", "forwardpass_attention_weights"] - attention_weights_batch = torch.from_numpy(attention_weights_batch) mtopdivs = get_mtopdivs( self._heads, length_responses, diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py index f5c741f41..edb1d1ac6 100644 --- a/src/lm_polygraph/stat_calculators/train_mtopdiv.py +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -63,14 +63,13 @@ def calculate_attention_forward_pass( output_attentions=True, ).attentions mask = batch["attention_mask"] - mask = mask[:, None, None, None, :] & mask[:, None, None, :, None] - attns = torch.stack(attns, dim=1) - attns = attns.masked_fill(mask == 0, float("nan")).cpu() - + mask = mask[:, None, None, :] & mask[:, None, :, None] + attns = [attn.masked_fill(mask == 0, float("nan")).cpu() for attn in attns] + torch.cuda.empty_cache() gc.collect() - return attns + return torch.stack(attns, dim=1).numpy() def select_heads(self, scores, labels): grounded_scores, hal_scores = scores[labels == 0], scores[labels == 1] @@ -92,11 +91,6 @@ def train(self, model, train_dataloader): train_labels = [] for inp_texts, target_texts, labels in tqdm(train_dataloader): - forwardpass_attention_weights = self.calculate_attention_forward_pass( - inp_texts, - target_texts, - model, - ) length_responses = model.tokenizer( target_texts, add_special_tokens=False, @@ -104,6 +98,11 @@ def train(self, model, train_dataloader): return_tensors="pt", ) length_responses = length_responses["attention_mask"].sum(dim=1).tolist() + forwardpass_attention_weights = self.calculate_attention_forward_pass( + inp_texts, + target_texts, + model, + ) _, num_layers, num_heads, _, _ = forwardpass_attention_weights.shape heads = product(range(num_layers), range(num_heads)) mtopdivs = get_mtopdivs( @@ -112,9 +111,13 @@ def train(self, model, train_dataloader): forwardpass_attention_weights, self.n_jobs ) + train_mtopdivs.append(mtopdivs) train_labels += labels + del forwardpass_attention_weights + gc.collect() + train_mtopdivs = np.concatenate(train_mtopdivs, axis=0) train_labels = np.array(train_labels) return { From b5501114aaf33e1c83d1169e969850d0d0fcbeec Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Mon, 30 Jun 2025 17:18:35 +0300 Subject: [PATCH 11/23] use standard lm-polygraph's data handling for train --- .../configs/topological_divergence_coqa.yaml | 13 +- ...ical_divergence_with_heads_selection.ipynb | 63 +++++--- .../default_TrainMTopDivCalculator.py | 112 +++++-------- .../{mtopdiv.py => mtopdiv_utils.py} | 2 + .../estimators/topological_divergence.py | 2 +- .../stat_calculators/train_mtopdiv.py | 147 +++++++----------- 6 files changed, 142 insertions(+), 197 deletions(-) rename src/lm_polygraph/estimators/{mtopdiv.py => mtopdiv_utils.py} (98%) diff --git a/examples/configs/topological_divergence_coqa.yaml b/examples/configs/topological_divergence_coqa.yaml index f3d59d747..5fb80db4c 100644 --- a/examples/configs/topological_divergence_coqa.yaml +++ b/examples/configs/topological_divergence_coqa.yaml @@ -32,21 +32,16 @@ generation_params: generate_until: - "\n" subsample_eval_dataset: -1 -batch_size: 16 +batch_size: 1 generation_metrics: null ignore_exceptions: false seed: - 1 +subsample_train_dataset: 100 +heads_extraction_priority: "cache" model_heads_cache: "model_heads_cache.yaml" max_heads: 6 n_jobs: -1 -train_dataset: - train_data_path: "coqa_Llama-3.1-8B-Instruct.csv" - context_column: context - question_column: question - prompt_column: prompt - response_column: generated_answer - label_column: hallucination - subsample: 100 + diff --git a/examples/topological_divergence_with_heads_selection.ipynb b/examples/topological_divergence_with_heads_selection.ipynb index 21f3f65c5..f0d558328 100644 --- a/examples/topological_divergence_with_heads_selection.ipynb +++ b/examples/topological_divergence_with_heads_selection.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "bd1b5660", "metadata": {}, "outputs": [ @@ -69,28 +69,49 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "430d4c18", "metadata": {}, "outputs": [], "source": [ "class Config:\n", " model_name = 'Qwen/Qwen2.5-0.5B-Instruct'\n", - " mtopdiv_train_data_path = '../../coqa_Meta-Llama-3-8B-Instruct.csv'\n", - " mtopdiv_heads_extraction_priority = \"config\"\n", - " mtopdiv_heads_path = \"../../config.yaml\"\n", - " context_column = 'context'\n", - " question_column = 'question'\n", - " prompt_column = 'prompt'\n", - " response_column = 'generated_answer'\n", - " label_column = 'hallucination'\n", - " batch_size = 1\n", + " \n", + " heads_extraction_priority = \"cache\"\n", + " model_heads_cache = \"model_heads_cache.yaml\"\n", + " \n", " subsample_train_dataset = 4\n", - " seed = 52\n", + " seed = [52]\n", " device_map = 'mps' \n", - " mtopdiv_max_heads = 6\n", + " max_heads = 6\n", " n_jobs = -1\n", "\n", + " cache_path = \"./workdir/output\"\n", + "\n", + " task = \"qa\"\n", + " instruct = False\n", + "\n", + " dataset = ['LM-polygraph/coqa', 'continuation']\n", + " text_column = \"input\"\n", + " label_column = \"output\"\n", + " train_split = \"train\"\n", + " eval_split = \"test\"\n", + " n_shot = 0\n", + " few_shot_split = \"train\"\n", + " few_shot_prompt = None\n", + " trust_remote_code = False\n", + " size = None\n", + "\n", + " max_new_tokens = 20\n", + " load_from_disk = False\n", + " generation_params = {\"generate_until\": [\"\\n\"]}\n", + " subsample_eval_dataset = -1\n", + " batch_size = 1\n", + "\n", + " generation_metrics = None\n", + " ignore_exceptions = False\n", + "\n", + "\n", "\n", "cfg = Config()\n", "\n", @@ -124,16 +145,18 @@ "name": "stderr", "output_type": "stream", "text": [ - "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to eager attention. This warning can be removed using the argument `attn_implementation=\"eager\"` when loading the model.\n" + " 0%| | 0/4 [00:00= len(self): - return self - rng = np.random.default_rng(self.seed[0]) - selected_indices = rng.choice( - len(self), - size=self.subsample_train_dataset, - replace=False, - ).tolist() - return self.select(selected_indices) def load_stat_calculator(config, builder): priority = config.heads_extraction_priority - cache_path = config.cache_path + cache_path = os.path.join(config.cache_path, config.model_heads_cache) max_heads = config.max_heads n_jobs = config.n_jobs - - train_dataset = MTopDivHoldoutDataset( - csv_path=config.train_data_path, - context_column=config.context_column, - question_column=config.question_column, - prompt_column=config.prompt_column, - response_column=config.response_column, - label_column=config.label_column, - batch_size=config.batch_size, - subsample_train_dataset=config.subsample_train_dataset, - seed=config.seed, - ) + load_train_dataset_fn = lambda: load_dataset(config) return TrainMTopDivCalculator( priority, - train_dataset, + load_train_dataset_fn, cache_path, max_heads, n_jobs, diff --git a/src/lm_polygraph/estimators/mtopdiv.py b/src/lm_polygraph/estimators/mtopdiv_utils.py similarity index 98% rename from src/lm_polygraph/estimators/mtopdiv.py rename to src/lm_polygraph/estimators/mtopdiv_utils.py index 928fad218..971c0cb64 100644 --- a/src/lm_polygraph/estimators/mtopdiv.py +++ b/src/lm_polygraph/estimators/mtopdiv_utils.py @@ -186,6 +186,8 @@ def save_model_heads( models.append({model_path: f"{heads}"}) config["models"] = models + os.makedirs(os.path.dirname(cache_path), exist_ok=True) + with open(cache_path, "w") as f: yaml.safe_dump(config, f, sort_keys=False) print(f"Saved heads for '{model_path}' to {cache_path}") diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 4a6d7d13e..1a03fec0e 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -3,7 +3,7 @@ import torch from .estimator import Estimator -from .mtopdiv import get_mtopdivs +from .mtopdiv_utils import get_mtopdivs class TopologicalDivergence(Estimator): diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py index edb1d1ac6..90308f806 100644 --- a/src/lm_polygraph/stat_calculators/train_mtopdiv.py +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -1,16 +1,13 @@ -import gc -import torch -from torch.utils.data import DataLoader import numpy as np -import pandas as pd from sklearn.metrics import roc_auc_score from tqdm import tqdm -from typing import Dict, List, Tuple, Literal +from typing import Dict, List, Tuple, Literal, Callable from itertools import product -from .stat_calculator import StatCalculator +from ..stat_calculators import StatCalculator, GreedyProbsCalculator, AttentionForwardPassCalculator +from ..generation_metrics import RougeMetric from lm_polygraph.utils.model import WhiteboxModel -from lm_polygraph.estimators.mtopdiv import ( +from lm_polygraph.estimators.mtopdiv_utils import ( get_mtopdivs, load_model_heads, save_model_heads, @@ -28,48 +25,17 @@ def meta_info() -> Tuple[List[str], List[str]]: def __init__(self, priority: Literal["train", "cache"] = "cache", - train_dataset: "MTopDivHoldoutDataset" = None, + load_train_dataset_fn: Callable = None, cache_path: str = None, max_heads: int = 6, n_jobs: int = -1): super().__init__() - if train_dataset is None and cache_path is None: - raise Exception("Either train_dataset or cache_path must be provided.") + self.priority = priority self.cache_path = cache_path - self.train_dataset = train_dataset + self.load_train_dataset_fn = load_train_dataset_fn self.max_heads = max_heads self.n_jobs = n_jobs - - @staticmethod - def calculate_attention_forward_pass( - prompts: List[str], - responses: List[str], - model: WhiteboxModel, - ) -> Tuple[List[np.ndarray], np.ndarray]: - - if model.tokenizer.chat_template is not None: - chats = [ - [{"role": "user", "content": p}, {"role": "assistant", "content": r}] - for p, r in zip(prompts, responses) - ] - else: - chats = [p + r for p, r in zip(prompts, responses)] - - with torch.no_grad(): - batch = model.tokenize(chats).to(model.device()) - attns = model.model( - **batch, - output_attentions=True, - ).attentions - mask = batch["attention_mask"] - mask = mask[:, None, None, :] & mask[:, None, :, None] - attns = [attn.masked_fill(mask == 0, float("nan")).cpu() for attn in attns] - - torch.cuda.empty_cache() - gc.collect() - - return torch.stack(attns, dim=1).numpy() def select_heads(self, scores, labels): grounded_scores, hal_scores = scores[labels == 0], scores[labels == 1] @@ -86,47 +52,6 @@ def select_heads(self, scores, labels): n_opt = n return heads[:n_opt] - def train(self, model, train_dataloader): - train_mtopdivs = [] - train_labels = [] - - for inp_texts, target_texts, labels in tqdm(train_dataloader): - length_responses = model.tokenizer( - target_texts, - add_special_tokens=False, - padding=True, - return_tensors="pt", - ) - length_responses = length_responses["attention_mask"].sum(dim=1).tolist() - forwardpass_attention_weights = self.calculate_attention_forward_pass( - inp_texts, - target_texts, - model, - ) - _, num_layers, num_heads, _, _ = forwardpass_attention_weights.shape - heads = product(range(num_layers), range(num_heads)) - mtopdivs = get_mtopdivs( - heads, - length_responses, - forwardpass_attention_weights, - self.n_jobs - ) - - train_mtopdivs.append(mtopdivs) - train_labels += labels - - del forwardpass_attention_weights - gc.collect() - - train_mtopdivs = np.concatenate(train_mtopdivs, axis=0) - train_labels = np.array(train_labels) - return { - "train_mtopdivs": train_mtopdivs, - "train_labels": train_labels, - "num_layers": num_layers, - "num_heads": num_heads, - } - def __call__( self, dependencies: Dict[str, np.array], @@ -139,19 +64,51 @@ def __call__( heads = np.array(heads) return {"topological_divergence_heads": heads[:self.max_heads]} - self.train_dataset.from_csv() - self.train_dataset.subsample() - train_dataloader = DataLoader( - self.train_dataset, - batch_size=self.train_dataset.batch_size, - shuffle=True - ) - output = self.train(model, train_dataloader) - heads = self.select_heads( - output["train_mtopdivs"], - output["train_labels"], - ) - heads = np.unravel_index(heads, (output["num_layers"], output["num_heads"])) + train_dataset = self.load_train_dataset_fn() + + mtopdivs = [] + labels = [] + + greedy_calc = GreedyProbsCalculator(False, False) + attn_forward_pass_calc = AttentionForwardPassCalculator() + generation_metric = RougeMetric("rougeL") + + for input_text, target in tqdm(train_dataset): + stats = greedy_calc( + dependencies=None, + texts=input_text, + model=model, + max_new_tokens=max_new_tokens, + ) + attn_weights_batch = attn_forward_pass_calc( + dependencies=stats, + texts=input_text, + model=model, + max_new_tokens=max_new_tokens, + )["forwardpass_attention_weights"] + length_responses = list(map(len, stats["greedy_tokens"])) + + _, num_layers, num_heads, _, _ = attn_weights_batch.shape + heads = product(range(num_layers), range(num_heads)) + + mtopdivs.append(get_mtopdivs( + heads, + length_responses, + attn_weights_batch, + n_jobs=self.n_jobs, + )) + labels.append(generation_metric( + stats=stats, + target_texts=target, + )) + + mtopdivs = np.concatenate(mtopdivs, axis=0) + + labels = np.concatenate(labels) + labels = ~(labels > 0.3) + + heads = self.select_heads(mtopdivs, labels) + heads = np.unravel_index(heads, (num_layers, num_heads)) heads = np.stack(heads, axis=1) save_model_heads(self.cache_path, model.model_path, heads.tolist()) From 9ebddda938ccd9411cba6443f78f9818b3c3fb11 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Thu, 3 Jul 2025 11:23:48 +0300 Subject: [PATCH 12/23] adapt configs for a newer script --- .../stat_calculators/default_calculators.yaml | 21 ++++++++++++------- .../configs/topological_divergence_coqa.yaml | 1 - 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/examples/configs/stat_calculators/default_calculators.yaml b/examples/configs/stat_calculators/default_calculators.yaml index aec93f200..84c0eecdf 100644 --- a/examples/configs/stat_calculators/default_calculators.yaml +++ b/examples/configs/stat_calculators/default_calculators.yaml @@ -32,18 +32,23 @@ builder: lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator cfg: heads_extraction_priority: cache - cache_path: '${model_heads_cache}' + cache_path: '${cache_path}' + model_heads_cache: '${model_heads_cache}' max_heads: '${max_heads}' n_jobs: '${n_jobs}' - train_data_path: '${train_dataset.train_data_path}' - context_column: '${train_dataset.context_column}' - question_column: '${train_dataset.question_column}' - prompt_column: '${train_dataset.prompt_column}' - response_column: '${train_dataset.response_column}' - label_column: '${train_dataset.label_column}' + dataset: '${dataset}' + text_column: '${text_column}' + label_column: '${label_column}' + description: '' + prompt: '' + few_shot_split: "train" + train_split: '${train_split}' + load_from_disk: '${load_from_disk}' + subsample_train_dataset: '${subsample_train_dataset}' + n_shot: '${n_shot}' batch_size: '${batch_size}' - subsample_train_dataset: '${train_dataset.subsample}' seed: '${seed}' + size: '${size}' stats: - "topological_divergence_heads" dependencies: \ No newline at end of file diff --git a/examples/configs/topological_divergence_coqa.yaml b/examples/configs/topological_divergence_coqa.yaml index 5fb80db4c..997168a72 100644 --- a/examples/configs/topological_divergence_coqa.yaml +++ b/examples/configs/topological_divergence_coqa.yaml @@ -44,4 +44,3 @@ heads_extraction_priority: "cache" model_heads_cache: "model_heads_cache.yaml" max_heads: 6 n_jobs: -1 - From ca522a299885d636e3f300623446b213c5725500 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Thu, 3 Jul 2025 11:50:28 +0300 Subject: [PATCH 13/23] flake8-black polishing --- .../default_TrainMTopDivCalculator.py | 6 +- src/lm_polygraph/estimators/mtopdiv_utils.py | 42 ++++++------- .../estimators/topological_divergence.py | 3 - .../stat_calculators/train_mtopdiv.py | 61 +++++++++++-------- 4 files changed, 58 insertions(+), 54 deletions(-) diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py b/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py index 1a4da965e..6acf7bb10 100644 --- a/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py +++ b/src/lm_polygraph/defaults/stat_calculator_builders/default_TrainMTopDivCalculator.py @@ -11,7 +11,7 @@ def load_dataset(args): - + log.info("=" * 100) log.info("Loading train dataset...") @@ -50,7 +50,9 @@ def load_stat_calculator(config, builder): cache_path = os.path.join(config.cache_path, config.model_heads_cache) max_heads = config.max_heads n_jobs = config.n_jobs - load_train_dataset_fn = lambda: load_dataset(config) + + def load_train_dataset_fn(): + return load_dataset(config) return TrainMTopDivCalculator( priority, diff --git a/src/lm_polygraph/estimators/mtopdiv_utils.py b/src/lm_polygraph/estimators/mtopdiv_utils.py index 971c0cb64..3159db1d5 100644 --- a/src/lm_polygraph/estimators/mtopdiv_utils.py +++ b/src/lm_polygraph/estimators/mtopdiv_utils.py @@ -2,8 +2,6 @@ import yaml import warnings import numpy as np -import pandas as pd -import torch from typing import List, Tuple, Optional import ast @@ -46,20 +44,15 @@ def transform_attention_scores_to_distances( """ attention_weights = attention_weights.astype(np.float32) n_tokens = attention_weights.shape[-1] - - # Convert attention weights to distances distance_mx = 1 - np.clip(attention_weights, a_min=0.0, a_max=None) - - # Create mask to zero out diagonal zero_diag = np.ones((n_tokens, n_tokens)) - np.eye(n_tokens) - - # Apply mask and ensure symmetry + distance_mx *= np.broadcast_to(zero_diag, distance_mx.shape) distance_mx = np.minimum( - np.swapaxes(distance_mx, -1, -2), - distance_mx + np.swapaxes(distance_mx, -1, -2), + distance_mx, ) - + return distance_mx @@ -82,6 +75,7 @@ def transform_distances_to_mtopdiv(distance_mx: np.ndarray) -> float: return barcodes[0][:-1, 1].sum() return 0 + def get_mtopdivs( heads: List[Tuple[int, int]], length_responses: List[int], @@ -90,41 +84,40 @@ def get_mtopdivs( ) -> np.ndarray: batch_size = attention_weights_batch.shape[0] padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) - + def job(layer_head_pair): layer, head = layer_head_pair mtopdivs = [] distance_matrices = transform_attention_scores_to_distances( attention_weights_batch[:, layer, head] ) - + for sample_id in range(batch_size): distance_matrix = distance_matrices[sample_id] padding_length = padding_lengths[sample_id] response_length = length_responses[sample_id] - + if padding_length > 0: distance_matrix = distance_matrix[:-padding_length, :-padding_length] distance_matrix[:-response_length, :-response_length] = 0 mtopdiv = transform_distances_to_mtopdiv(distance_matrix) / response_length mtopdivs.append(mtopdiv) - + return np.array(mtopdivs, dtype=float) - # Use loky backend with shared memory if IS_PARALLEL_AVAILABLE: - with Parallel(n_jobs=n_jobs, prefer='processes') as parallel: + with Parallel(n_jobs=n_jobs, prefer="processes") as parallel: mtopdivs = parallel(delayed(job)((layer, head)) for layer, head in heads) else: mtopdivs = [job((layer, head)) for layer, head in heads] mtopdivs = np.stack(mtopdivs, axis=1) - return mtopdivs + def load_model_heads( - cache_path: Optional[str], - model_path: str + cache_path: Optional[str], + model_path: str, ) -> Optional[List[Tuple[int, int]]]: """ Load model heads. @@ -155,10 +148,11 @@ def load_model_heads( print(f"Failed to load heads from cache: {e}") return None + def save_model_heads( - cache_path: str, - model_path: str, - heads: List[Tuple[int, int]] + cache_path: str, + model_path: str, + heads: List[Tuple[int, int]], ) -> None: """ Save a list of heads for a model. @@ -182,7 +176,7 @@ def save_model_heads( for entry in models: if isinstance(entry, dict) and model_path in entry: return - + models.append({model_path: f"{heads}"}) config["models"] = models diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 1a03fec0e..10496d37b 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -1,6 +1,5 @@ from typing import Dict, List, Tuple, Optional import numpy as np -import torch from .estimator import Estimator from .mtopdiv_utils import get_mtopdivs @@ -86,5 +85,3 @@ def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: ) return np.mean(mtopdivs, axis=1) - - diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py index 90308f806..47136781d 100644 --- a/src/lm_polygraph/stat_calculators/train_mtopdiv.py +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -4,7 +4,11 @@ from typing import Dict, List, Tuple, Literal, Callable from itertools import product -from ..stat_calculators import StatCalculator, GreedyProbsCalculator, AttentionForwardPassCalculator +from ..stat_calculators import ( + StatCalculator, + GreedyProbsCalculator, + AttentionForwardPassCalculator, +) from ..generation_metrics import RougeMetric from lm_polygraph.utils.model import WhiteboxModel from lm_polygraph.estimators.mtopdiv_utils import ( @@ -23,12 +27,14 @@ def meta_info() -> Tuple[List[str], List[str]]: return ["topological_divergence_heads"], [] - def __init__(self, - priority: Literal["train", "cache"] = "cache", - load_train_dataset_fn: Callable = None, - cache_path: str = None, - max_heads: int = 6, - n_jobs: int = -1): + def __init__( + self, + priority: Literal["train", "cache"] = "cache", + load_train_dataset_fn: Callable = None, + cache_path: str = None, + max_heads: int = 6, + n_jobs: int = -1, + ): super().__init__() self.priority = priority @@ -36,7 +42,7 @@ def __init__(self, self.load_train_dataset_fn = load_train_dataset_fn self.max_heads = max_heads self.n_jobs = n_jobs - + def select_heads(self, scores, labels): grounded_scores, hal_scores = scores[labels == 0], scores[labels == 1] deltas = hal_scores.mean(0) - grounded_scores.mean(0) @@ -51,7 +57,7 @@ def select_heads(self, scores, labels): best_auroc = roc_auc n_opt = n return heads[:n_opt] - + def __call__( self, dependencies: Dict[str, np.array], @@ -62,13 +68,14 @@ def __call__( heads = load_model_heads(self.cache_path, model.model_path) if heads and self.priority == "cache": heads = np.array(heads) - return {"topological_divergence_heads": heads[:self.max_heads]} - + heads = heads[: self.max_heads] + return {"topological_divergence_heads": heads} + train_dataset = self.load_train_dataset_fn() - + mtopdivs = [] labels = [] - + greedy_calc = GreedyProbsCalculator(False, False) attn_forward_pass_calc = AttentionForwardPassCalculator() generation_metric = RougeMetric("rougeL") @@ -87,23 +94,27 @@ def __call__( max_new_tokens=max_new_tokens, )["forwardpass_attention_weights"] length_responses = list(map(len, stats["greedy_tokens"])) - + _, num_layers, num_heads, _, _ = attn_weights_batch.shape heads = product(range(num_layers), range(num_heads)) - mtopdivs.append(get_mtopdivs( - heads, - length_responses, - attn_weights_batch, - n_jobs=self.n_jobs, - )) - labels.append(generation_metric( - stats=stats, - target_texts=target, - )) + mtopdivs.append( + get_mtopdivs( + heads, + length_responses, + attn_weights_batch, + n_jobs=self.n_jobs, + ) + ) + labels.append( + generation_metric( + stats=stats, + target_texts=target, + ) + ) mtopdivs = np.concatenate(mtopdivs, axis=0) - + labels = np.concatenate(labels) labels = ~(labels > 0.3) From 2299722ce3ad50bae773c440c1aaab1e6978b88e Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Thu, 3 Jul 2025 12:20:12 +0300 Subject: [PATCH 14/23] add test --- test/test_estimators.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/test/test_estimators.py b/test/test_estimators.py index 456f5aaca..77233fda7 100644 --- a/test/test_estimators.py +++ b/test/test_estimators.py @@ -250,3 +250,9 @@ def test_attentionscore(model): estimator = AttentionScore() ue = estimate_uncertainty(model, estimator, INPUT) assert isinstance(ue.uncertainty, float) + + +def test_topological_divergence(model): + estimator = TopologicalDivergence(heads=[[0, 0], [0, 1], [1, 0], [1, 1]]) + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) From 3a04f7041ea6c7be26c77fcd332d63ea37929e44 Mon Sep 17 00:00:00 2001 From: ahdr3w <71148619+ahdr3w@users.noreply.github.com> Date: Wed, 9 Jul 2025 21:47:50 +0300 Subject: [PATCH 15/23] Update requirements.txt Add ripser and joblib in requirements.txt --- requirements.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/requirements.txt b/requirements.txt index cf37ac3c9..273b4873d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -35,3 +35,5 @@ evaluate>=0.4.2 spacy>=3.4.0,<3.8.0 fastchat diskcache>=5.6.3 +joblib +ripser From 9a5bc9087cd6c23c0b73685c36a4b9a196bfc777 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Sat, 19 Jul 2025 02:23:18 +0300 Subject: [PATCH 16/23] move topological_divergence estimator into default_estimators --- .../default_TopologicalDivergence.yaml | 5 -- .../estimators/default_estimators.yaml | 4 ++ .../stat_calculators/default_calculators.yaml | 8 ++-- .../configs/topological_divergence_coqa.yaml | 46 ------------------- .../estimators/topological_divergence.py | 6 +-- 5 files changed, 10 insertions(+), 59 deletions(-) delete mode 100644 examples/configs/estimators/default_TopologicalDivergence.yaml delete mode 100644 examples/configs/topological_divergence_coqa.yaml diff --git a/examples/configs/estimators/default_TopologicalDivergence.yaml b/examples/configs/estimators/default_TopologicalDivergence.yaml deleted file mode 100644 index 6b87c9f29..000000000 --- a/examples/configs/estimators/default_TopologicalDivergence.yaml +++ /dev/null @@ -1,5 +0,0 @@ -- name: TopologicalDivergence - cfg: - heads: null - max_heads: '${max_heads}' - n_jobs: '${n_jobs}' \ No newline at end of file diff --git a/examples/configs/estimators/default_estimators.yaml b/examples/configs/estimators/default_estimators.yaml index 4b8de1076..fb7cbcb81 100644 --- a/examples/configs/estimators/default_estimators.yaml +++ b/examples/configs/estimators/default_estimators.yaml @@ -88,3 +88,7 @@ cfg: layer: 16 gen_only: False +- name: TopologicalDivergence + cfg: + heads: null + n_jobs: -1 \ No newline at end of file diff --git a/examples/configs/stat_calculators/default_calculators.yaml b/examples/configs/stat_calculators/default_calculators.yaml index 84c0eecdf..82ab0f6bd 100644 --- a/examples/configs/stat_calculators/default_calculators.yaml +++ b/examples/configs/stat_calculators/default_calculators.yaml @@ -33,9 +33,9 @@ cfg: heads_extraction_priority: cache cache_path: '${cache_path}' - model_heads_cache: '${model_heads_cache}' - max_heads: '${max_heads}' - n_jobs: '${n_jobs}' + model_heads_cache: "model_heads_cache.yaml" + max_heads: 6 + n_jobs: -1 dataset: '${dataset}' text_column: '${text_column}' label_column: '${label_column}' @@ -44,7 +44,7 @@ few_shot_split: "train" train_split: '${train_split}' load_from_disk: '${load_from_disk}' - subsample_train_dataset: '${subsample_train_dataset}' + subsample_train_dataset: 100 n_shot: '${n_shot}' batch_size: '${batch_size}' seed: '${seed}' diff --git a/examples/configs/topological_divergence_coqa.yaml b/examples/configs/topological_divergence_coqa.yaml deleted file mode 100644 index 997168a72..000000000 --- a/examples/configs/topological_divergence_coqa.yaml +++ /dev/null @@ -1,46 +0,0 @@ -hydra: - run: - dir: ${cache_path}/${task}/${model.path}/${dataset}/${now:%Y-%m-%d}/${now:%H-%M-%S} - -defaults: - - model: bloomz-560m - - estimators: default_TopologicalDivergence - - stat_calculators: default_calculators - - base_processing_coqa - - _self_ - -cache_path: ./workdir/output -save_path: '${hydra:run.dir}' - -task: qa -instruct: false - -dataset: ['LM-polygraph/coqa', 'continuation'] -text_column: input -label_column: output -train_split: train -eval_split: test -n_shot: 0 -few_shot_split: train -few_shot_prompt: null -trust_remote_code: false -size: null - -max_new_tokens: 20 -load_from_disk: false -generation_params: - generate_until: - - "\n" -subsample_eval_dataset: -1 -batch_size: 1 - -generation_metrics: null -ignore_exceptions: false -seed: - - 1 - -subsample_train_dataset: 100 -heads_extraction_priority: "cache" -model_heads_cache: "model_heads_cache.yaml" -max_heads: 6 -n_jobs: -1 diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 10496d37b..20805fb97 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -18,17 +18,16 @@ class TopologicalDivergence(Estimator): def __init__( self, heads: Optional[List[Tuple[int, int]]] = None, - max_heads: Optional[int] = 6, n_jobs: int = -1, ): """ Initializes TopologicalDivergence estimator. Parameters: - selected_heads (List[Tuple[int, int]]): List of attention heads to calculate MTopDiv for. + heads (List[Tuple[int, int]]): List of attention heads to calculate MTopDiv for. First integer is layer index, second is head index. If not provided or empty, all heads will be used. - n_jobs (int): Number of jobs for parallel processing. Default: 16. + n_jobs (int): Number of jobs for parallel processing. Default: -1. """ if not heads: calculators = [ @@ -41,7 +40,6 @@ def __init__( super().__init__(calculators, "sequence") self._heads = heads - self._max_heads = max_heads self._n_jobs = n_jobs def __str__(self): From 69d3ee9ebff8422415406941e6afda4c708e67ca Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Sat, 19 Jul 2025 03:39:53 +0300 Subject: [PATCH 17/23] remove explicit declaration of TrainMTopDivCalculator from configs and register it in default calculators --- .../stat_calculators/default_calculators.yaml | 24 ------------------- .../register_default_stat_calculators.py | 24 +++++++++++++++++++ .../stat_calculators/train_mtopdiv.py | 5 ++-- 3 files changed, 27 insertions(+), 26 deletions(-) diff --git a/examples/configs/stat_calculators/default_calculators.yaml b/examples/configs/stat_calculators/default_calculators.yaml index 82ab0f6bd..b9617b723 100644 --- a/examples/configs/stat_calculators/default_calculators.yaml +++ b/examples/configs/stat_calculators/default_calculators.yaml @@ -28,27 +28,3 @@ - "background_train_embeddings" - "train_greedy_log_likelihoods" dependencies: -- name: TrainMTopDivCalculator - builder: lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator - cfg: - heads_extraction_priority: cache - cache_path: '${cache_path}' - model_heads_cache: "model_heads_cache.yaml" - max_heads: 6 - n_jobs: -1 - dataset: '${dataset}' - text_column: '${text_column}' - label_column: '${label_column}' - description: '' - prompt: '' - few_shot_split: "train" - train_split: '${train_split}' - load_from_disk: '${load_from_disk}' - subsample_train_dataset: 100 - n_shot: '${n_shot}' - batch_size: '${batch_size}' - seed: '${seed}' - size: '${size}' - stats: - - "topological_divergence_heads" - dependencies: \ No newline at end of file diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index b7fa94fc7..815d9bbdd 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -148,6 +148,30 @@ def _register( }, ) _register(AttentionForwardPassCalculator) + _register( + TrainMTopDivCalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator", + { + "heads_extraction_priority": "cache", + "cache_path": './workdir/output', + "model_heads_cache": "model_heads_cache.yaml", + "max_heads": 6, + "n_jobs": -1, + "dataset": ["LM-polygraph/coqa", "continuation"], + "text_column": "input", + "label_column": "output", + "description": '', + "prompt": '', + "few_shot_split": "train", + "train_split": "train", + "load_from_disk": False, + "subsample_train_dataset": 100, + "n_shot": 0, + "batch_size": 1, + "seed": [1], + "size": None, + }, + ) elif model_type == "VisualLM": _register( GreedyProbsVisualCalculator, diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py index 47136781d..6915853f4 100644 --- a/src/lm_polygraph/stat_calculators/train_mtopdiv.py +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -65,7 +65,8 @@ def __call__( model: WhiteboxModel, max_new_tokens: int = 100, ) -> Dict[str, np.ndarray]: - heads = load_model_heads(self.cache_path, model.model_path) + name_or_path = model.model.config.name_or_path + heads = load_model_heads(self.cache_path, name_or_path) if heads and self.priority == "cache": heads = np.array(heads) heads = heads[: self.max_heads] @@ -122,5 +123,5 @@ def __call__( heads = np.unravel_index(heads, (num_layers, num_heads)) heads = np.stack(heads, axis=1) - save_model_heads(self.cache_path, model.model_path, heads.tolist()) + save_model_heads(self.cache_path, name_or_path, heads.tolist()) return {"topological_divergence_heads": heads} From 0d3e053d780114d1a52bf10a6c51e5e743536414 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Sat, 19 Jul 2025 03:40:18 +0300 Subject: [PATCH 18/23] update test --- test/test_estimators.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/test/test_estimators.py b/test/test_estimators.py index f58f16865..86915212c 100644 --- a/test/test_estimators.py +++ b/test/test_estimators.py @@ -252,12 +252,22 @@ def test_attentionscore(model): assert isinstance(ue.uncertainty, float) -def test_topological_divergence(model): +def test_topological_divergence_fixed_heads(model): estimator = TopologicalDivergence(heads=[[0, 0], [0, 1], [1, 0], [1, 1]]) ue = estimate_uncertainty(model, estimator, INPUT) assert isinstance(ue.uncertainty, float) +def test_topological_divergence_select_heads(model): + estimator = TopologicalDivergence() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + # second time heads should be loaded from cache + estimator = TopologicalDivergence() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + def test_cocoamsp(model): estimator = CocoaMSP() ue = estimate_uncertainty(model, estimator, INPUT) From 85845dbd93b6c16c0b8d5676d08fb3f4f8d79e53 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Sat, 19 Jul 2025 03:42:24 +0300 Subject: [PATCH 19/23] fix: explicitly cast indices to int in select function --- src/lm_polygraph/utils/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lm_polygraph/utils/dataset.py b/src/lm_polygraph/utils/dataset.py index 6d99437f0..4e4a36018 100644 --- a/src/lm_polygraph/utils/dataset.py +++ b/src/lm_polygraph/utils/dataset.py @@ -50,8 +50,8 @@ def select(self, indices: List[int]): Parameters: indices (List[int]): indices to left in the dataset.Must have the same length as input texts. """ - self.x = [self.x[i] for i in indices] - self.y = [self.y[i] for i in indices] + self.x = [self.x[int(i)] for i in indices] + self.y = [self.y[int(i)] for i in indices] return self def train_test_split(self, test_size: int, seed: int, split: str = "train"): From 042fda6a440ae4e6a719cf6dc8595433e6b13caf Mon Sep 17 00:00:00 2001 From: ahdr3w <71148619+ahdr3w@users.noreply.github.com> Date: Sat, 19 Jul 2025 03:49:25 +0300 Subject: [PATCH 20/23] Delete examples/topological_divergence_with_heads_selection.ipynb --- ...ical_divergence_with_heads_selection.ipynb | 234 ------------------ 1 file changed, 234 deletions(-) delete mode 100644 examples/topological_divergence_with_heads_selection.ipynb diff --git a/examples/topological_divergence_with_heads_selection.ipynb b/examples/topological_divergence_with_heads_selection.ipynb deleted file mode 100644 index f0d558328..000000000 --- a/examples/topological_divergence_with_heads_selection.ipynb +++ /dev/null @@ -1,234 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "id": "bd1b5660", - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/ahdr3w/.pyenv/versions/3.10.11/lib/python3.10/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": [ - "from time import time\n", - "from transformers import AutoModelForCausalLM, AutoTokenizer\n", - "\n", - "from lm_polygraph.estimators import TopologicalDivergence\n", - "\n", - "from lm_polygraph.utils import WhiteboxModel, UEManager, Dataset\n", - "from lm_polygraph.utils.builder_enviroment_stat_calculator import BuilderEnvironmentStatCalculator\n", - "from lm_polygraph.utils.factory_stat_calculator import StatCalculatorContainer\n", - "from lm_polygraph.utils.estimate_uncertainty import UncertaintyOutput\n", - "\n", - "from lm_polygraph.defaults.register_default_stat_calculators import (\n", - " register_default_stat_calculators,\n", - ")\n", - "from lm_polygraph.stat_calculators import TrainMTopDivCalculator\n" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "8d4148d7", - "metadata": {}, - "outputs": [], - "source": [ - "def estimate_uncertainty(\n", - " model, estimator, available_stat_calculators, input_text: str\n", - ") -> UncertaintyOutput:\n", - " man = UEManager(\n", - " Dataset([input_text], [\"\"], batch_size=1),\n", - " model,\n", - " [estimator],\n", - " available_stat_calculators=available_stat_calculators,\n", - " builder_env_stat_calc=BuilderEnvironmentStatCalculator(model),\n", - " generation_metrics=[],\n", - " ue_metrics=[],\n", - " processors=[],\n", - " ignore_exceptions=False,\n", - " verbose=False,\n", - " )\n", - " man()\n", - " ue = man.estimations[estimator.level, str(estimator)]\n", - " texts = man.stats.get(\"greedy_texts\", None)\n", - " tokens = man.stats.get(\"greedy_tokens\", None)\n", - " if tokens is not None and len(tokens) > 0:\n", - " # Remove last token, which is the end of the sequence token\n", - " # since we don't include it's uncertainty in the estimator's output\n", - " tokens = tokens[0][:-1]\n", - " return UncertaintyOutput(\n", - " ue[0], input_text, texts[0], tokens, model.model_path, str(estimator)\n", - " )\n" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "430d4c18", - "metadata": {}, - "outputs": [], - "source": [ - "class Config:\n", - " model_name = 'Qwen/Qwen2.5-0.5B-Instruct'\n", - " \n", - " heads_extraction_priority = \"cache\"\n", - " model_heads_cache = \"model_heads_cache.yaml\"\n", - " \n", - " subsample_train_dataset = 4\n", - " seed = [52]\n", - " device_map = 'mps' \n", - " max_heads = 6\n", - " n_jobs = -1\n", - "\n", - " cache_path = \"./workdir/output\"\n", - "\n", - " task = \"qa\"\n", - " instruct = False\n", - "\n", - " dataset = ['LM-polygraph/coqa', 'continuation']\n", - " text_column = \"input\"\n", - " label_column = \"output\"\n", - " train_split = \"train\"\n", - " eval_split = \"test\"\n", - " n_shot = 0\n", - " few_shot_split = \"train\"\n", - " few_shot_prompt = None\n", - " trust_remote_code = False\n", - " size = None\n", - "\n", - " max_new_tokens = 20\n", - " load_from_disk = False\n", - " generation_params = {\"generate_until\": [\"\\n\"]}\n", - " subsample_eval_dataset = -1\n", - " batch_size = 1\n", - "\n", - " generation_metrics = None\n", - " ignore_exceptions = False\n", - "\n", - "\n", - "\n", - "cfg = Config()\n", - "\n", - "base_model = AutoModelForCausalLM.from_pretrained(\n", - " cfg.model_name,\n", - " device_map=cfg.device_map,\n", - ")\n", - "tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)\n", - "model = WhiteboxModel(base_model, tokenizer, model_path=cfg.model_name)\n", - "available_stat_calculators = (\n", - " register_default_stat_calculators(\"Whitebox\")\n", - ")\n", - "sc = StatCalculatorContainer(\n", - " name=TrainMTopDivCalculator.__name__,\n", - " obj=TrainMTopDivCalculator,\n", - " builder=\"lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator\",\n", - " cfg=cfg,\n", - " dependencies=TrainMTopDivCalculator.meta_info()[1],\n", - " stats=TrainMTopDivCalculator.meta_info()[0],\n", - ")\n", - "available_stat_calculators.append(sc)\n" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "5f927bc7", - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - " 0%| | 0/4 [00:00 Date: Sat, 19 Jul 2025 03:59:09 +0300 Subject: [PATCH 21/23] flake8-black --- .../defaults/register_default_stat_calculators.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index 815d9bbdd..983a41622 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -153,15 +153,15 @@ def _register( "lm_polygraph.defaults.stat_calculator_builders.default_TrainMTopDivCalculator", { "heads_extraction_priority": "cache", - "cache_path": './workdir/output', + "cache_path": "./workdir/output", "model_heads_cache": "model_heads_cache.yaml", "max_heads": 6, "n_jobs": -1, "dataset": ["LM-polygraph/coqa", "continuation"], "text_column": "input", "label_column": "output", - "description": '', - "prompt": '', + "description": "", + "prompt": "", "few_shot_split": "train", "train_split": "train", "load_from_disk": False, From 2d5e5e29437be031ade3cbe85fb790e2b8d840f4 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Sun, 20 Jul 2025 18:47:36 +0300 Subject: [PATCH 22/23] Set n_jobs=1 in Parallel to prevent process leaks and ensure safe exit --- examples/configs/estimators/default_estimators.yaml | 2 +- src/lm_polygraph/defaults/register_default_stat_calculators.py | 2 +- src/lm_polygraph/estimators/mtopdiv_utils.py | 2 +- src/lm_polygraph/estimators/topological_divergence.py | 2 +- src/lm_polygraph/stat_calculators/train_mtopdiv.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/configs/estimators/default_estimators.yaml b/examples/configs/estimators/default_estimators.yaml index c16ccece0..db1e380c8 100644 --- a/examples/configs/estimators/default_estimators.yaml +++ b/examples/configs/estimators/default_estimators.yaml @@ -103,4 +103,4 @@ - name: TopologicalDivergence cfg: heads: null - n_jobs: -1 + n_jobs: 1 diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index 983a41622..cd7fcbc51 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -156,7 +156,7 @@ def _register( "cache_path": "./workdir/output", "model_heads_cache": "model_heads_cache.yaml", "max_heads": 6, - "n_jobs": -1, + "n_jobs": 1, "dataset": ["LM-polygraph/coqa", "continuation"], "text_column": "input", "label_column": "output", diff --git a/src/lm_polygraph/estimators/mtopdiv_utils.py b/src/lm_polygraph/estimators/mtopdiv_utils.py index 3159db1d5..edeeef8d1 100644 --- a/src/lm_polygraph/estimators/mtopdiv_utils.py +++ b/src/lm_polygraph/estimators/mtopdiv_utils.py @@ -80,7 +80,7 @@ def get_mtopdivs( heads: List[Tuple[int, int]], length_responses: List[int], attention_weights_batch: np.array, - n_jobs: Optional[float] = -1, + n_jobs: Optional[float] = 1, ) -> np.ndarray: batch_size = attention_weights_batch.shape[0] padding_lengths = np.isnan(attention_weights_batch[:, 0, 0, 0]).sum(axis=-1) diff --git a/src/lm_polygraph/estimators/topological_divergence.py b/src/lm_polygraph/estimators/topological_divergence.py index 20805fb97..7f5a0fdc2 100644 --- a/src/lm_polygraph/estimators/topological_divergence.py +++ b/src/lm_polygraph/estimators/topological_divergence.py @@ -18,7 +18,7 @@ class TopologicalDivergence(Estimator): def __init__( self, heads: Optional[List[Tuple[int, int]]] = None, - n_jobs: int = -1, + n_jobs: int = 1, ): """ Initializes TopologicalDivergence estimator. diff --git a/src/lm_polygraph/stat_calculators/train_mtopdiv.py b/src/lm_polygraph/stat_calculators/train_mtopdiv.py index 6915853f4..9ec809216 100644 --- a/src/lm_polygraph/stat_calculators/train_mtopdiv.py +++ b/src/lm_polygraph/stat_calculators/train_mtopdiv.py @@ -33,7 +33,7 @@ def __init__( load_train_dataset_fn: Callable = None, cache_path: str = None, max_heads: int = 6, - n_jobs: int = -1, + n_jobs: int = 1, ): super().__init__() From 62a3966a429f5339fecdbb80608b28eec27b7670 Mon Sep 17 00:00:00 2001 From: ahdr3w Date: Fri, 25 Jul 2025 15:34:18 +0300 Subject: [PATCH 23/23] set back n_jobs=-1 --- examples/configs/estimators/default_estimators.yaml | 2 +- .../defaults/register_default_stat_calculators.py | 2 +- test/test_estimators.py | 13 ++++++++++--- 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/examples/configs/estimators/default_estimators.yaml b/examples/configs/estimators/default_estimators.yaml index db1e380c8..c16ccece0 100644 --- a/examples/configs/estimators/default_estimators.yaml +++ b/examples/configs/estimators/default_estimators.yaml @@ -103,4 +103,4 @@ - name: TopologicalDivergence cfg: heads: null - n_jobs: 1 + n_jobs: -1 diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index cd7fcbc51..983a41622 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -156,7 +156,7 @@ def _register( "cache_path": "./workdir/output", "model_heads_cache": "model_heads_cache.yaml", "max_heads": 6, - "n_jobs": 1, + "n_jobs": -1, "dataset": ["LM-polygraph/coqa", "continuation"], "text_column": "input", "label_column": "output", diff --git a/test/test_estimators.py b/test/test_estimators.py index 86915212c..b8ebed98d 100644 --- a/test/test_estimators.py +++ b/test/test_estimators.py @@ -1,5 +1,6 @@ import torch import pytest +from joblib.externals.loky import get_reusable_executor from transformers import AutoModelForCausalLM, AutoTokenizer @@ -253,18 +254,24 @@ def test_attentionscore(model): def test_topological_divergence_fixed_heads(model): - estimator = TopologicalDivergence(heads=[[0, 0], [0, 1], [1, 0], [1, 1]]) + estimator = TopologicalDivergence( + heads=[[0, 0], [0, 1], [1, 0], [1, 1]], + n_jobs=-1, + ) ue = estimate_uncertainty(model, estimator, INPUT) + get_reusable_executor().shutdown(wait=True) assert isinstance(ue.uncertainty, float) def test_topological_divergence_select_heads(model): - estimator = TopologicalDivergence() + estimator = TopologicalDivergence(n_jobs=-1) ue = estimate_uncertainty(model, estimator, INPUT) + get_reusable_executor().shutdown(wait=True) assert isinstance(ue.uncertainty, float) # second time heads should be loaded from cache - estimator = TopologicalDivergence() + estimator = TopologicalDivergence(n_jobs=-1) ue = estimate_uncertainty(model, estimator, INPUT) + get_reusable_executor().shutdown(wait=True) assert isinstance(ue.uncertainty, float)