Fix truncation not keeping the last message per the documented rule

[Osamaali313]

Problem

_truncate_for_max_tokens (tokens/tokenizers/instruct.py) documents three rules for fitting a conversation into max_tokens, including that the last message must be kept:

# drop some messages to fit in max_tokens. Rules:
# - don't drop any system messages
# - when a user message is dropped, all following assistant|tool message should be dropped until the next
#   user message
# - we never drop the last message

But the inner drop() only guards two cases — system messages and the last user message (last_user_message_index). The last message itself is never guarded. When the final message is an assistant or trailing tool message — finetuning examples always end with the assistant target — the "clear everything until the next user message" cascade can silently exclude it. The result is a truncated example that ends at a user message with no assistant response, plus a Tokenized.prefix_ids that references tokens no longer present in tokens.

This is reachable from the public API: MistralTokenizer.encode_chat_completion(request, max_model_input_len=N) with request.truncate_for_context_length=True.

Fix

Add the missing guard so drop() honors the documented rule:

if idx == len(messages) - 1:
    # never drop the last message
    return

When the last message genuinely cannot fit, to_drop stays positive and the method raises TokenizerException("Input couldn't fit in truncate_at_max_token") — the documented loud failure — instead of returning a corrupted result.

Reproduction

conversation (last msg = final assistant target), tight budget	before	after
[user, FINAL assistant]	last message excluded	raises "couldn't fit"
[user, asst, user, FINAL asst]	last message excluded	raises "couldn't fit"
normal case (a non-final message can be removed)	correct	correct (unchanged)

Existing test_truncation cases all retain the last message, so they are unaffected; the all-or-nothing case now matches test_truncation_failed's raise.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.