From d6e5877d039b5b58fd5e404ec238fdf19786632a Mon Sep 17 00:00:00 2001 From: Ekaterina Fadeeva Date: Tue, 24 Feb 2026 12:32:37 +0100 Subject: [PATCH 1/2] Graybox CCP --- .../claim_conditioned_probability.py | 8 +- .../model_adapters/togetherai_adapter.py | 4 +- .../stat_calculators/graybox_stats.py | 108 ++++++++++++++++++ 3 files changed, 112 insertions(+), 8 deletions(-) create mode 100644 src/lm_polygraph/stat_calculators/graybox_stats.py diff --git a/src/lm_polygraph/estimators/claim_conditioned_probability.py b/src/lm_polygraph/estimators/claim_conditioned_probability.py index d2640a8e8..0c26a7c67 100644 --- a/src/lm_polygraph/estimators/claim_conditioned_probability.py +++ b/src/lm_polygraph/estimators/claim_conditioned_probability.py @@ -9,7 +9,6 @@ class ClaimConditionedProbability(Estimator): def __init__(self): super().__init__( [ - "greedy_tokens", "greedy_tokens_alternatives", "greedy_tokens_alternatives_nli", ], @@ -38,18 +37,15 @@ def _combine_nli(self, forward: str, backward: str): return "neutral" def __call__(self, stats: Dict[str, np.ndarray]) -> np.ndarray: - words = stats["greedy_tokens"] alternatives = stats["greedy_tokens_alternatives"] alternatives_nli = stats["greedy_tokens_alternatives_nli"] prob_nli = [] - for sample_words, sample_alternatives, sample_alternatives_nli in zip( - words, + for sample_alternatives, sample_alternatives_nli in zip( alternatives, alternatives_nli, ): sample_mnlis = [] - for word, word_alternatives, word_alternatives_nli in zip( - sample_words, + for word_alternatives, word_alternatives_nli in zip( sample_alternatives, sample_alternatives_nli, ): diff --git a/src/lm_polygraph/model_adapters/togetherai_adapter.py b/src/lm_polygraph/model_adapters/togetherai_adapter.py index ebdc9626f..a9c89a4b6 100644 --- a/src/lm_polygraph/model_adapters/togetherai_adapter.py +++ b/src/lm_polygraph/model_adapters/togetherai_adapter.py @@ -52,7 +52,7 @@ def generate_texts( break except Exception as e: if retries > 4: - raise Exception from e + raise e retries += 1 continue @@ -60,7 +60,7 @@ def generate_texts( [self.parse_response(resp) for resp in response.choices] ) - return parsed_responses + return parsed_responses @register_adapter("together_ai") class TogetherAIAdapter(TogetherAIChatCompletionMixin, APIProviderAdapter): diff --git a/src/lm_polygraph/stat_calculators/graybox_stats.py b/src/lm_polygraph/stat_calculators/graybox_stats.py new file mode 100644 index 000000000..98dd7c12b --- /dev/null +++ b/src/lm_polygraph/stat_calculators/graybox_stats.py @@ -0,0 +1,108 @@ +import numpy as np +from typing import Dict, List, Tuple, Any + +from .stat_calculator import StatCalculator +from lm_polygraph.model_adapters.api_provider_adapter import APIProviderAdapter + + +class GrayBoxStatsCalculator(StatCalculator): + """ + Provider-based graybox stats calculator (API models). + + Computes a minimal set of generation statistics for a batch of input texts: + + - greedy_texts: + The generated text for each input (typically the first returned choice). + + - greedy_tokens_alternatives: + Per generated token position, the adapter-provided top-k alternatives as + (token, logprob) pairs. + + Shape: + List[batch] -> List[position] -> List[(token_str, logprob_float)] + """ + + @staticmethod + def meta_info() -> Tuple[List[str], List[str]]: + return [ + "greedy_texts", + "greedy_tokens_alternatives", + ], [] + + def __init__(self, n_alternatives: int = 10): + super().__init__() + self.n_alternatives = n_alternatives + + def __call__( + self, + dependencies: Dict[str, Any], + texts: List[str], + model: APIProviderAdapter, + max_new_tokens: int = 100, + ) -> Dict[str, Any]: + """ + Generate continuations for `texts` using an API adapter and store minimal stats. + + Parameters + ---------- + dependencies: dict used to pass inputs and store computed statistics. + Must contain: + - dependencies["model"]: the underlying model wrapper/object used by the adapter + (e.g., has .prepare_input(...) and .model_path, depending on your design). + + texts: Batch of input prompts. + + model: The APIProviderAdapter instance. + + max_new_tokens: + Max number of tokens to generate. + + Returns + ------- + Dict[str, Any] + The updated dependencies dict with: + - "greedy_texts" + - "greedy_tokens_alternatives" + """ + out = model.generate_texts( + model=dependencies["model"], + input_texts=texts, + args={ + "max_tokens": max_new_tokens, + "n": 1, + "output_scores": True, + "top_logprobs": self.n_alternatives, + }, + ) + + greedy_texts: List[str] = [] + greedy_tokens_alternatives: List[List[List[Tuple[str, float]]]] = [] + + # out is expected to be: List[prompt] -> List[choice] -> StandardizedResponse + for prompt_choices in out: + if not prompt_choices: + greedy_texts.append("") + greedy_tokens_alternatives.append([]) + continue + + # Take the first choice as "greedy" (since n=1) + resp = prompt_choices[0] + + greedy_texts.append(getattr(resp, "text", "") or "") + + alts_per_pos: List[List[Tuple[str, float]]] = [] + alt_tokens = getattr(resp, "alternative_tokens", None) + top_lps = getattr(resp, "top_logprobs", None) + + if alt_tokens and top_lps: + # alt_tokens: List[List[str]] + # top_lps: List[List[float]] + for toks, lps in zip(alt_tokens, top_lps): + # Pair token strings with their logprobs + alts_per_pos.append(list(zip(toks, lps))) + + greedy_tokens_alternatives.append(alts_per_pos) + + dependencies["greedy_texts"] = greedy_texts + dependencies["greedy_tokens_alternatives"] = greedy_tokens_alternatives + return dependencies \ No newline at end of file From cf45b6a5ac2776bca4e28fb03c82229496944ec3 Mon Sep 17 00:00:00 2001 From: Ekaterina Fadeeva Date: Tue, 24 Feb 2026 12:48:39 +0100 Subject: [PATCH 2/2] Example --- examples/other/example_graybox_ccp.ipynb | 131 +++++++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 examples/other/example_graybox_ccp.ipynb diff --git a/examples/other/example_graybox_ccp.ipynb b/examples/other/example_graybox_ccp.ipynb new file mode 100644 index 000000000..24a9ad16a --- /dev/null +++ b/examples/other/example_graybox_ccp.ipynb @@ -0,0 +1,131 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "53110d78-413d-4d0d-a22d-696b99a18531", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], + "source": [ + "from lm_polygraph.model_adapters.togetherai_adapter import TogetherAIAdapter\n", + "import os\n", + "\n", + "os.environ[\"TOGETHER_API_KEY\"] = \"MY_TOKEN\"\n", + "\n", + "adapter = TogetherAIAdapter()" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "51955d69-6678-4a7b-a894-48d3821d6c50", + "metadata": {}, + "outputs": [], + "source": [ + "class MyModel:\n", + " def __init__(self, model_path: str = \"google/gemma-3n-E4B-it\"):\n", + " self.model_path = model_path\n", + " \n", + " def prepare_input(self, prompt):\n", + " return [\n", + " {\"role\": \"user\", \"content\": prompt},\n", + " ]" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "c37830de-d567-407d-b5b8-53dbeef6fb7f", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Some weights of the model checkpoint at microsoft/deberta-large-mnli were not used when initializing DebertaForSequenceClassification: ['config']\n", + "- This IS expected if you are initializing DebertaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n", + "- This IS NOT expected if you are initializing DebertaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n" + ] + } + ], + "source": [ + "from lm_polygraph.stat_calculators.graybox_stats import GrayBoxStatsCalculator\n", + "from lm_polygraph.stat_calculators import GreedyAlternativesNLICalculator\n", + "from lm_polygraph.estimators import ClaimConditionedProbability\n", + "from lm_polygraph.utils.deberta import Deberta\n", + "\n", + "nli_model = Deberta()\n", + "\n", + "model = MyModel()\n", + "stat_calculators = [\n", + " GrayBoxStatsCalculator(),\n", + " GreedyAlternativesNLICalculator(nli_model),\n", + "]\n", + "estimator = ClaimConditionedProbability()" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a6048572-471b-4442-839d-318b1ff26645", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([-0.99399466])" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "deps = {'model': model}\n", + "texts = ['2+2=?']\n", + "for stat_calculator in stat_calculators:\n", + " deps.update(stat_calculator(deps, texts, adapter))\n", + "estimator(deps)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ca8d48a7-b110-422e-ab68-b1ae6ca6f500", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}