guided decoding with grammar (#488)

* support guided decoding with grammar

* 0.4.1 fixes

Author: saiatmakuri

docs/guides/completions.md

@@ -246,6 +246,22 @@ print(response.json())
 # {"request_id":"5b184654-96b6-4932-9eb6-382a51fdb3d5","output":{"text":"{\"myString\" : \"John Doe","num_prompt_tokens":6,"num_completion_tokens":10,"tokens":null}}
 ```
 
+=== "Guided decoding with Context-Free Grammar"
+
+```python
+from llmengine import Completion
+
+response = Completion.create(
+    model="llama-2-7b",
+    prompt="Hello, my name is",
+    max_new_tokens=10,
+    temperature=0.2,
+    guided_grammar="start: \"John\""
+)
+
+print(response.json())
+# {"request_id": "34621b44-c655-402c-a459-f108b3e49b12", "output": {"text": "John", "num_prompt_tokens": 6, "num_completion_tokens": 4, "tokens": None}}
+
 ## Which model should I use?
 
 See the [Model Zoo](../../model_zoo) for more information on best practices for which model to use for Completions.

model-engine/model_engine_server/common/dtos/llms.py

@@ -186,15 +186,19 @@ class CompletionSyncV1Request(BaseModel):
     """
     guided_json: Optional[Dict[str, Any]] = None
     """
-    JSON schema for guided decoding.
+    JSON schema for guided decoding. Only supported in vllm.
     """
     guided_regex: Optional[str] = None
     """
-    Regex for guided decoding.
+    Regex for guided decoding. Only supported in vllm.
     """
     guided_choice: Optional[List[str]] = None
     """
-    Choices for guided decoding.
+    Choices for guided decoding. Only supported in vllm.
+    """
+    guided_grammar: Optional[str] = None
+    """
+    Context-free grammar for guided decoding. Only supported in vllm.
     """
 
 
@@ -272,6 +276,10 @@ class CompletionStreamV1Request(BaseModel):
     """
     Choices for guided decoding. Only supported in vllm.
     """
+    guided_grammar: Optional[str] = None
+    """
+    Context-free grammar for guided decoding. Only supported in vllm.
+    """
 
 
 class CompletionStreamOutput(BaseModel):

model-engine/model_engine_server/domain/use_cases/llm_model_endpoint_use_cases.py

@@ -1381,16 +1381,19 @@ def validate_and_update_completion_params(
         guided_count += 1
     if request.guided_regex is not None:
         guided_count += 1
+    if request.guided_grammar is not None:
+        guided_count += 1
 
     if guided_count > 1:
         raise ObjectHasInvalidValueException(
-            "Only one of guided_json, guided_choice, guided_regex can be enabled."
+            "Only one of guided_json, guided_choice, guided_regex, guided_grammar can be enabled."
         )
 
     if (
         request.guided_choice is not None
         or request.guided_regex is not None
         or request.guided_json is not None
+        or request.guided_grammar is not None
     ) and not inference_framework == LLMInferenceFramework.VLLM:
         raise ObjectHasInvalidValueException("Guided decoding is only supported in vllm.")
 
@@ -1691,6 +1694,8 @@ async def execute(
                 vllm_args["guided_regex"] = request.guided_regex
             if request.guided_json is not None:
                 vllm_args["guided_json"] = request.guided_json
+            if request.guided_grammar is not None:
+                vllm_args["guided_grammar"] = request.guided_grammar
 
             inference_request = SyncEndpointPredictV1Request(
                 args=vllm_args,
@@ -1959,6 +1964,8 @@ async def execute(
                 args["guided_regex"] = request.guided_regex
             if request.guided_json is not None:
                 args["guided_json"] = request.guided_json
+            if request.guided_grammar is not None:
+                args["guided_grammar"] = request.guided_grammar
             args["stream"] = True
         elif model_content.inference_framework == LLMInferenceFramework.LIGHTLLM:
             args = {

model-engine/model_engine_server/inference/vllm/requirements.txt

@@ -1,2 +1,2 @@
-vllm==0.4.0.post1
+vllm==0.4.1
 pydantic>=2.0

model-engine/model_engine_server/inference/vllm/vllm_server.py

@@ -45,6 +45,7 @@ async def generate(request: Request) -> Response:
     guided_json = request_dict.pop("guided_json", None)
     guided_regex = request_dict.pop("guided_regex", None)
     guided_choice = request_dict.pop("guided_choice", None)
+    guided_grammar = request_dict.pop("guided_grammar", None)
     sampling_params = SamplingParams(**request_dict)
 
     # Dummy request to get guided decode logit processor
@@ -56,15 +57,17 @@ async def generate(request: Request) -> Response:
                 "guided_json": guided_json,
                 "guided_regex": guided_regex,
                 "guided_choice": guided_choice,
+                "guided_grammar": guided_grammar,
             }
         )
     except Exception:
         raise HTTPException(
             status_code=400, detail="Bad request: failed to parse guided decoding parameters."
         )
 
+    guided_decoding_backend = engine.engine.decoding_config.guided_decoding_backend
     guided_decode_logit_processor = await get_guided_decoding_logits_processor(
-        partial_openai_request, engine.get_tokenizer()
+        guided_decoding_backend, partial_openai_request, await engine.get_tokenizer()
     )
     if guided_decode_logit_processor is not None:
         if sampling_params.logits_processors is None:

model-engine/tests/unit/domain/test_llm_use_cases.py

@@ -1108,11 +1108,13 @@ async def test_validate_and_update_completion_params():
     completion_sync_request.guided_regex = ""
     completion_sync_request.guided_json = {}
     completion_sync_request.guided_choice = [""]
+    completion_sync_request.guided_grammar = ""
     with pytest.raises(ObjectHasInvalidValueException):
         validate_and_update_completion_params(LLMInferenceFramework.VLLM, completion_sync_request)
 
     completion_sync_request.guided_regex = None
     completion_sync_request.guided_choice = None
+    completion_sync_request.guided_grammar = None
     with pytest.raises(ObjectHasInvalidValueException):
         validate_and_update_completion_params(
             LLMInferenceFramework.TEXT_GENERATION_INFERENCE, completion_sync_request