A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 23:30 UTC
Add outlines fallback when JSON schema has enum (#13449)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Feb 18, 2025 at 01:49 UTC
b53d79983c273b2775456d99c0e0890aea073512
3 files changed
+86
tests/entrypoints/conftest.py
+41
@@ -141,6 +141,47 @@ def sample_definition_json_schema():
141
}
142
143
144
+@pytest.fixture
145
+def sample_enum_json_schema():
146
+ return {
147
+ "type": "object",
148
+ "properties": {
149
+ "status": {
150
+ "type": "string",
151
+ "enum": ["active", "inactive",
152
+ "pending"] # Literal values using enum
153
+ },
154
+ "priority": {
155
+ "type": "string",
156
+ "enum": ["low", "medium", "high", "critical"]
157
+ },
158
+ "category": {
159
+ "type": "object",
160
+ "properties": {
161
+ "type": {
162
+ "type": "string",
163
+ "enum": ["bug", "feature", "improvement"]
164
+ },
165
+ "severity": {
166
+ "type": "integer",
167
+ "enum": [1, 2, 3, 4,
168
+ 5] # Enum can also contain numbers
169
+ }
170
+ },
171
+ "required": ["type", "severity"]
172
+ },
173
+ "flags": {
174
+ "type": "array",
175
+ "items": {
176
+ "type": "string",
177
+ "enum": ["urgent", "blocked", "needs_review", "approved"]
178
+ }
179
+ }
180
+ },
181
+ "required": ["status", "priority", "category", "flags"]
182
+ }
183
+
184
+
185
@pytest.fixture
186
def sample_guided_choice():
187
return [
tests/entrypoints/llm/test_guided_generate.py
+41
@@ -146,6 +146,47 @@ def test_guided_definition_json_completion(sample_definition_json_schema, llm,
146
schema=sample_definition_json_schema)
147
148
149
+@pytest.mark.skip_global_cleanup
150
+@pytest.mark.parametrize("guided_decoding_backend", GUIDED_DECODING_BACKENDS)
151
+def test_guided_enum_json_completion(sample_enum_json_schema, llm,
152
+ guided_decoding_backend: str):
153
+ sampling_params = SamplingParams(temperature=1.0,
154
+ max_tokens=1000,
155
+ guided_decoding=GuidedDecodingParams(
156
+ json=sample_enum_json_schema,
157
+ backend=guided_decoding_backend))
158
+ outputs = llm.generate(prompts=[
159
+ "Create a bug report JSON that fits this schema: "
160
+ f"{sample_enum_json_schema}. Make it for a high priority critical bug."
161
+ ] * 2,
162
+ sampling_params=sampling_params,
163
+ use_tqdm=True)
164
+
165
+ assert outputs is not None
166
+
167
+ for output in outputs:
168
+ assert output is not None
169
+ assert isinstance(output, RequestOutput)
170
+ prompt = output.prompt
171
+
172
+ generated_text = output.outputs[0].text
173
+ assert generated_text is not None
174
+ print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
175
+ output_json = json.loads(generated_text)
176
+ jsonschema.validate(instance=output_json,
177
+ schema=sample_enum_json_schema)
178
+
179
+ # Additional assertions to verify enum values
180
+ assert output_json["status"] in ["active", "inactive", "pending"]
181
+ assert output_json["priority"] in ["low", "medium", "high", "critical"]
182
+ assert output_json["category"]["type"] in [
183
+ "bug", "feature", "improvement"
184
+ ]
185
+ assert output_json["category"]["severity"] in [1, 2, 3, 4, 5]
186
+ for flag in output_json["flags"]:
187
+ assert flag in ["urgent", "blocked", "needs_review", "approved"]
188
+
189
+
190
@pytest.mark.skip_global_cleanup
191
@pytest.mark.parametrize("guided_decoding_backend", GUIDED_DECODING_BACKENDS)
192
def test_guided_choice_completion(sample_guided_choice, llm,
vllm/model_executor/guided_decoding/utils.py
+4
@@ -14,6 +14,10 @@ def has_xgrammar_unsupported_json_features(schema: dict) -> bool:
14
if "pattern" in obj:
15
return True
16
17
+ # Check for enum restrictions
18
+ if "enum" in obj:
19
+ return True
20
+
21
# Check for numeric ranges
22
if obj.get("type") in ("integer", "number") and any(
23
key in obj for key in [