Graybox CCP

Author: cant-access-rediska0123

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