review updates

Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>

Author: AngeloDanducci

mellea/backends/backend.py

@@ -47,3 +47,4 @@ def __init__(
         self.model_id = model_id
         self.model_options = model_options if model_options is not None else {}
         self.formatter: ChatFormatter = formatter
+        self._warned_about: set[str] = set()

mellea/backends/huggingface.py

@@ -651,13 +651,48 @@ async def _generate_from_intrinsic(
             user_params["tokenizer"] = self._tokenizer
         generate_input.update(user_params)
 
+        want_scores = bool(
+            model_options.get(ModelOption.LOGITS)
+            or model_options.get(ModelOption.RAW_LOGITS)
+        )
+        if model_options.get(ModelOption.LOGITS):
+            generate_input["output_scores"] = True
+        if model_options.get(ModelOption.RAW_LOGITS):
+            generate_input["output_logits"] = True
+
+        # When logits are requested, intercept the raw GenerateDecoderOnlyOutput that
+        # generate_with_transformers produces internally but never returns (it wraps
+        # everything into a ChatCompletionResponse). We proxy self._model so that
+        # .generate() stores the raw output in raw_hf_output_cell before returning it;
+        # granite_formatters_processing then writes it to mot._meta["hf_output"] for
+        # _surface_logits in post_processing.
+        raw_hf_output_cell: list[GenerateDecoderOnlyOutput | None] = [None]
+
+        model_arg = self._model
+        if want_scores:
+            _real_model = self._model
+
+            class _CapturingModelProxy:
+                def generate(self_proxy, *args: Any, **kwargs: Any) -> Any:
+                    result = cast(Callable[..., Any], _real_model.generate)(
+                        *args, **kwargs
+                    )
+                    if isinstance(result, GenerateDecoderOnlyOutput):
+                        raw_hf_output_cell[0] = result
+                    return result
+
+                def __getattr__(self_proxy, name: str) -> Any:
+                    return getattr(_real_model, name)
+
+            model_arg = _CapturingModelProxy()  # type: ignore[assignment]
+
         chat_response = asyncio.to_thread(
             self._generate_with_adapter_lock,
             adapter.qualified_name,
             granite_formatters.base.util.generate_with_transformers,  # type: ignore
             # Passed as args/kwargs to generate.
             self._tokenizer,
-            self._model,
+            model_arg,
             generate_input,
             other_input,
         )
@@ -681,8 +716,18 @@ async def granite_formatters_processing(
             except json.JSONDecodeError as e:
                 raise Exception(f"Intrinsic did not return a JSON: {chunk}") from e
 
-            # TODO: If we want to support caches, we need to get the GenerateDecoderOnlyOutput. This means we
-            #       probably need to break out the pieces from `generate_with_transformers`.
+            # If logits were requested, stash the intercepted raw output so that
+            # post_processing/_surface_logits can populate generation.logits/raw_logits.
+            if want_scores:
+                if raw_hf_output_cell[0] is not None:
+                    mot._meta["hf_output"] = raw_hf_output_cell[0]
+                else:
+                    MelleaLogger.get_logger().warning(
+                        "ModelOption.LOGITS/RAW_LOGITS requested on intrinsic path but "
+                        "generate_with_transformers did not return a GenerateDecoderOnlyOutput; "
+                        "generation.logits and generation.raw_logits will be None."
+                    )
+
             # processing expects a str or a GenerateDecoderOnlyOutput. Extract the str.
             return await self.processing(
                 mot, res.choices[0].message.content, input_ids=input_ids
@@ -941,6 +986,12 @@ async def _generate_from_context_with_kv_cache(
                 # transformers' generate() requires a tokenizer to decode stop_strings.
                 generate_kwargs["tokenizer"] = self._tokenizer
 
+            kv_scores_kwargs: dict[str, Any] = {}
+            if model_options.get(ModelOption.LOGITS):
+                kv_scores_kwargs["output_scores"] = True
+            if model_options.get(ModelOption.RAW_LOGITS):
+                kv_scores_kwargs["output_logits"] = True
+
             chat_response = asyncio.to_thread(
                 self._generate_with_adapter_lock,
                 "",  # Empty for no adapters.
@@ -951,7 +1002,7 @@ async def _generate_from_context_with_kv_cache(
                 past_key_values=merged_cache,
                 attention_mask=attention_mask.to(self._device),
                 return_dict_in_generate=True,
-                output_scores=True,
+                **kv_scores_kwargs,
                 **generate_kwargs,
                 **streaming_kwargs,  # type: ignore
                 **format_kwargs,  # type: ignore
@@ -1112,6 +1163,12 @@ async def _generate_from_context_standard(
                 # transformers' generate() requires a tokenizer to decode stop_strings.
                 generate_kwargs["tokenizer"] = self._tokenizer
 
+            scores_kwargs: dict[str, Any] = {}
+            if model_options.get(ModelOption.LOGITS):
+                scores_kwargs["output_scores"] = True
+            if model_options.get(ModelOption.RAW_LOGITS):
+                scores_kwargs["output_logits"] = True
+
             chat_response = asyncio.to_thread(
                 self._generate_with_adapter_lock,
                 "",  # Empty for no adapters.
@@ -1121,6 +1178,7 @@ async def _generate_from_context_standard(
                 attention_mask=input_ids["attention_mask"],
                 return_dict_in_generate=True,
                 use_cache=self._use_caches,  # Only create KV cache if caching is enabled
+                **scores_kwargs,
                 **generate_kwargs,
                 **streaming_kwargs,  # type: ignore
                 **format_kwargs,  # type: ignore
@@ -1220,6 +1278,50 @@ async def processing(
                 ),
             )
 
+    def _surface_logits(
+        self, mot: ModelOutputThunk, hf_output: GenerateDecoderOnlyOutput
+    ) -> None:
+        """Populate mot.generation.logits and/or raw_logits from hf_output when requested.
+
+        Checks ModelOption.LOGITS (processed scores) and ModelOption.RAW_LOGITS (raw
+        LM-head logits) in mot._model_options. Skips population when
+        num_return_sequences > 1, logging a warning the first time per backend instance.
+
+        Args:
+            mot: The output thunk whose generation metadata will be updated.
+            hf_output: The HuggingFace generation output containing scores/logits.
+        """
+        hf_field = {ModelOption.LOGITS: "scores", ModelOption.RAW_LOGITS: "logits"}
+        for opt, tensors, attr in (
+            (ModelOption.LOGITS, hf_output.scores, "logits"),
+            (ModelOption.RAW_LOGITS, hf_output.logits, "raw_logits"),
+        ):
+            if not (mot._model_options and mot._model_options.get(opt)):
+                continue
+            if tensors is None:
+                MelleaLogger.get_logger().debug(
+                    "%s requested but hf_output.%s is None; generation.%s will not be populated.",
+                    opt,
+                    hf_field[opt],
+                    attr,
+                )
+                continue
+            warn_key = f"nrs_{opt}"
+            if tensors[0].shape[0] > 1:
+                if warn_key not in self._warned_about:
+                    self._warned_about.add(warn_key)
+                    MelleaLogger.get_logger().warning(
+                        "%s is set but num_return_sequences > 1; "
+                        "logit tensors are ambiguous across sequences and will not be populated.",
+                        opt,
+                    )
+            else:
+                setattr(
+                    mot.generation,
+                    attr,
+                    tuple(s.squeeze(0).detach().clone() for s in tensors),
+                )
+
     async def post_processing(
         self,
         mot: ModelOutputThunk,
@@ -1262,7 +1364,11 @@ class used during generation, if any.
         if (
             self._use_caches
             and isinstance(hf_output, GenerateDecoderOnlyOutput)
-            and (hf_output.past_key_values is not None or hf_output.scores is not None)
+            and (
+                hf_output.past_key_values is not None
+                or hf_output.scores is not None
+                or hf_output.logits is not None
+            )
         ):
             output_complete = hf_output.sequences[0]
             kv_cache: DynamicCache | None = hf_output.past_key_values  # type: ignore
@@ -1282,29 +1388,19 @@ class used during generation, if any.
             cache_key = id(mot.value)
             self.cache_put(cache_key, cache_info)
 
-            # Surface logits to caller before clearing — scores move to LRU cache;
-            # when LOGITS=True the tuple views also keep them alive via generation.logits.
-            if (
-                hf_output.scores is not None
-                and mot._model_options
-                and mot._model_options.get(ModelOption.LOGITS)
-            ):
-                # squeeze(0): hf_output.scores is (1, vocab_size) per token; normalise to (vocab_size,)
-                mot.generation.logits = tuple(s.squeeze(0) for s in hf_output.scores)
+            # Surface logits before clearing — scores move to LRU cache; when
+            # LOGITS/RAW_LOGITS=True the tuple views also keep them alive via generation.
+            if mot._model_options:
+                self._surface_logits(mot, hf_output)
 
-            # Clear KV cache and scores from HF output - they're retained via the LRU cache
-            # (and, when LOGITS=True, also via the views in mot.generation.logits).
+            # Clear KV cache and scores from HF output; retained via LRU cache above.
             hf_output.past_key_values = None
             hf_output.scores = None
-        elif (
-            isinstance(hf_output, GenerateDecoderOnlyOutput)
-            and hf_output.scores is not None
-            and mot._model_options
-            and mot._model_options.get(ModelOption.LOGITS)
-        ):
-            # Caching disabled — scores were not moved to LRU, surface them directly
-            # squeeze(0): normalise (1, vocab_size) per token to (vocab_size,)
-            mot.generation.logits = tuple(s.squeeze(0) for s in hf_output.scores)
+            hf_output.logits = None
+        elif isinstance(hf_output, GenerateDecoderOnlyOutput) and mot._model_options:
+            # Caching disabled — surface scores/raw logits directly if requested.
+            self._surface_logits(mot, hf_output)
+            hf_output.logits = None
 
         # Only scan for tools if we are not doing structured output and tool calls were provided to the model.
         if _format is None and tool_calls:
@@ -1526,6 +1622,12 @@ async def generate_from_raw(
             # transformers' generate() requires a tokenizer to decode stop_strings.
             generate_kwargs["tokenizer"] = self._tokenizer
 
+        raw_scores_kwargs: dict[str, Any] = {}
+        if model_opts.get(ModelOption.LOGITS):
+            raw_scores_kwargs["output_scores"] = True
+        if model_opts.get(ModelOption.RAW_LOGITS):
+            raw_scores_kwargs["output_logits"] = True
+
         _start = time.perf_counter()
         try:
             outputs = await asyncio.to_thread(
@@ -1536,7 +1638,7 @@ async def generate_from_raw(
                 input_ids=inputs["input_ids"],
                 attention_mask=inputs["attention_mask"],
                 return_dict_in_generate=True,
-                output_scores=True,
+                **raw_scores_kwargs,
                 **generate_kwargs,
                 **format_kwargs,
             )
@@ -1571,6 +1673,7 @@ async def generate_from_raw(
         agg_prompt = 0
         agg_completion = 0
         want_logits = bool(model_opts and model_opts.get(ModelOption.LOGITS))
+        want_raw_logits = bool(model_opts and model_opts.get(ModelOption.RAW_LOGITS))
         for i, decoded_result in enumerate(decoded_results):
             n_prompt_tokens = inputs["input_ids"][i].size(0)  # type: ignore
             n_completion_tokens = len(sequences_to_decode[i])
@@ -1593,6 +1696,11 @@ async def generate_from_raw(
                     step_scores[i].detach().clone() for step_scores in outputs.scores
                 )
 
+            if want_raw_logits and outputs.logits is not None:
+                result.generation.raw_logits = tuple(
+                    step_logits[i].detach().clone() for step_logits in outputs.logits
+                )
+
             action = actions[i]
             result.parsed_repr = (
                 action.parse(result) if isinstance(action, Component) else result.value

mellea/backends/litellm.py

@@ -251,10 +251,15 @@ def _make_backend_specific_and_remove(
             # OpenAI compatible endpoints should accept both (and Watsonx does accept both).
             model_opts_remapping[ModelOption.MAX_NEW_TOKENS] = "max_tokens"
 
-        if model_options.get(ModelOption.LOGITS):
-            MelleaLogger.get_logger().warning(
-                "ModelOption.LOGITS is not supported by the LiteLLM backend; generation.logits will be None."
-            )
+        for opt, field in (
+            (ModelOption.LOGITS, "generation.logits"),
+            (ModelOption.RAW_LOGITS, "generation.raw_logits"),
+        ):
+            if model_options.get(opt) and opt not in self._warned_about:
+                self._warned_about.add(opt)
+                MelleaLogger.get_logger().warning(
+                    f"{opt!r} is not supported by the LiteLLM backend; {field} will be None."
+                )
 
         backend_specific = ModelOption.replace_keys(model_options, model_opts_remapping)
         backend_specific = ModelOption.remove_special_keys(backend_specific)

mellea/backends/model_options.py

@@ -32,7 +32,8 @@ class ModelOption:
         STREAM (str): Sentinel key for enabling streaming responses.
         STOP_SEQUENCES (str): Sentinel key for a `list[str]` of strings that, when
             encountered in the model output, cause generation to halt.
-        LOGITS (str): Sentinel key for requesting per-token logit scores from the backend.
+        LOGITS (str): Sentinel key for requesting per-token processed logit scores (post-LogitsProcessor).
+        RAW_LOGITS (str): Sentinel key for requesting per-token raw LM-head logits (pre-LogitsProcessor).
     """
 
     TOOLS = "@@@tools@@@"
@@ -71,19 +72,38 @@ class ModelOption:
     text-generation endpoint).
     """
     LOGITS = "@@@logits@@@"
-    """When `True`, request that the backend return per-token logit scores.
+    """When `True`, request per-token processed logit scores (``output_scores``).
 
-    Scores are exposed on `mot.generation.logits` as a tuple of 1-D
-    tensors of shape `(vocab_size,)`, one per generated token. This shape
-    is consistent across both the standard and batch (`generate_from_raw`)
-    HuggingFace paths.
+    These are logits *after* the ``LogitsProcessor`` chain has run —
+    temperature scaling, top-k/top-p masking, repetition penalty, etc.
+    Exposed on ``mot.generation.logits`` as a tuple of 1-D tensors of shape
+    ``(vocab_size,)``, one per generated token.
 
     Only supported by the HuggingFace local backend. Backends that cannot
     return logits (OpenAI, Ollama, LiteLLM, WatsonX) log a warning when this
-    option is set and leave `generation.logits` as `None`.
+    option is set and leave ``generation.logits`` as ``None``.
 
-    **Streaming not supported**: when `ModelOption.STREAM=True`, logit
-    scores are not available and `mot.generation.logits` will be `None`.
+    **Streaming not supported**: when ``ModelOption.STREAM=True``, logit
+    scores are not available and ``mot.generation.logits`` will be ``None``.
+
+    See also ``ModelOption.RAW_LOGITS`` for unprocessed LM-head output.
+    """
+
+    RAW_LOGITS = "@@@raw_logits@@@"
+    """When `True`, request per-token raw logits (``output_logits``).
+
+    These are the raw, unprocessed logits straight from the LM head —
+    the model's actual output *before* any ``LogitsProcessor`` transforms
+    (temperature, top-k/top-p, repetition penalty, etc.).
+    Exposed on ``mot.generation.raw_logits`` as a tuple of 1-D tensors of
+    shape ``(vocab_size,)``, one per generated token.
+
+    Only supported by the HuggingFace local backend.
+
+    **Streaming not supported**: when ``ModelOption.STREAM=True``, raw
+    logits are not available and ``mot.generation.raw_logits`` will be ``None``.
+
+    See also ``ModelOption.LOGITS`` for processor-transformed scores.
     """
 
     @staticmethod

mellea/backends/ollama.py

@@ -286,10 +286,15 @@ def _make_backend_specific_and_remove(
         Returns:
             a new dict
         """
-        if model_options.get(ModelOption.LOGITS):
-            MelleaLogger.get_logger().warning(
-                "ModelOption.LOGITS is not supported by the Ollama backend; generation.logits will be None."
-            )
+        for opt, field in (
+            (ModelOption.LOGITS, "generation.logits"),
+            (ModelOption.RAW_LOGITS, "generation.raw_logits"),
+        ):
+            if model_options.get(opt) and opt not in self._warned_about:
+                self._warned_about.add(opt)
+                MelleaLogger.get_logger().warning(
+                    f"{opt!r} is not supported by the Ollama backend; {field} will be None."
+                )
 
         backend_specific = ModelOption.replace_keys(
             model_options, self.from_mellea_model_opts_map

mellea/backends/openai.py

@@ -425,10 +425,15 @@ def _make_backend_specific_and_remove(
 
         backend_specific = ModelOption.replace_keys(model_options, remap_dict)
 
-        if model_options.get(ModelOption.LOGITS):
-            MelleaLogger.get_logger().warning(
-                "ModelOption.LOGITS is not supported by the OpenAI backend; generation.logits will be None."
-            )
+        for opt, field in (
+            (ModelOption.LOGITS, "generation.logits"),
+            (ModelOption.RAW_LOGITS, "generation.raw_logits"),
+        ):
+            if model_options.get(opt) and opt not in self._warned_about:
+                self._warned_about.add(opt)
+                MelleaLogger.get_logger().warning(
+                    f"{opt!r} is not supported by the OpenAI backend; {field} will be None."
+                )
 
         # OpenAI Backend has specific filtering functionality.
         if is_chat_context: