BUG: Stop sequences not working properly in TransformersModel

[MaxBourdon]

Problem
Hi!
I have discovered that stop_sequences used in TransformersModel.generate do not always work as intended.
In particular, when running CodeAgent with default stop_sequences = ["Observation:", "Calling tools:"] I still see (sometimes multiple) "Calling tools:" strings in assistant's messages.

Steps to reproduce
Here is an example to reproduce the bug.

from smolagents import CodeAgent, TransformersModel


model = TransformersModel("Qwen/Qwen2.5-Coder-7B-Instruct")

agent = CodeAgent(
    tools=[],
    model=model,
)

# Ensure multi-turn code generation so that the agent sees "Calling tools:" in the message history
agent.run("Concatenate strings 'foo' and 'bar'")
agent.run("Concatenate strings 'foo' and 'bar' again", reset=False)

# Print only assistant's messages (generated by the model)
for m in agent.write_memory_to_messages():
    if m.role == "assistant":
        print(m.content[0]["text"])
        print(100 * '-')

Actual behavior
Here is what printed:

Thought: To concatenate the strings 'foo' and 'bar', I will use Python's string concatenation operator `+`.
<code>
result = 'foo' + 'bar'
final_answer(result)
</code>
----------------------------------------------------------------------------------------------------
Thought: Since we already successfully concatenated 'foo' and 'bar' in the previous task, I will repeat the same process to ensure consistency.
<code>
result = 'foo' + 'bar'
final_answer(result)
</code>
Calling tools:
[{'id': 'call_2', 'type': 'function', 'function': {'name': 'python_interpreter', 'arguments': "result = 'foo' + 'bar'\nfinal_answer(result)"}}]</code>
----------------------------------------------------------------------------------------------------

As can be seen, the agent adds "Calling tools:" to the last generated message, which should not be there.

Environment:
I am using the latest (as of now) released version smolagents[toolkit,transformers]==1.21.2.
Other environment details are not relevant here.

Reason:
The reason for such behavior lies in the current implementation of make_stopping_criteria.

To be more precise, StopOnStrings.__call__ checks for the presence of stop strings at the very end of the generation stream, which is appended with newly generated tokens. However, since tokenizer decodes tokens in chunks of chars, a stop string can be added to the stream with some extra trailing chars.
For instance, a "Calling tools:\n" string would be tokenized as ["Calling", " tools", ":\n"], so it is added to the stream with trailing '\n'; therefore, exact string matching via endswith fails to detect the stop string "Calling tools:".

Fix:
I propose the following version of StopOnStrings, which solves this problem:

class StopOnStrings(StoppingCriteria):
    def __init__(self, stop_strings: list[str], tokenizer):
        self.stop_strings = stop_strings
        self.tokenizer = tokenizer
        self.prev_tail = ""
        if any(s == "" for s in self.stop_strings):
            raise ValueError("stop_strings must be non-empty")
        # Keep only last (max_len - 1) chars
        self.max_tail_len = max(len(s) for s in self.stop_strings) - 1

    def reset(self):
        self.prev_tail = ""

    def __call__(self, input_ids, scores, **kwargs):
        """
        Return True if any stop string appears in the generated tokens stream; otherwise return False.
        """
        generated = self.tokenizer.decode(
            input_ids[0][-1], skip_special_tokens=True
        )

        if not generated:
            return False

        # Build the minimal search window that contains every possible
        # newly created occurrence of any stop string
        window = self.prev_tail + generated

        # Update the tail
        self.prev_tail = window[-self.max_tail_len :]

        # Check each stop string -- we only check inside window
        for s in self.stop_strings:
            if s in window:
                return True
        return False

For similar reasons, I suggest fixing remove_stop_sequences as well:

def remove_stop_sequences(content: str, stop_sequences: list[str], delta: int = 128) -> str:
    """
    Removes stop sequences if they occur at the end of the content string, but accounts for extra trailing characters due to tokenization artefacts.
    `delta` sets the allowed right offset for extra characters.
    """
    for stop_seq in stop_sequences:
        subseq_len = len(stop_seq) + delta
        if stop_seq in content[-subseq_len:]:
            idx = content.rindex(stop_seq)
            content = content[:idx]
    return content

In addition, I would like to note that after adding these simple fixes to my project, I noticed a significant reduction in the average step execution time of my agent and even some improvements in quality. It seems that generating these fake “Calling tools:” lines actually wastes a lot of computing resources.

I would be glad to submit a pull request that resolves this issue.

---

Checklist

• I have searched the existing issues and have not found a similar bug report.
• I have provided a minimal, reproducible example.
• I have provided the full traceback of the error.
• I have provided my environment details.
• I am willing to work on this issue and submit a pull request.

[albertvillanova]

Thanks for reporting, @MaxBourdon.

I am wondering why not just using the builtin transformers StopStringCriteria: https://huggingface.co/docs/transformers/v4.53.3/en/internal/generation_utils#transformers.StopStringCriteria

CC: @aymeric-roucher who implemented current code:

• c4f3885#diff-0319e64c7d0f9e302a7c7a0dd60f04dedb7d4dcf7dbe9e046543335654a9efa1

[MaxBourdon]

Thanks for your reply, @albertvillanova

I am wondering why not just using the builtin transformers StopStringCriteria: https://huggingface.co/docs/transformers/v4.53.3/en/internal/generation_utils#transformers.StopStringCriteria

At first, I also tried using the default StopStringCriteria but then found that it drastically increases agent's step execution time.
For some reason, the string-level implementation works much faster in practice.

Here is the comparison of execution time for the default generation, generation with stop_strings, and generation with stopping_criteria based on the example from the docs. For stopping_criteria, I use my implementation of StopOnStrings from the original post.

[albertvillanova]

CC @huggingface/transformers, in case the StopStringCriteria could be improved or we need to adjust something in our integration?

[ydshieh]

Probably @gante would have some insights

[gante]

@MaxBourdon I believe the root cause was in the custom StoppingCriteria class defined in smolagent, and transformers has a better native solution 🤗

Check my PR linked above for an explanation -- it should also fix the issue!

[MaxBourdon]

@gante Thank you for your PR.
Please note that, as discussed above, the native transformers stop_strings solution works much slower than (seemingly naive) string-level comparison. I use the latter in my project, it noticeably reduces per-step time in practice.

[gante]

@MaxBourdon ah, I missed the discussion halfway through this issue (speed issues).

Your solution has the limitation that it can't be compiled (because it calls the tokenizer). This trait is not important for smolagents, so it can be used as a replacement here, yes. It is also mostly correct, although I would drop skip_special_tokens=True to be able to stop on special tokens as well. For full correctness, it would need to handle batch size > 1.

I'm going to have a look at the speed issue for the transformers StopStringCriteria class. If I can't get speedups, I'll edit my PR to smolagents to use your class instead -- and maybe add it in transformers too 🤗

[gante]

@MaxBourdon benchmarked 🤗

It is indeed slower due to the torch.compile requirements, especially on the first run -- it relies on cached data at initialization time, so the times you see using timeit correspond to the worst-case scenario. But even with cache hits, its runtime at decoding time is significantly slower (~0.5ms per generated token on my setup, which is not negligible). If we apply torch.compile, the run times are similar, but your solution is still faster.

Here what I'm going to do:

1. In transformers: add your class, but with batch support and skip_special_tokens=False. In generate, if we are not compiling (default), your class will be used.
2. In smolagents: leave my PR as is. Because of the changes described in 1., when we merge them smolagents will automatically get the benefits :)

---

For completeness, benchmark script:

from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
from time import time
import torch

N_RUNS = 10
MAX_NEW_TOKENS = 100
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
STOP_STRINGS = ["Potato", "Carrots", "Onions", "Garlic", "Tomatoes", "Lettuce", "Cucumbers"]
BATCH_SIZE = 8

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", dtype=torch.bfloat16)
inputs = tokenizer(["The quick brown"] * BATCH_SIZE, return_tensors="pt").to(model.device)

# warmup
for i in range(3):
    gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=MAX_NEW_TOKENS, min_new_tokens=MAX_NEW_TOKENS)

all_times = []
for i in range(N_RUNS):
    start_time = time()
    gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=MAX_NEW_TOKENS, min_new_tokens=MAX_NEW_TOKENS)
    end_time = time()
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]
    if i == 0:
        continue
    all_times.append(end_time - start_time)

print(f"[No stopping criteria] Average time: {sum(all_times) / (N_RUNS - 1)} seconds")


all_times = []
for i in range(N_RUNS):
    start_time = time()
    gen_out = model.generate(
        **inputs,
        do_sample=False,
        stop_strings=STOP_STRINGS,
        tokenizer=tokenizer,
        max_new_tokens=MAX_NEW_TOKENS,
        min_new_tokens=MAX_NEW_TOKENS,
    )
    end_time = time()
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]
    if i == 0:
        continue
    all_times.append(end_time - start_time)

print(f"[transformers stopping criteria] Average time: {sum(all_times) / (N_RUNS - 1)} seconds")


class StopOnStrings(StoppingCriteria):
    def __init__(self, stop_strings: list[str], tokenizer):
        self.stop_strings = stop_strings
        self.tokenizer = tokenizer
        self.prev_tail = None
        if any(s == "" for s in self.stop_strings):
            raise ValueError("stop_strings must be non-empty")
        # Keep only last (max_len - 1) chars
        self.max_tail_len = max(len(s) for s in self.stop_strings) - 1

    def reset(self):
        self.prev_tail = None

    def __call__(self, input_ids, scores, **kwargs):
        """
        Return True if any stop string appears in the generated tokens stream; otherwise return False.
        """
        should_stop = torch.zeros_like(input_ids[:, -1], dtype=torch.bool, device=input_ids.device)
        generated = self.tokenizer.batch_decode(input_ids[:, -1])

        if all(g == "" for g in generated):
            return should_stop

        # Build the minimal search window that contains every possible
        # newly created occurrence of any stop string
        if self.prev_tail is None:
            self.prev_tail = generated
        else:
            for batch_idx in range(len(generated)):
                self.prev_tail[batch_idx] = self.prev_tail[batch_idx] + generated[batch_idx]

        # Update the tail
        for batch_idx in range(len(generated)):
            self.prev_tail[batch_idx] = self.prev_tail[batch_idx][-self.max_tail_len :]

        # Check each stop string -- we only check inside window
        for batch_idx in range(len(generated)):
            for s in self.stop_strings:
                if s in self.prev_tail[batch_idx]:
                    should_stop[batch_idx] = True
        return should_stop

custom_stopping_criteria = StoppingCriteriaList([StopOnStrings(STOP_STRINGS, tokenizer)])

all_times = []
for i in range(N_RUNS):
    start_time = time()
    gen_out = model.generate(
        **inputs,
        do_sample=False,
        stopping_criteria=custom_stopping_criteria,
        max_new_tokens=MAX_NEW_TOKENS,
        min_new_tokens=MAX_NEW_TOKENS,
    )
    end_time = time()
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]
    if i == 0:
        continue
    all_times.append(end_time - start_time)

print(f"[custom stopping criteria] Average time: {sum(all_times) / (N_RUNS - 1)} seconds")

[MaxBourdon]

@gante Thank you, I really appreciate it!

[gante]

huggingface/transformers#40520 -- transformers-side PR

@MaxBourdon I ended up adding some additional logic to handle corner cases :) (we need to keep more than your self.max_tail_len in the buffer, in the case the last token contains extra characters)