diff --git a/.gitignore b/.gitignore index 238da6459..046586a0d 100644 --- a/.gitignore +++ b/.gitignore @@ -179,3 +179,6 @@ cython_debug/ # Agent instructions **/AGENTS.md + +#vscode +.vscode/ \ No newline at end of file diff --git a/README.md b/README.md index 3483d484e..836848e25 100644 --- a/README.md +++ b/README.md @@ -186,6 +186,8 @@ UE methods such as `EigValLaplacian()` support fully blackbox LLMs that do not p | Eccentricity (Ecc) [(Lin et al., 2023)](https://arxiv.org/abs/2305.19187) | Black-box | Meaning Diversity | High | Low | No | sequence | | Lexical similarity (LexSim) [(Fomicheva et al., 2020a)](https://direct.mit.edu/tacl/article/doi/10.1162/tacl_a_00330/96475/Unsupervised-Quality-Estimation-for-Neural-Machine) | Black-box | Meaning Diversity | High | Low | No | sequence | | Kernel Language Entropy [(Nikitin et al., 2024)](https://arxiv.org/pdf/2405.20003) | Black-box | Meaning Diversity | High | Low | No | sequence | +| Predictive Kernel Entropy (PKE) [(Gruber et al., 2024)](https://arxiv.org/abs/2310.05833) | Black-box | Meaning Diversity | High | Low | No | sequence | +| Spectral Uncertainty [(Walha et al., 2025)](https://arxiv.org/abs/2509.22272) | Black-box | Meaning Diversity | High | Low | No | sequence | | LUQ [(Zhang et al., 2024)](https://aclanthology.org/2024.emnlp-main.299/) | Black-box | Meaning diversity | High | Low | No | sequence | | Verbalized Uncertainty 1S [(Tian et al., 2023)](https://arxiv.org/abs/2305.14975) | Black-box | Reflexive | Low | Low | No | sequence | | Verbalized Uncertainty 2S [(Tian et al., 2023)](https://arxiv.org/abs/2305.14975) | Black-box | Reflexive | Medium | Low | No | sequence | diff --git a/examples/basic_example_visual.ipynb b/examples/basic_example_visual.ipynb index d6de0a83b..22c2974d3 100644 --- a/examples/basic_example_visual.ipynb +++ b/examples/basic_example_visual.ipynb @@ -98,7 +98,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "6958a441", "metadata": {}, "outputs": [ @@ -115,7 +115,12 @@ "%load_ext autoreload\n", "%autoreload 2\n", "\n", - "from transformers import AutoModelForVision2Seq, AutoProcessor\n", + "from transformers import AutoProcessor\n", + "try:\n", + " from transformers import AutoModelForVision2Seq\n", + "except ImportError:\n", + " # transformers >= 5.0 renamed AutoModelForVision2Seq → AutoModelForImageTextToText\n", + " from transformers import AutoModelForImageTextToText as AutoModelForVision2Seq\n", "from lm_polygraph.model_adapters.visual_whitebox_model import VisualWhiteboxModel\n", "from lm_polygraph import estimate_uncertainty\n", "from lm_polygraph.estimators import MaximumTokenProbability, MaximumSequenceProbability, SemanticEntropy, EigValLaplacian" diff --git a/examples/configs/estimators/default_estimators.yaml b/examples/configs/estimators/default_estimators.yaml index e175947f3..8fac1cb17 100644 --- a/examples/configs/estimators/default_estimators.yaml +++ b/examples/configs/estimators/default_estimators.yaml @@ -9,6 +9,8 @@ - name: PTrueSampling - name: MonteCarloSequenceEntropy - name: MonteCarloNormalizedSequenceEntropy +- name: PredictiveKernelEntropy +- name: SpectralUncertainty - name: LexicalSimilarity cfg: metric: "rouge1" diff --git a/examples/configs/model/default_visual.py b/examples/configs/model/default_visual.py index 0e45b3fff..4002a9544 100644 --- a/examples/configs/model/default_visual.py +++ b/examples/configs/model/default_visual.py @@ -1,4 +1,10 @@ -from transformers import AutoModelForVision2Seq, AutoProcessor +from transformers import AutoProcessor + +try: + from transformers import AutoModelForVision2Seq +except ImportError: + # transformers >= 5.0 renamed AutoModelForVision2Seq → AutoModelForImageTextToText + from transformers import AutoModelForImageTextToText as AutoModelForVision2Seq def load_model(model_path: str, device_map: str): diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index f3ddabd82..49c921510 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -56,6 +56,19 @@ def _register( "device": None, } + # Shared sentence embedding model defaultconfig + sentence_embedding_model_cfg = { + "model_name": "all-mpnet-base-v2", + "device": None, + "cache_folder": None, + "encoding_arguments": { + "batch_size": 256, + "convert_to_numpy": True, + "convert_to_tensor": False, + "show_progress_bar": False, + }, + } + _register(InitialStateCalculator) _register(RawInputCalculator) _register( @@ -74,7 +87,11 @@ def _register( {"nli_model": nli_model_cfg}, ) _register(SemanticClassesCalculator) - + _register( + SampleSentenceEmbeddingsCalculator, + "lm_polygraph.defaults.stat_calculator_builders.default_SampleSentenceEmbeddingsCalculator", + {"sentence_embedding_model": sentence_embedding_model_cfg}, + ) if model_type == "Blackbox": _register(BlackboxGreedyTextsCalculator) _register(BlackboxSamplingGenerationCalculator) diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/default_SampleSentenceEmbeddingsCalculator.py b/src/lm_polygraph/defaults/stat_calculator_builders/default_SampleSentenceEmbeddingsCalculator.py new file mode 100644 index 000000000..edc97d3d0 --- /dev/null +++ b/src/lm_polygraph/defaults/stat_calculator_builders/default_SampleSentenceEmbeddingsCalculator.py @@ -0,0 +1,12 @@ +from lm_polygraph.stat_calculators.sample_sentence_embeddings import ( + SampleSentenceEmbeddingsCalculator, +) +from .utils import load_sentence_embedding_model + + +def load_stat_calculator(config, builder): + if not hasattr(builder, "sentence_embedding_model"): + builder.sentence_embedding_model = load_sentence_embedding_model( + **config.sentence_embedding_model + ) + return SampleSentenceEmbeddingsCalculator(builder.sentence_embedding_model) diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/utils.py b/src/lm_polygraph/defaults/stat_calculator_builders/utils.py index bff96354e..03d5356f2 100644 --- a/src/lm_polygraph/defaults/stat_calculator_builders/utils.py +++ b/src/lm_polygraph/defaults/stat_calculator_builders/utils.py @@ -1,6 +1,8 @@ import logging +from typing import Dict, Optional from lm_polygraph.utils.deberta import Deberta, MultilingualDeberta +from lm_polygraph.utils.sentence_embedder import SentenceEmbedder log = logging.getLogger("lm_polygraph") @@ -21,3 +23,19 @@ def load_nli_model( f"Initialized {nli_model.deberta_path} on {nli_model.device} with batch_size={nli_model.batch_size}" ) return nli_model + + +def load_sentence_embedding_model( + model_name: str = "all-mpnet-base-v2", + device: Optional[str] = None, + cache_folder: Optional[str] = None, + encoding_arguments: Optional[Dict] = None, +) -> SentenceEmbedder: + model = SentenceEmbedder( + model_name, + device=device, + cache_folder=cache_folder, + encoding_arguments=encoding_arguments, + ) + log.info(f"Initialized SentenceEmbedder({model_name}) on {model.device}") + return model diff --git a/src/lm_polygraph/estimators/__init__.py b/src/lm_polygraph/estimators/__init__.py index ee93389a5..c05cef6df 100644 --- a/src/lm_polygraph/estimators/__init__.py +++ b/src/lm_polygraph/estimators/__init__.py @@ -90,3 +90,5 @@ from .rauq import RAUQ from .csl import CSL from .semantic_density import SemanticDensity +from .predictive_kernel_entropy import PredictiveKernelEntropy +from .spectral_uncertainty import SpectralUncertainty diff --git a/src/lm_polygraph/estimators/predictive_kernel_entropy.py b/src/lm_polygraph/estimators/predictive_kernel_entropy.py new file mode 100644 index 000000000..18ee7983d --- /dev/null +++ b/src/lm_polygraph/estimators/predictive_kernel_entropy.py @@ -0,0 +1,59 @@ +import numpy as np + +from typing import Dict + +from .estimator import Estimator +import sklearn.metrics + + +class PredictiveKernelEntropy(Estimator): + """ + Predictive Kernel Entropy (PKE) uncertainty estimator. + + For each input, n answers are sampled and encoded with an external sentence + embedding model. Uncertainty is the negative mean pairwise kernel similarity: + + PKE = -1 / (n*(n-1)) * sum_{i != j} k(X_i, X_j) + + Higher values indicate higher uncertainty. + + Reference: https://arxiv.org/abs/2310.05833 + + Parameters: + kernel (str): Kernel function — one of "rbf", "laplacian", "cosine". Default "rbf". + gamma (float): Gamma parameter for the RBF and the laplacian kernels. Ignored for the cosine kernel. + """ + + def __init__(self, kernel: str = "rbf", gamma: float = 1.0): + super().__init__(["sample_sentence_embeddings"], "sequence") + if kernel not in ("rbf", "laplacian", "cosine"): + raise ValueError( + f"Unknown kernel '{kernel}'. Choose from: rbf, laplacian, cosine." + ) + self.kernel = kernel + self.gamma = gamma + if kernel == "rbf": + self.kernel_function = lambda x: sklearn.metrics.pairwise.rbf_kernel( + x, gamma=gamma + ) + elif kernel == "laplacian": + self.kernel_function = lambda x: sklearn.metrics.pairwise.laplacian_kernel( + x, gamma=gamma + ) + elif kernel == "cosine": + self.kernel_function = sklearn.metrics.pairwise.linear_kernel + + def __str__(self) -> str: + return f"PredictiveKernelEntropy(kernel={self.kernel}, gamma={self.gamma})" + + def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: + scores = [] + for embeddings in stats["sample_sentence_embeddings"]: + E = np.array(embeddings) # (n, d) + n = len(E) + if n < 2: + scores.append(0.0) + continue + K = self.kernel_function(E) # (n, n) + scores.append((K.diagonal().sum() - K.sum()) / (n * (n - 1))) + return np.array(scores) diff --git a/src/lm_polygraph/estimators/spectral_uncertainty.py b/src/lm_polygraph/estimators/spectral_uncertainty.py new file mode 100644 index 000000000..09783c878 --- /dev/null +++ b/src/lm_polygraph/estimators/spectral_uncertainty.py @@ -0,0 +1,85 @@ +import numpy as np + +from typing import Dict + +from .estimator import Estimator +import sklearn.metrics +import scipy + + +class SpectralUncertainty(Estimator): + """ + Spectral Uncertainty estimator. + + For each input, n answers are sampled and encoded with an external sentence + embedding model. The empirical kernel matrix K is built from pairwise kernel + values between embeddings, the kernel matrix is normalized to unit trace, and + as a result, eigenvalues of K can be interpreted as a probability distribution + over latent semantic directions. Finally, uncertainty is computed as the Shannon + entropy of that distribution: + + H = -Σ λᵢ log(λᵢ) + + Higher values indicate higher uncertainty (more spread in the spectrum). + + Reference: https://arxiv.org/abs/2509.22272 + + Parameters: + kernel (str): Kernel function — one of "rbf", "laplacian", "cosine". Default "rbf". + gamma (float): Gamma parameter for rbf/laplacian kernels. Ignored for cosine. + eps (float): Small value to be added to the diagonal to avoid numerical issues. + """ + + def __init__(self, kernel: str = "rbf", gamma: float = 1.0, eps: float = 1e-6): + super().__init__(["sample_sentence_embeddings"], "sequence") + if kernel not in ("rbf", "laplacian", "cosine"): + raise ValueError( + f"Unknown kernel '{kernel}'. Choose from: rbf, laplacian, cosine." + ) + self.kernel = kernel + self.gamma = gamma + self.eps = eps + if kernel == "rbf": + self.kernel_function = lambda x: sklearn.metrics.pairwise.rbf_kernel( + x, gamma=gamma + ) + elif kernel == "laplacian": + self.kernel_function = lambda x: sklearn.metrics.pairwise.laplacian_kernel( + x, gamma=gamma + ) + elif kernel == "cosine": + self.kernel_function = sklearn.metrics.pairwise.linear_kernel + + def __str__(self) -> str: + return f"SpectralUncertainty(kernel={self.kernel}, gamma={self.gamma})" + + def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: + scores = [] + for embeddings in stats["sample_sentence_embeddings"]: + E = np.array(embeddings) # (n, d) + n = len(E) + if n < 2: + scores.append(0.0) + continue + K = self.kernel_function(E) # (n, n) symmetric + K /= n # Normalizing to trace =1. Here we assume that the embeddings are normalized, so the trace is n. + try: + eigenvalues, _ = scipy.linalg.eigh(K) + except np.linalg.LinAlgError: + print( + f"Linalg error happened because of condition {np.linalg.cond(K)}. Trying numpy function." + ) + try: + eigenvalues, _ = np.linalg.eigh(K) + except np.linalg.LinAlgError: + epsilon = 1e-6 # You can tune this + K += epsilon * np.eye(K.shape[0]) + print( + f"Numy function failed. Adding {epsilon} to the diagonal. New condition is {np.linalg.cond(K)}" + ) + eigenvalues, _ = scipy.linalg.eigh(K) + # Due to numerical issues, sometimes negative eigenvalues appear (although theoretically impossible because PSD). These are very marginal and are clipped to 0. + eigenvalues = np.array([(ev if ev >= 0 else 0.0) for ev in eigenvalues]) + H = scipy.stats.entropy(eigenvalues) + scores.append(H) + return np.array(scores) diff --git a/src/lm_polygraph/stat_calculators/__init__.py b/src/lm_polygraph/stat_calculators/__init__.py index 793ee34fd..ef6c1502b 100644 --- a/src/lm_polygraph/stat_calculators/__init__.py +++ b/src/lm_polygraph/stat_calculators/__init__.py @@ -62,3 +62,4 @@ from .semantic_classes import SemanticClassesCalculator from .attention_forward_pass_visual import AttentionForwardPassCalculatorVisual from .vllm_logprobs_extraction import VLLMLogprobsExtractionCalculator +from .sample_sentence_embeddings import SampleSentenceEmbeddingsCalculator diff --git a/src/lm_polygraph/stat_calculators/sample_sentence_embeddings.py b/src/lm_polygraph/stat_calculators/sample_sentence_embeddings.py new file mode 100644 index 000000000..10072126a --- /dev/null +++ b/src/lm_polygraph/stat_calculators/sample_sentence_embeddings.py @@ -0,0 +1,39 @@ +import numpy as np + +from typing import Dict, List, Tuple + +from .stat_calculator import StatCalculator +from lm_polygraph.utils.model import Model +from lm_polygraph.utils.sentence_embedder import SentenceEmbedder + + +class SampleSentenceEmbeddingsCalculator(StatCalculator): + """ + Encodes each sampled generation with an external sentence embedding model. + Produces `sample_sentence_embeddings`: List[np.ndarray] of shape (n_samples, embedding_dim). + Used by PredictiveKernelEntropy and any other estimator that needs dense + embeddings of sampled texts rather than the LLM's own hidden states. + """ + + @staticmethod + def meta_info() -> Tuple[List[str], List[str]]: + return ["sample_sentence_embeddings"], ["sample_texts"] + + def __init__(self, model: SentenceEmbedder): + super().__init__() + self._embedder = model + + def __call__( + self, + dependencies: Dict[str, np.ndarray], + texts: List[str], + model: Model, + max_new_tokens: int = 100, + **kwargs, + ) -> Dict[str, np.ndarray]: + return { + "sample_sentence_embeddings": [ + self._embedder.encode(samples) + for samples in dependencies["sample_texts"] + ] + } diff --git a/src/lm_polygraph/utils/factory_estimator.py b/src/lm_polygraph/utils/factory_estimator.py index 7a8a52bfb..4603eb2c3 100644 --- a/src/lm_polygraph/utils/factory_estimator.py +++ b/src/lm_polygraph/utils/factory_estimator.py @@ -60,6 +60,8 @@ def load_simple_estimators(name: str, config): CSL, SemanticDensity, BoostedProbSequence, + PredictiveKernelEntropy, + SpectralUncertainty, ] try: diff --git a/src/lm_polygraph/utils/sentence_embedder.py b/src/lm_polygraph/utils/sentence_embedder.py new file mode 100644 index 000000000..2dbcc2b07 --- /dev/null +++ b/src/lm_polygraph/utils/sentence_embedder.py @@ -0,0 +1,53 @@ +import numpy as np +import torch + +from typing import Dict, List, Optional +from sentence_transformers import SentenceTransformer + + +class SentenceEmbedder: + """ + Wrapper around SentenceTransformer that manages model loading and encoding + arguments. Analogous to the Deberta wrapper — intended to be constructed + once and shared across stat calculators via BuilderEnvironmentStatCalculator. + """ + + _DEFAULT_ENCODING_ARGUMENTS: Dict = { + "batch_size": 256, + "convert_to_numpy": True, + "convert_to_tensor": False, + "show_progress_bar": False, + } + + def __init__( + self, + model_name: str = "all-mpnet-base-v2", + device: Optional[str] = None, + cache_folder: Optional[str] = None, + encoding_arguments: Optional[Dict] = None, + ): + self.model_name = model_name + if device is None: + self.device = "cuda:0" if torch.cuda.is_available() else "cpu" + else: + self.device = device + self.cache_folder = cache_folder + self.encoding_arguments = { + **self._DEFAULT_ENCODING_ARGUMENTS, + **(encoding_arguments or {}), + } + self._model = None + self.setup() + + def setup(self): + if self._model is not None: + return + + self._model = SentenceTransformer( + self.model_name, + device=self.device, + cache_folder=self.cache_folder, + ) + + def encode(self, samples: List[str]) -> np.ndarray: + return self._model.encode(samples, **self.encoding_arguments) diff --git a/test/configs/test_default_estimators.yaml b/test/configs/test_default_estimators.yaml index 1fada8483..d8ce529c1 100644 --- a/test/configs/test_default_estimators.yaml +++ b/test/configs/test_default_estimators.yaml @@ -10,6 +10,8 @@ estimators: - name: PTrueSampling - name: MonteCarloSequenceEntropy - name: MonteCarloNormalizedSequenceEntropy + - name: PredictiveKernelEntropy + - name: SpectralUncertainty - name: LexicalSimilarity cfg: metric: "rouge1" diff --git a/test/local/test_estimators_visual.py b/test/local/test_estimators_visual.py index 53ec106dd..dc84c2d63 100644 --- a/test/local/test_estimators_visual.py +++ b/test/local/test_estimators_visual.py @@ -1,8 +1,13 @@ import torch import pytest -from transformers import AutoModelForVision2Seq, AutoProcessor +from transformers import AutoProcessor +try: + from transformers import AutoModelForVision2Seq +except ImportError: + # transformers >= 5.0 renamed AutoModelForVision2Seq → AutoModelForImageTextToText + from transformers import AutoModelForImageTextToText as AutoModelForVision2Seq from lm_polygraph import estimate_uncertainty from lm_polygraph.estimators import * from lm_polygraph.model_adapters.visual_whitebox_model import VisualWhiteboxModel @@ -178,6 +183,18 @@ def test_eigenscore(model): assert isinstance(ue.uncertainty, float) +def test_predictive_kernel_entropy(model): + estimator = PredictiveKernelEntropy() + ue = estimate_uncertainty(model, estimator, INPUT, IMAGES) + assert isinstance(ue.uncertainty, float) + + +def test_spectral_uncertainty(model): + estimator = SpectralUncertainty() + ue = estimate_uncertainty(model, estimator, INPUT, IMAGES) + assert isinstance(ue.uncertainty, float) + + def test_focus(model): model_name = model.model.config._name_or_path estimator = Focus( diff --git a/test/test_estimators.py b/test/test_estimators.py index 86ca66e95..bba78de09 100644 --- a/test/test_estimators.py +++ b/test/test_estimators.py @@ -288,6 +288,12 @@ def test_semantic_density_concat(model): assert isinstance(ue.uncertainty, float) +def test_predictive_kernel_entropy(model): + estimator = PredictiveKernelEntropy() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float) + + def test_semantic_density(model): estimator = SemanticDensity(concat_input=False) ue = estimate_uncertainty(model, estimator, INPUT) @@ -298,3 +304,9 @@ def test_boostedprob_sequence(model): estimator = BoostedProbSequence() ue = estimate_uncertainty(model, estimator, INPUT) assert isinstance(ue.uncertainty, float) + + +def test_spectral_uncertainty(model): + estimator = SpectralUncertainty() + ue = estimate_uncertainty(model, estimator, INPUT) + assert isinstance(ue.uncertainty, float)