Add support for additional constrained decoding options

README.md

@@ -224,8 +224,30 @@ curl -X 'POST' 'http://127.0.0.1:8000/stream/process' \
 ```
 will result in a response like {"doc_name": "DOC", "start": INT, "end": INT, "label_name": "STR", "label_id": "STR", ...}\n...
 
+#### Chat with served models
+You can also "chat" with the running model using the `/stream/ws` endpoint. For example:
+```html
+<form action="" onsubmit="send_doc(event)">
+    <input type="text" id="cms-input" autocomplete="off"/>
+    <button>Send</button>
+</form>
+<ul id="cms-output"></ul>
+<script>
+    var ws = new WebSocket("ws://localhost:8000/stream/ws");
+    ws.onmessage = function(event) {
+        document.getElementById("cms-output").appendChild(
+            Object.assign(document.createElement('li'), { textContent: event.data })
+        );
+    };
+    function send_doc(event) {
+        ws.send(document.getElementById("cms-input").value);
+        event.preventDefault();
+    };
+</script>
+```
+
 ### Serve causal language models
-You can serve causal LLMs (e.g., LLaMa 3 or DeepSeek R1) using the `huggingface_llm` model type. To do so, run:
+You can serve open-weight causal LLMs (e.g., LLaMa 3, DeepSeek R1, Qwen3 or Mistral) using the `huggingface_llm` model type. To do so, run:
 ```commandline
 cms serve --model-type huggingface_llm --model-path PATH/TO/MODEL_PACKAGE.zip --host 127.0.0.1 --port 8000
 ```
@@ -236,11 +258,11 @@ curl -N -X 'POST' 'http://127.0.0.1:8000/stream/generate?max_tokens=512' \
     -d 'What is hypertension?'
 ```
 
-CMS also provides APIs that are compatible with the OpenAI client. For example, using its Python SDK, you can generate text and retrieve embeddings with the following snippet:
+CMS also provides APIs that are compatible with OpenAI and Ollama clients. For example, using their Python SDKs, you can generate text and retrieve embeddings with the following snippet:
 ```python
 from openai import OpenAI
 
-client = OpenAI(api_key="dummy", base_url="http://localhost:8000/v1")
+client = OpenAI(api_key="dummy", base_url="http://localhost:8000/openai/v1")
 completion = client.chat.completions.create(
     model="Huggingface LLM Model",
     messages=[
@@ -256,6 +278,29 @@ embeddings = client.embeddings.create(
     model="Huggingface LLM Model", input="What is hypertension?"
 )
 ```
+```python
+from ollama import Client
+
+client = Client(host="http://localhost:8000//ollama")
+
+response = client.chat(
+    model="Huggingface LLM Model",
+    messages=[
+        {"role": "system", "content": "You are bot answering questions from the user"},
+        {"role": "user", "content": "What is the normal LVEF range?"},
+    ],
+    stream=False,
+    options={
+        "num_predict": 512,
+        "temperature": 0.7,
+        "top_p": 0.9
+    },
+)
+print(f"CMS => {response["message"]["content"]}")
+embeddings = client.embed(
+    model="Huggingface LLM Model", input="What is the normal LVEF range?"
+)
+```
 Note that to enable quantization and training features, you need to install the extra dependencies:
 ```commandline
 pip install '.[llm]'
@@ -268,25 +313,3 @@ To that end, install the following extra dependencies:
 pip install '.[mcp]'
 ```
 For detailed configuration, please refer to the [CMS MCP Server docs](./app/mcp/README.md).
-
-#### Chat with served models
-You can also "chat" with the running model using the `/stream/ws` endpoint. For example:
-```html
-<form action="" onsubmit="send_doc(event)">
-    <input type="text" id="cms-input" autocomplete="off"/>
-    <button>Send</button>
-</form>
-<ul id="cms-output"></ul>
-<script>
-    var ws = new WebSocket("ws://localhost:8000/stream/ws");
-    ws.onmessage = function(event) {
-        document.getElementById("cms-output").appendChild(
-            Object.assign(document.createElement('li'), { textContent: event.data })
-        );
-    };
-    function send_doc(event) {
-        ws.send(document.getElementById("cms-input").value);
-        event.preventDefault();
-    };
-</script>
-```

app/api/routers/generative.py

@@ -55,6 +55,16 @@
     cms_tpot_milliseconds,
 )
 from app.exception import GenerationException, ClientException
+from app.processors.constrained_decoder import ConstrainedDecoder
+try:
+    import xgrammar as xgr
+except ImportError:
+    xgr = None  # type: ignore[assignment]
+
+try:
+    from llguidance import LLMatcher
+except ImportError:
+    LLMatcher = None  # type: ignore[assignment,misc]
 
 
 PATH_GENERATE = "/generate"
@@ -134,6 +144,7 @@ def generate_text(
                 override_template=chat_template if chat_template else None,
             ),
             max_tokens=max_tokens,
+            num_beams=config.DECODING_NUM_BEAMS,
             temperature=temperature,
             top_p=top_p,
             stop_sequences=stop_sequences,
@@ -214,6 +225,7 @@ async def _stream(prompt: str) -> AsyncGenerator:
         async for generated in model_service.generate_async(
             prompt=prompt,
             max_tokens=max_tokens,
+            num_beams=config.DECODING_NUM_BEAMS,
             temperature=temperature,
             top_p=top_p,
             stop_sequences=stop_sequences,
@@ -231,6 +243,12 @@ async def _stream(prompt: str) -> AsyncGenerator:
                        f"TPOT in milliseconds: {str(generated.tpot_ms)}"
                        "</cms_token_usage>"
                     )
+                if generated.timed_out:
+                    yield (
+                        "\n\n<cms_generation_timed_out>"
+                        "Generation timed out before completion so the returned content may be partial."
+                        "</cms_generation_timed_out>"
+                    )
                 continue
             yield generated
 
@@ -293,7 +311,7 @@ def generate_chat_completions(
     max_tokens = request_data.max_tokens
     temperature = request_data.temperature
     top_p = request_data.top_p
-    json_schema_parser = _get_parser_for_response_format(request_data.response_format)
+    json_schema_parser = _get_parser_for_openai_compatible(request_data.response_format, model_service=model_service)
 
     if isinstance(request_data.stop, str):
         stop_sequences = [request_data.stop]
@@ -340,6 +358,7 @@ async def _stream(
             async for generated in model_service.generate_async(
                 prompt,
                 max_tokens=max_tokens,
+                num_beams=config.DECODING_NUM_BEAMS,
                 temperature=temperature,
                 top_p=top_p,
                 stop_sequences=stop_sequences,
@@ -358,6 +377,14 @@ async def _stream(
                             }
                         }
                         yield f"data: {json.dumps(data)}\n\n"
+                    if generated.timed_out:
+                        error_data = {
+                            "error": {
+                                "message": "Generation timed out before completion and the returned content may be partial.",
+                                "type": "generation_error",
+                            }
+                        }
+                        yield f"data: {json.dumps(error_data)}\n\n"
                     continue
                 if tool_call_emitted:
                     continue
@@ -452,6 +479,7 @@ def _report_tokens(
         generation_result = model_service.generate(
             prompt,
             max_tokens=max_tokens,
+            num_beams=config.DECODING_NUM_BEAMS,
             temperature=temperature,
             top_p=top_p,
             stop_sequences=stop_sequences or [],
@@ -587,6 +615,7 @@ async def _stream(
             async for generated in model_service.generate_async(
                 prompt,
                 max_tokens=max_tokens,
+                num_beams=config.DECODING_NUM_BEAMS,
                 temperature=temperature,
                 top_p=top_p,
                 stop_sequences=stop_sequences,
@@ -656,6 +685,7 @@ def _report_tokens(
         generation_result = model_service.generate(
             prompt,
             max_tokens=max_tokens,
+            num_beams=config.DECODING_NUM_BEAMS,
             temperature=temperature,
             top_p=top_p,
             stop_sequences=stop_sequences,
@@ -972,7 +1002,7 @@ async def ollama_generate(
     temperature = options.get("temperature", 0.7)
     top_p = options.get("top_p", 0.9)
     stop_sequences = _normalise_stop_sequences(options.get("stop", None))
-    json_schema_parser = _get_parser_for_json_schema(request_data.format)
+    json_schema_parser = _get_parser_for_ollama_compatible(request_data.format, model_service=model_service)
 
     def _report_tokens(
         prompt_token_num: int,
@@ -1001,6 +1031,7 @@ async def _stream() -> AsyncGenerator[str, None]:
                 async for generated in model_service.generate_async(
                     prompt,
                     max_tokens=max_tokens,
+                    num_beams=config.DECODING_NUM_BEAMS,
                     temperature=temperature,
                     top_p=top_p,
                     stop_sequences=stop_sequences,
@@ -1034,7 +1065,7 @@ async def _stream() -> AsyncGenerator[str, None]:
                 "created_at": _iso_utc_now(),
                 "response": "",
                 "done": True,
-                "done_reason": "stop",
+                "done_reason": "timed_out" if generation_result is not None and generation_result.timed_out else "stop",
                 "prompt_eval_count": generation_result.prompt_token_num if generation_result is not None else 0,
                 "eval_count": generation_result.completion_token_num if generation_result is not None else 0,
                 "ttft_in_milliseconds": generation_result.ttft_ms if generation_result is not None else -1,
@@ -1048,6 +1079,7 @@ async def _stream() -> AsyncGenerator[str, None]:
         generation_result = model_service.generate(
             prompt,
             max_tokens=max_tokens,
+            num_beams=config.DECODING_NUM_BEAMS,
             temperature=temperature,
             top_p=top_p,
             stop_sequences=stop_sequences,
@@ -1095,7 +1127,7 @@ async def ollama_chat(
     temperature = options.get("temperature", 0.7)
     top_p = options.get("top_p", 0.9)
     stop_sequences = _normalise_stop_sequences(options.get("stop", None))
-    json_schema_parser = _get_parser_for_json_schema(request_data.format)
+    json_schema_parser = _get_parser_for_ollama_compatible(request_data.format, model_service=model_service)
     _ensures_chat_template(model_service, chat_template)
 
     def _report_tokens(
@@ -1142,6 +1174,7 @@ async def _stream() -> AsyncGenerator[str, None]:
                 async for generated in model_service.generate_async(
                     prompt,
                     max_tokens=max_tokens,
+                    num_beams=config.DECODING_NUM_BEAMS,
                     temperature=temperature,
                     top_p=top_p,
                     stop_sequences=stop_sequences,
@@ -1176,7 +1209,7 @@ async def _stream() -> AsyncGenerator[str, None]:
                 "created_at": _iso_utc_now(),
                 "message": {"role": PromptRole.ASSISTANT.value, "content": ""},
                 "done": True,
-                "done_reason": "stop",
+                "done_reason": "timed_out" if generated_result is not None and generated_result.timed_out else "stop",
                 "prompt_eval_count": generated_result.prompt_token_num if generated_result is not None else 0,
                 "eval_count": generated_result.completion_token_num if generated_result is not None else 0,
                 "ttft_in_milliseconds": generated_result.ttft_ms if generated_result is not None else -1,
@@ -1190,6 +1223,7 @@ async def _stream() -> AsyncGenerator[str, None]:
         generated_result = model_service.generate(
             prompt,
             max_tokens=max_tokens,
+            num_beams=config.DECODING_NUM_BEAMS,
             temperature=temperature,
             top_p=top_p,
             stop_sequences=stop_sequences,
@@ -1267,33 +1301,35 @@ def _build_prompt_text(
     )
 
 
-def _get_parser_for_response_format(response_format: Optional[OpenAIResponseFormat]) -> Optional[Any]:
+def _get_parser_for_openai_compatible(
+    response_format: Optional[OpenAIResponseFormat],
+    model_service: Optional[AbstractModelService] = None,
+) -> Optional[Any]:
     if response_format is None:
         return None
     if response_format.type == "json_schema":
-        try:
-            from lmformatenforcer import JsonSchemaParser
-            parser = JsonSchemaParser(response_format.json_schema.schema_)
-            setattr(parser, "schema", response_format.json_schema.schema_)
-            return parser
-        except ImportError as e:
-            raise ClientException("lmformatenforcer package is not installed; required for JSON schema support") from e
-        except Exception as exc:
-            raise ClientException("Invalid JSON schema in response_format") from exc
+        return ConstrainedDecoder.get_parser(
+            response_format.json_schema.schema_,
+            model_service=model_service,
+            source="response_format",
+            backend=config.DECODING_BACKEND,
+        )
     else:
         raise ClientException("Unsupported response_format type; only 'json_schema' is supported")
 
 
-def _get_parser_for_json_schema(json_schema: Optional[Dict[str, Any]]) -> Optional[Any]:
+def _get_parser_for_ollama_compatible(
+    json_schema: Optional[Dict[str, Any]],
+    model_service: Optional[AbstractModelService] = None,
+) -> Optional[Any]:
     if json_schema is None:
         return None
-    try:
-        from lmformatenforcer import JsonSchemaParser
-        return JsonSchemaParser(json_schema)
-    except ImportError as e:
-        raise ClientException("lmformatenforcer package is not installed; required for JSON schema support") from e
-    except Exception as exc:
-        raise ClientException("Invalid JSON schema") from exc
+    return ConstrainedDecoder.get_parser(
+        json_schema,
+        model_service=model_service,
+        source="json_schema",
+        backend=config.DECODING_BACKEND,
+    )
 
 
 def _send_usage_metrics(

app/api/routers/stream.py

@@ -176,7 +176,11 @@ async def _monitor_idle() -> None:
             logger.debug(str(e))
 
 
-@router.get(PATH_SSE_EVENTS)
+@router.get(
+    PATH_SSE_EVENTS,
+    tags=[Tags.Annotations.name],
+    description="Server-Sent Events (SSE) endpoint to receive NER entities as stream events for a specific client",
+)
 @limiter.exempt
 async def get_entities_stream_from_sse(
     request: Request,

app/config.py

@@ -48,6 +48,8 @@ class Settings(BaseSettings):   # type: ignore
     ENABLE_SPDA_ATTN: str = "true"                    # if "true", attempt to use SPDA attention for HuggingFace LLM loading
     ASSISTANT_MODEL_FULL_PATH: str = ""               # the full path to the assistant model package for speculative decoding
     OVERRIDE_CHAT_TEMPLATE: str = ""                  # if set, override the chat template used for prompt formatting
+    DECODING_BACKEND: str = "lm_format_enforcer"      # the decoding backend to use for LLM constrained decoding, one of "lm_format_enforcer", "xgrammar" or "llguidance"
+    DECODING_NUM_BEAMS: int = 1                       # the number of beams to used for beam search during LLM decoding
     DEBUG: str = "false"                              # if "true", the debug mode is switched on
 
     class Config:

app/domain.py

@@ -137,6 +137,13 @@ class LlmDatasetType(Enum):
     CSV = "csv"
 
 
+class LLMDecodingBackend(Enum):
+    LM_FORMAT_ENFORCER = "lm_format_enforcer"
+    XGRAMMAR = "xgrammar"
+    LLGUIDANCE = "llguidance"
+
+
+
 class Annotation(BaseModel):
     doc_name: Optional[str] = Field(default=None, description="The name of the document to which the annotation belongs")
     start: int = Field(description="The start index of the annotation span")
@@ -220,6 +227,10 @@ class GenerationResult(BaseModel):
     completion_token_num: int = Field(..., description="The number of tokens in the generated text")
     ttft_ms: int = Field(default=-1, description="Time to first token in milliseconds")
     tpot_ms: int = Field(default=-1, description="Average time per output token in milliseconds")
+    timed_out: bool = Field(
+        default=False,
+        description="Whether the generation was stopped by the generation timeout criteria before completing",
+    )
 
 
 class OpenAIStreamOptions(BaseModel):

app/envs/.env

@@ -53,7 +53,7 @@ PROCESS_RATE_LIMIT=180/minute
 PROCESS_BULK_RATE_LIMIT=90/minute
 
 # The rate limit on the text generation routes
-GENERATION_RATE_LIMIT=10/minute
+GENERATION_RATE_LIMIT=1000/minute
 
 # The timeout in seconds on the text generation requests
 GENERATION_TIMEOUT_SECONDS=180
@@ -110,5 +110,11 @@ ASSISTANT_MODEL_FULL_PATH=
 # If set, override the chat template used for prompt formatting
 OVERRIDE_CHAT_TEMPLATE=
 
+# the decoding backend to use for LLM constrained decoding, either "lm_format_enforcer", "xgrammar", or "llguidance"
+DECODING_BACKEND=lm_format_enforcer
+
+# the number of beams to used for beam search during LLM decoding
+DECODING_NUM_BEAMS=1
+
 # If "true", the debug mode is switched on
 DEBUG=false