Skip to content

fix: retain end-of-generation token logprobs - #2383

Open
Kumar-Abhimanyu wants to merge 2 commits into
abetlen:mainfrom
Kumar-Abhimanyu:fix/eos-logprobs
Open

Kumar-Abhimanyu wants to merge 2 commits into
abetlen:mainfrom
Kumar-Abhimanyu:fix/eos-logprobs

Conversation

@Kumar-Abhimanyu

Copy link
Copy Markdown

When the sampler immediately produces an end-of-generation token, completions currently return empty probability arrays even with logprobs enabled. The generation loop exits before retaining that token's probability.

This change preserves the sampled end-of-generation token in probability metadata. Regular completions append its special-token spelling and probability; streamed completions carry it in the final chunk. Completion text, token usage, and cache keys retain their existing behavior. Requests without logprobs remain unchanged.

Fixes #1528.

Validation:

  • Built the native CPU library on Linux with Python 3.12, GCC, GGML_NATIVE=OFF and LLAMA_CURL=OFF.
  • 24 regression cases cover immediate EOS and EOS after text, streaming and regular responses, echo, and logprobs=None/0/3. Sampling is controlled; tokenization and EOG detection use the real native vocabulary model.
  • On the original code, 16 probability cases fail and 8 cases without logprobs pass. All 24 pass with the fix.
  • Existing tokenizer/version tests and grammar/chat-format/speculative tests pass: 32 tests total.
  • Repository-wide Ruff lint and formatting checks pass.

Full model-download/inference tests and GPU backends were not run. The proposed behavior includes special tokens only in probability metadata; maintainer feedback on the EOS probability contract is welcome.

This branch has not been deployed

No deployments
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.

logprobs aren't returned if sampler will reach EOS

1 participant