feat: standardize logits in mot by AngeloDanducci · Pull Request #1261 · generative-computing/mellea

AngeloDanducci · 2026-06-12T06:13:51Z

Pull Request

Issue

Fixes #123

Description

standardize logits in MOT

Testing

Tests added to the respective file if code was changed
New code has 100% coverage if code was added
Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

AI coding assistants used

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

Component
Requirement
Sampling Strategy
Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

psschwei · 2026-06-12T11:50:27Z

cc @ajbozarth as it relates to your MOT redesign work

ajbozarth · 2026-06-12T15:02:45Z

While reviewing I realized that when I reworked #909 into an epic it called out an intent to put logits into generation when implementing #123 but didn't add a comment on that issue to do so.

when implemented, logits go in mot.generation (e.g. generation.logprobs), not mot.raw. Per #793 precedent.

ajbozarth

a couple follow up nits to my above comment

ajbozarth

Some feedback from Claude:

Main ask is to land the field on GenerationMetadata rather than as a top-level attr on ModelOutputThunk — per #909 and the #793 precedent, that's the standardized home for backend-execution metadata. The implementation logic (squeeze, clone-per-batch-item, cached vs. non-cached branches) looks correct; the rest of my comments are smaller items.

ajbozarth

LGTM, just a couple nits

akihikokuroda

LGTM

ajbozarth

Some feedback from Claude.

ajbozarth

LGTM, I confirmed all of my comments were addressed in addition to @jakelorocco comments

jakelorocco · 2026-06-23T13:23:12Z

+        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
+
+            # Two load-bearing assumptions: (1) __getattr__ falls through for attributes
+            # accessed by generate_with_transformers (model.device, model.vocab_size,
+            # model.generation_config). (2) chat_completion_request_to_transformers_inputs
+            # always sets return_dict_in_generate=True, so .generate() always returns a
+            # GenerateDecoderOnlyOutput — if that ever changes, the cell stays None and
+            # logits silently won't be populated.
+            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]


Can you please confirm that you manually ran the intrinsic tests (minus the openai ones) to make sure this doesn't cause issues? If it does, I think we can just say it's not supported behavior for this code path. I didn't realize that supporting this for adapter components would cause this complexity.

Yes I had manually run the intrinsics tests, and the only failures were also the same ones I had on main.

I just double checked main against my 4 previously troubled tests, which now have 3 passing and 1 xfail on main.

test_intrinsics_formatters.py::test_canned_input[context_relevance] test_intrinsics_formatters.py::test_run_transformers[context_relevance_alora] test_intrinsics_formatters.py::test_run_transformers[context-attribution] test_intrinsics_formatters.py::test_run_transformers[uncertainty_alora]

After rebasing these all seem to mirror main now (all passing, 1 xfail).

I'll run against this again as well just to double check

uv run pytest test/formatters/granite/test_intrinsics_formatters.py test/formatters/granite/test_intrinsics_canned_output.py test/backends/test_adapters/ test/stdlib/components/intrinsic/ -m "not openai and not qualitative"

The above all pass, will run a full suite again - would not be surprised if the rebase has introduced a few other issues.

jakelorocco · 2026-06-23T13:40:08Z

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


I looked at this with fresh eyes. I don't think we want to cache things here if neither the kv_cache nor the scores are not None.

Would it be possible to separate out the caching logic from the logit surfacing logic? I know the cache handles deleting references from the hf_output object, but maybe this can be extracted as well?

I think so, will take a stab at that.

jakelorocco · 2026-06-23T13:45:31Z

Does logits need to be added here as well? To be explicitly deleted? I don't think so but would like to confirm?

If my understanding is correct it's not /necessary/ but I should probably add it anyway - after del mot.eta["hf_output"] garbage collection should clean it up (but I think this is also the case for the scores and sequences?).

jakelorocco

small issue with the rebase I think

jakelorocco · 2026-06-24T15:34:00Z

-        allowed_types = tuple(at.value for at in action.adapter_types)
-        adapter = self._find_adapter(action.intrinsic_name, allowed_types)
+        adapter = get_adapter_for_intrinsic(
+            action.intrinsic_name, action.adapter_types, self._added_adapters
+        )


I think these lines were rebased incorrectly? main shows the allowed_types... version.

Updated, re-ran intrinsics test locally 👍

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

Co-authored-by: Alex Bozarth <ajbozart@us.ibm.com> Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>

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

AngeloDanducci requested review from a team, jakelorocco and nrfulton as code owners June 12, 2026 06:13

AngeloDanducci requested a review from akihikokuroda June 12, 2026 06:13

github-actions Bot added the enhancement New feature or request label Jun 12, 2026

AngeloDanducci enabled auto-merge June 12, 2026 06:16

ajbozarth reviewed Jun 12, 2026

View reviewed changes

Comment thread mellea/core/base.py Outdated

Comment thread mellea/core/base.py Outdated

ajbozarth reviewed Jun 12, 2026

View reviewed changes

Comment thread mellea/core/base.py Outdated

Comment thread mellea/backends/huggingface.py Outdated

Comment thread mellea/backends/model_options.py

Comment thread test/backends/test_huggingface_unit.py

Comment thread test/core/test_logits.py Outdated

AngeloDanducci requested a review from ajbozarth June 15, 2026 18:10

ajbozarth reviewed Jun 15, 2026

View reviewed changes

Comment thread mellea/backends/model_options.py

Comment thread mellea/backends/model_options.py Outdated

Comment thread mellea/core/base.py Outdated

Comment thread mellea/backends/huggingface.py Outdated

ajbozarth approved these changes Jun 15, 2026

View reviewed changes

Comment thread mellea/backends/huggingface.py Outdated

Comment thread test/backends/test_huggingface_unit.py Outdated

AngeloDanducci force-pushed the ad-123 branch from 894489e to 94a3d61 Compare June 15, 2026 19:46

akihikokuroda approved these changes Jun 16, 2026

View reviewed changes

ajbozarth mentioned this pull request Jun 19, 2026

refactor(core): consolidate backend raw responses onto mot.raw namespace #1215

Closed

6 tasks

jakelorocco reviewed Jun 22, 2026

View reviewed changes

Comment thread mellea/backends/huggingface.py Outdated

Comment thread mellea/backends/huggingface.py

Comment thread mellea/backends/litellm.py Outdated

Comment thread mellea/backends/huggingface.py

AngeloDanducci force-pushed the ad-123 branch from 6f2668d to 91a801a Compare June 22, 2026 15:20

ajbozarth requested changes Jun 22, 2026

View reviewed changes

Comment thread mellea/backends/huggingface.py Outdated

Comment thread mellea/backends/huggingface.py

Comment thread mellea/backends/huggingface.py Outdated

ajbozarth approved these changes Jun 22, 2026

View reviewed changes

AngeloDanducci requested a review from jakelorocco June 22, 2026 18:01

jakelorocco reviewed Jun 23, 2026

View reviewed changes

AngeloDanducci force-pushed the ad-123 branch from 7173b34 to 04e358d Compare June 23, 2026 19:55

AngeloDanducci requested a review from jakelorocco June 23, 2026 22:03

jakelorocco requested changes Jun 24, 2026

View reviewed changes

jakelorocco approved these changes Jun 24, 2026

View reviewed changes

AngeloDanducci added this pull request to the merge queue Jun 24, 2026

github-merge-queue Bot removed this pull request from the merge queue due to a conflict with the base branch Jun 24, 2026

AngeloDanducci added 2 commits June 24, 2026 15:00

standardize logits in mot

3176fe2

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

address review feedback

56cbf18

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

AngeloDanducci and others added 12 commits June 24, 2026 15:00

adds warnings and removes double backticks

942e5be

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

Update mellea/backends/huggingface.py

1bcda1a

Co-authored-by: Alex Bozarth <ajbozart@us.ibm.com> Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>

Update test/backends/test_huggingface_unit.py

f43f8b6

Co-authored-by: Alex Bozarth <ajbozart@us.ibm.com> Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>

review updates

5110e8a

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

rebase cleanup and missing imports

a1adefa

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

additional review feedback

71831f1

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

change logit surfacing and cleanup

2febf40

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

fix logit surfacing

c3c07a4

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

update test model ids suite

992c3b4

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

fix rebase issue

58a14d0

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

fix adapter test

e95c06c

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

rebase and update tests to match

d3d55ac

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

AngeloDanducci force-pushed the ad-123 branch from b84bbbb to d3d55ac Compare June 24, 2026 19:12

AngeloDanducci enabled auto-merge June 24, 2026 19:12

AngeloDanducci added this pull request to the merge queue Jun 24, 2026

Merged via the queue into generative-computing:main with commit 8c951a8 Jun 24, 2026
9 checks passed

AngeloDanducci deleted the ad-123 branch June 24, 2026 20:05

Uh oh!

Conversation

AngeloDanducci commented Jun 12, 2026

Pull Request

Issue

Description

Testing

Attribution

Adding a new component, requirement, sampling strategy, or tool?

Uh oh!

psschwei commented Jun 12, 2026

Uh oh!

ajbozarth commented Jun 12, 2026

Uh oh!

ajbozarth left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ajbozarth left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ajbozarth left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

akihikokuroda left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ajbozarth left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ajbozarth left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jakelorocco left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development