transformers で特定のキーワードが生成されたら生成を止める

from transformers import StopStringCriteria, pipeline

pipe = pipeline(
    "text-generation", model="HuggingFaceTB/SmolLM2-135M-Instruct", torch_dtype=torch.bfloat16
)

stop_string = StopStringCriteria(
    tokenizer=pipe.tokenizer, # トークナイズしないと検出できないのでトークナイザーが必要
    stop_strings=["キーワード", "文字列を直接指定するか配列で複数渡せる", "\n"],
)

prompt = "hello, "

# 生成時の `stopping_criteria` で指定する
# GenerationMixin 系は generate() するときに渡せば良い。
# 今回は pipeline だが、それも他の 生成オプションと一緒に渡せばいい
pipe(
    prompt,
    max_new_tokens=128, 
    stopping_criteria=[stop_string], # stopping_criteria に **配列** として渡す
)