Repeated line breaks when using JSON Logits Processor (vLLM)

[LilianJim]

Describe the issue as clearly as possible:

When using a JSONLogitsProcessor to constrain the output, a bug that often occurs is the model generated line breaks (\n) indefinitely.
I don't think this is the expected behavior, and for now I've only witnessed it with Mixtral 7x8B

Steps/code to reproduce the bug:

reproduce.py

# Set seeds for reproducibility
import torch
torch.manual_seed(0)
torch.cuda.manual_seed(0)
torch.backends.cudnn.deterministic = True
import numpy
numpy.random.seed(0)

import argparse
import asyncio
import json

from types import SimpleNamespace

from outlines.serve.vllm import JSONLogitsProcessor
from transformers import AutoTokenizer
from pydantic import BaseModel

from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.sampling_params import SamplingParams
from vllm.utils import random_uuid


class Question(BaseModel):
    class Statement(BaseModel):
        statement: str
        source: int
    statements: list[Statement]


async def main():
    parser = argparse.ArgumentParser()
    parser = AsyncEngineArgs.add_cli_args(parser)
    args = parser.parse_args()
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    # Double-wrap the tokenizer so it can be accessed with `llm.tokenizer.tokenizer`
    dummy_llm = SimpleNamespace(
        tokenizer=SimpleNamespace(
            tokenizer=tokenizer,
        )
    )
    schema = Question.model_json_schema()
    print(json.dumps(Question.model_json_schema(), indent=2))
    json_logits_processor = JSONLogitsProcessor(schema, dummy_llm)
    engine_args = AsyncEngineArgs.from_cli_args(args)
    engine = AsyncLLMEngine.from_engine_args(engine_args)
    sampling_params = SamplingParams(
        logits_processors=[json_logits_processor],
        max_tokens=256,
    )
    request_id = random_uuid()
    results_generator = engine.generate(
        prompt=tokenizer.apply_chat_template([{
            "role": "user",
            "content": (
                "What's the country with the largest population on Earth? "
                "And what's special about it? About its culture? "
                "Cite sources 1 and 6 about in your answer. "
                "Don't make any assumption about the sources, just cite them naturally."
            )
        }], tokenize=False),
        sampling_params=sampling_params,
        request_id=request_id,
    )
    # Pretty-print the results by only printing the next token
    previous_output_cursor = 0
    async for request_output in results_generator:
        text = request_output.outputs[0].text
        print(text[previous_output_cursor:], end="")
        previous_output_cursor = len(text)

if __name__ == "__main__":
    asyncio.run(main())

Expected result:

A valid JSON schema, e.g.

{"statements": [{"statement": "The country with the largest population on Earth is China. According to Source 1, China has a population of over 1.4 billion people.", "source": 1]}

Error message:

Output of python3 reproduce.py --model casperhansen/mixtral-instruct-awq --quantization awq

{
  "$defs": {
    "Statement": {
      "properties": {
        "statement": {
          "title": "Statement",
          "type": "string"
        },
        "source": {
          "title": "Source",
          "type": "integer"
        }
      },
      "required": [
        "statement",
        "source"
      ],
      "title": "Statement",
      "type": "object"
    }
  },
  "properties": {
    "statements": {
      "items": {
        "$ref": "#/$defs/Statement"
      },
      "title": "Statements",
      "type": "array"
    }
  },
  "required": [
    "statements"
  ],
  "title": "Question",
  "type": "object"
}
WARNING 02-28 18:10:00 config.py:186] awq quantization is not fully optimized yet. The speed can be slower than non-quantized models.
INFO 02-28 18:10:00 llm_engine.py:79] Initializing an LLM engine with config: model='casperhansen/mixtral-instruct-awq', tokenizer='casperhansen/mixtral-instruct-awq', tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=32768, download_dir=None, load_format=auto, tensor_parallel_size=1, disable_custom_all_reduce=False, quantization=awq, enforce_eager=False, kv_cache_dtype=auto, device_config=cuda, seed=0)
INFO 02-28 18:10:02 weight_utils.py:163] Using model weights format ['*.safetensors']
INFO 02-28 18:10:12 llm_engine.py:337] # GPU blocks: 22278, # CPU blocks: 2048
INFO 02-28 18:10:13 model_runner.py:676] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
INFO 02-28 18:10:13 model_runner.py:680] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
INFO 02-28 18:10:25 model_runner.py:748] Graph capturing finished in 11 secs.
INFO 02-28 18:10:25 async_llm_engine.py:433] Received request e52a4b7b96a144fc8de912d23fe14c1d: prompt: "<s>[INST] What's the country with the largest population on Earth? And what's special about it? About its culture? Cite sources 1 and 6 about in your answer. Don't make any assumption about the sources, just cite them naturally. [/INST]", prefix_pos: None,sampling_params: SamplingParams(n=1, best_of=1, presence_penalty=0.0, frequency_penalty=0.0, repetition_penalty=1.0, temperature=1.0, top_p=1.0, top_k=-1, min_p=0.0, seed=None, use_beam_search=False, length_penalty=1.0, early_stopping=False, stop=[], stop_token_ids=[], include_stop_str_in_output=False, ignore_eos=False, max_tokens=256, logprobs=None, prompt_logprobs=None, skip_special_tokens=True, spaces_between_special_tokens=True), prompt_token_ids: None, lora_request: None.
{"statements": [
{"statement": "The country with the largest population on Earth is China. According to Source 1, China has a population of over 1.4 billion people."
, "source":
 
  
 
  
 
 
  
  
 
  
 

 



 

 

 

INFO 02-28 18:10:30 metrics.py:161] Avg prompt throughput: 12.4 tokens/s, Avg generation throughput: 16.8 tokens/s, Running: 1 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 0.0%, CPU KV cache usage: 0.0%
 

 

 

 

 

 

 

 

 

 

 




 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

INFO 02-28 18:10:35 metrics.py:161] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 17.3 tokens/s, Running: 1 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 0.1%, CPU KV cache usage: 0.0%
 

 

 

 

 

 

 

 




 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

INFO 02-28 18:10:39 async_llm_engine.py:110] Finished request e52a4b7b96a144fc8de912d23fe14c1d.

Outlines/Python version information:

Version information

0.0.34.dev4+gd938678
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0]

vllm==0.3.2
tokenizers=0.15.2
transformers==4.38.1
outlines @ git+https://github.com/outlines-dev/outlines@d938678d6e6cb5ddd7b87e92d4badb14e9e56ae8

Context for the issue:

Found a workaround by passing whitespace_pattern=r"[\n ]?" to JSONLogitsProcessor, which seems to mitigate the issue. Despite this, the described default behavior is probably not wanted.

[rlouf]

The "workaround" is the solution to this; the default behavior of Outlines is to let the model choose the whitespace pattern but not all models seem to fare fine with that.

[jqueguiner]

is there a way to pass this parameter to vllm ?