Graybox CCP

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
+}

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,
             ):

src/lm_polygraph/model_adapters/togetherai_adapter.py

@@ -52,15 +52,15 @@ def generate_texts(
                     break
                 except Exception as e:
                     if retries > 4:
-                        raise Exception from e
+                        raise e
                     retries += 1
                     continue
 
             parsed_responses.append(
                 [self.parse_response(resp) for resp in response.choices]
             )
 
-            return parsed_responses
+        return parsed_responses
 
 @register_adapter("together_ai")
 class TogetherAIAdapter(TogetherAIChatCompletionMixin, APIProviderAdapter):

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