diff --git a/src/raglite/_chatml_function_calling.py b/src/raglite/_chatml_function_calling.py index e028744..d60ac82 100644 --- a/src/raglite/_chatml_function_calling.py +++ b/src/raglite/_chatml_function_calling.py @@ -296,11 +296,11 @@ def chatml_function_calling_with_streaming( # Collect the llama.create_completion keyword arguments so we don't have to repeat these with # each completion call stop = ( - [stop, "<|im_end|>", "|im_end|>"] + [stop, "<|im_end|>"] if isinstance(stop, str) - else [*stop, "<|im_end|>", "|im_end|>"] + else [*stop, "<|im_end|>"] if stop - else ["<|im_end|>", "|im_end|>"] + else ["<|im_end|>"] ) grammar = ( # It is assumed the grammar applies to messages only, not tool calls grammar