One thing that still bothers me is that:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3-4B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

messages = [{"role": "user", "content": "Give me a concise proof that sqrt(2) is irrational."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)

think_start = tokenizer.convert_tokens_to_ids("<think>")
think_end = tokenizer.convert_tokens_to_ids("</think>")
model.generation_config.begin_thinking_token_id = think_start
model.generation_config.end_thinking_token_id = think_end

out = model.generate(
    **inputs,
    max_new_tokens=128,
    max_thinking_tokens=64,
    do_sample=False,
)

print(tokenizer.decode(out[0], skip_special_tokens=False))

will result in:

<|im_start|>user
Give me a concise proof that sqrt(2) is irrational.<|im_end|>
<|im_start|>assistant
<think>
Okay, so I need to prove that the square root of 2 is irrational. Hmm, I remember that an irrational number is a number that can't be expressed as a fraction of two integers. So, I need to show that there are no integers a and b where a/b equals sqrt(2). Let</think>

To prove that √2 is irrational, we use a **proof by contradiction**. Here's a concise version:

1. **Assume** √2 is rational. Then it can be written as a fraction $ \frac{a}{b} $, where $ a $ and $ b

where the end_thinking_token is appended to the output, whereas by default \n<\think> is generated by the model, resulting in:

<|im_start|>user
Give me a concise proof that sqrt(2) is irrational.<|im_end|>
<|im_start|>assistant
<think>
Okay, so I need to prove that the square root of 2 is irrational. Hmm, I remember that an irrational number is a number that can't be expressed as a fraction of two integers. So, I need to show that there are no integers a and b where a/b equals sqrt(2). Let
</think>

To prove that √2 is irrational, we use a **proof by contradiction**. Here's a concise version:

1. **Assume** √2 is rational. Then it can be written as a fraction $ \frac{a}{b} $, where $ a $ and $ b

Additionally, the early stopping of the Qwen example (https://qwen.readthedocs.io/en/latest/getting_started/quickstart.html#thinking-budget) adds:
\n\nConsidering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>\n\n

Not sure whether such customization should be allowed? Kind of related to the comment above, but seems like a different logic and maybe for a future PR?

I might have overcomplicated the MaxThinkingTokensLogitsProcessor, but it is to avoid token overcount when 1) people use <think> in their user prompt, or 2) people feed previous reasoning traces which include <think>...</think> (although this is not best practice, see number 4 in https://huggingface.co/Qwen/Qwen3-4B#best-practices).

Another issue is when enable_thinking=True pre-fills the prompt with a <think> token as this becomes part of the prompt_length and will be ignored. Any thoughts on this? A simple solution here would be prompt_length-1, but I noticed that, for example, Qwen3-4B doesn't do the pre-filling.

While testing this, another problem is that the model might output a second <\think> token after \n\n (so not directly after the enforced <\think> token):

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3-4B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

messages = [
    {
        "role": "user",
        "content": "Give me a concise proof that sqrt(2) is irrational.",
    },
]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)

think_start = tokenizer.convert_tokens_to_ids("<think>")
think_end = tokenizer.convert_tokens_to_ids("</think>")
model.generation_config.begin_thinking_token_id = think_start
model.generation_config.end_thinking_token_id = think_end

out = model.generate(
    **inputs,
    max_new_tokens=64,
    max_thinking_tokens=32,
    do_sample=False,
)

print(tokenizer.decode(out[0], skip_special_tokens=False))

<|im_start|>user
Give me a concise proof that sqrt(2) is irrational.<|im_end|>
<|im_start|>assistant
<think>
Okay, so I need to prove that the square root of 2 is irrational. Hmm, I remember that an irrational number is a number that can't</think>

</think>

Sure! Here's a concise proof that $\sqrt{2}$ is irrational:

---

**Proof that $\sqrt{2}$ is irrational

Update should satisfy following requirements:

• When a user uses past reasoning traces (i.e., previous assistant chunks contain <think> ... </think>), it should ignore this.
• When a user uses <think>, </think>, or <think> ... </think>, in their prompt, it should be ignored.
• It should count the number of tokens only in the current reasoning chain.
  It should be robust against the model generating additional <think> tokens in the current reasoning chain (not restart the count).
• It should be robust against the model automatically appending the <think> token when enable_thinking=True (right now it becomes part of the prompt and gets ignored because of it).
• The model sometimes generates an additional </think> token after the enforced closure </think> token. Not always directly after, sometimes it first generated \n\n and then the </think> token. Right now, the </think> token can only be generated if a new <think> token has been generated.

@AndresAlgaba hey!
Great job, thanks a lot for the PR! I think it a super useful feature to have. I left a few general comment and question above, mostly to make sure I understand the flow and to simplify it when possible

I'm actually not sure about this PR! I think forcibly ending model thinking mid-stream will force them out of distribution and cause them to output very low-quality answers. We could add it as a feature for advanced users, but I worry that if we do that then people will assume that because it's a feature that it works well in practice, when I think it probably won't, right?

Note that in the Qwen example, they don't cut off thinking directly, but instead insert an "early stopping prompt" that is likely fine-tuned for Qwen and will not work well on other thinking models.

@zucchini-nlp Thank you for the feedback, working on it!

@Rocketknight1 Thanks for the comment! Indeed, I agree it is more of an experimental feature. It is in line with Gemini's thinking_budget (https://ai.google.dev/gemini-api/docs/thinking#set-budget- and Claude's budget_tokens (https://docs.claude.com/en/docs/build-with-claude/extended-thinking#how-to-use-extended-thinking). I remember that DeepSeek was planning to implement a max_thinking_tokens parameter at some point, but I can't find it anymore. I implemented it for my own research and thought others might be interested in this feature, but with the warranty that it is experimental. So not sure whether it should then be added to the library like this?

hi @AndresAlgaba, yes, but note that that feature in Gemini or Claude does not actually cap the reasoning tokens or forcibly stop thinking - they specifically say the model can overflow the budget. They've probably just trained the model so that it (roughly) obeys thinking budgets, and then they probably tell it the budget in the prompt somewhere

True, same for the reasoning effort in OpenAI models. I still believe some people might be interested in the feature, see a similar feature request in vllm: vllm-project/vllm#17887 (and PR: vllm-project/vllm#20859, where they also propose a logit processor). Is there any way we can indicate it's experimental? Is it an option to keep the processor without adding themax_thinking_tokens argument for now?

I believe that adding the \n\nConsidering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>\n\n part is interesting, but part of different PR?

Hmn, if we can mention in the docs that it's experimental and might reduce answer quality it's probably okay! cc @zucchini-nlp, are you happy with that?

Thanks again for the feedback! I have made some changes after further experimentation and have updated the docs (including a warning that the feature is experimental).

Seems like a good idea! I agree that hard truncation is often not the way to go. Thank you for offering help, I will look into it and keep you posted here!

Hi @zucchini-nlp, sorry for not following up on this but I decided to look further into the "early stopping" schemes and write a paper on it (now available on arxiv: https://arxiv.org/abs/2601.23163). I will experiment further and then start to add to the transformers-community org!

Amazing @AndresAlgaba ! If you have a paper/blog-post, feel free to link it directly in the hub page via hub papers :)