Change cfg key in vllm serving

Author: mory91

docs/reference/vllm.md

@@ -25,11 +25,7 @@ You can then query the model in shell by passing a prompt and either
 1. a [JSON Schema][jsonschema]{:target="_blank"} specification or
 2. a [Regex][regex]{:target="_blank"} pattern
 
-<<<<<<< HEAD
-with the `schema` or `regex` parameters, respectively, to the `/generate` endpoint. If both are specified, the schema will be used. If neither is specified, the generated text will be unconstrained.
-=======
-with the `schema`, `regex` or `cfg` parameters, respectively, to the `/generate` endpoint. If both are specified, the schema will be used. If neither is specified, the generated text will be unconstrained.
->>>>>>> 43ff5c5 (Expect vllm.LLMEngine as processor's argument)
+with the `schema`, `regex` or `grammar` parameters, respectively, to the `/generate` endpoint. If both are specified, the schema will be used. If neither is specified, the generated text will be unconstrained.
 
 For example, to generate a string that matches the schema `{"type": "string"}` (any string):
 
@@ -51,6 +47,16 @@ curl http://127.0.0.1:8000/generate \
         }'
 ```
 
+To generate a string that matches the grammar `<grammar>`:
+
+```bash
+curl http://127.0.0.1:8000/generate \
+    -d '{
+        "prompt": "What is Pi? Give me the first 15 digits: ",
+        "grammar": "start: DECIMAL \r\nDIGIT: \"0\"..\"9\"\r\nINT: DIGIT+\r\nDECIMAL: INT \".\" INT? | \".\" INT"
+        }'
+```
+
 Instead of `curl`, you can also use the [requests][requests]{:target="_blank"} library from another python program.
 
 Please consult the [vLLM documentation][vllm]{:target="_blank"} for details on additional request parameters. You can also [read the code](https://github.com/outlines-dev/outlines/blob/main/outlines/serve/serve.py) in case you need to customize the solution to your needs.

examples/vllm_integration.py

@@ -15,10 +15,14 @@ class User(BaseModel):
 
 llm = vllm.LLM(model="gpt2")
 logits_processor = JSONLogitsProcessor(User, llm.llm_engine)
-result = llm.generate(
+outputs = llm.generate(
     ["A prompt", "Another prompt"],
     sampling_params=vllm.SamplingParams(
         max_tokens=100, logits_processors=[logits_processor]
     ),
 )
-print(result)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

outlines/serve/serve.py

@@ -25,6 +25,7 @@
 from vllm.utils import random_uuid
 
 from .vllm import (
+    CFGLogitsProcessor,
     JSONLogitsProcessor,
     RegexLogitsProcessor,
     _patched_apply_logits_processors,
@@ -65,10 +66,14 @@ async def generate(request: Request) -> Response:
 
     json_schema = request_dict.pop("schema", None)
     regex_string = request_dict.pop("regex", None)
+    cfg_string = request_dict.pop("grammar", None)
+
     if json_schema is not None:
         logits_processors = [JSONLogitsProcessor(json_schema, engine.engine)]
     elif regex_string is not None:
         logits_processors = [RegexLogitsProcessor(regex_string, engine.engine)]
+    elif cfg_string is not None:
+        logits_processors = [CFGLogitsProcessor(cfg_string, engine.engine)]
     else:
         logits_processors = []
 

outlines/serve/vllm.py

@@ -2,12 +2,12 @@
 import json
 import math
 from collections import defaultdict
-from typing import DefaultDict, List, Callable
+from typing import Callable, DefaultDict, List
 
 import torch
 from vllm import LLMEngine
 
-from outlines.fsm.fsm import RegexFSM, CFGFSM, FSM
+from outlines.fsm.fsm import CFGFSM, FSM, RegexFSM
 from outlines.fsm.json_schema import build_regex_from_object