Skip to content

fix: return raw model logits from hooked generation - #1967

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/generation-raw-logits
Nov 2, 2026
Merged

jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/generation-raw-logits

Conversation

@emerardd

@emerardd emerardd commented Nov 2, 2026

Copy link
Copy Markdown
Contributor

Description

Fixes #1966.

Keep the raw last-position model logits separate from decoding scores inside _generate_tokens(), and return the raw tensor to generate()'s logits collector. This fixes output_logits=True returning forced-EOS, minimum-length, or no-repeat-ngram scores in .logits rather than the unprocessed model output.

Sampling and stopping criteria continue to receive the existing processed scores; their ordering, token selection, cache handling, streaming output, and the public return structure are unchanged. No output_scores API or dependency is added. The generation docstring now explicitly describes the returned logits as unprocessed.

The offline integration tests compare every returned step against an independent same-weight HF model's forward on that generated prefix. They cover all three constraint routes plus the unconstrained control, greedy/sampled decoding, single/batched input, unchanged sequences with output_logits enabled, and cached/uncached GPT-2 with temperature, top-k, frequency and repetition penalties. Constraint assertions verify that forced EOS, EOS suppression, and no-repeat processing still affect generated tokens.

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Validation

  • Pre-fix single-row matrix: 6 failed, 6 passed. All three processor routes fail in both greedy and sampled decoding; unconstrained BART and decoder-only controls pass.
  • Final affected surface: 72 passed across test_generation_raw_logits.py, test_seq2seq_generate_stream.py, test_generate_stream_args.py, and test_generate_stream_batch_decode.py.
  • Full mypy: 390 source files passed.
  • Scoped pinned pycln/isort/Black checks and git diff --check passed.
  • All models are tiny config-built CPU models; no downloads. The full unit/test-pr suite, pretrained-model parity, GPU and MPS were not run. Output includes the existing SWIG import/shutdown deprecation warnings; no new warning category was introduced.

Checklist

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes — only the affected surface above was run, not the full suite
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4

jlarson4 commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator

Looks good, merging as is!

@jlarson4
jlarson4 merged commit 3ce1fd4 into TransformerLensOrg:dev Nov 2, 2026
29 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug Report] generate(output_logits=True) returns decoding scores instead of raw model logits

2 participants