| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This file demonstrates the example usage of structured outputs |
| 4 | in vLLM. It shows how to apply different constraints such as choice, |
| 5 | regex, json schema, and grammar to produce structured and formatted |
| 6 | results based on specific prompts. |
| 7 | """ |
| 8 | |
| 9 | from enum import Enum |
| 10 | |
| 11 | from pydantic import BaseModel |
| 12 | |
| 13 | from vllm import LLM, SamplingParams |
| 14 | from vllm.sampling_params import StructuredOutputsParams |
| 15 | |
| 16 | MAX_TOKENS = 50 |
| 17 | |
| 18 | # Structured outputs by Choice (list of possible options) |
| 19 | structured_outputs_params_choice = StructuredOutputsParams( |
| 20 | choice=["Positive", "Negative"] |
| 21 | ) |
| 22 | sampling_params_choice = SamplingParams( |
| 23 | structured_outputs=structured_outputs_params_choice |
| 24 | ) |
| 25 | prompt_choice = "Classify this sentiment: vLLM is wonderful!" |
| 26 | |
| 27 | # Structured outputs by Regex |
| 28 | structured_outputs_params_regex = StructuredOutputsParams(regex=r"\w+@\w+\.com\n") |
| 29 | sampling_params_regex = SamplingParams( |
| 30 | structured_outputs=structured_outputs_params_regex, |
| 31 | stop=["\n"], |
| 32 | max_tokens=MAX_TOKENS, |
| 33 | ) |
| 34 | prompt_regex = ( |
| 35 | "Generate an email address for Alan Turing, who works in Enigma." |
| 36 | "End in .com and new line. Example result:" |
| 37 | "alan.turing@enigma.com\n" |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | # Structured outputs by JSON using Pydantic schema |
| 42 | class CarType(str, Enum): |
| 43 | sedan = "sedan" |
| 44 | suv = "SUV" |
| 45 | truck = "Truck" |
| 46 | coupe = "Coupe" |
| 47 | |
| 48 | |
| 49 | class CarDescription(BaseModel): |
| 50 | brand: str |
| 51 | model: str |
| 52 | car_type: CarType |
| 53 | |
| 54 | |
| 55 | json_schema = CarDescription.model_json_schema() |
| 56 | structured_outputs_params_json = StructuredOutputsParams(json=json_schema) |
| 57 | sampling_params_json = SamplingParams( |
| 58 | structured_outputs=structured_outputs_params_json, max_tokens=MAX_TOKENS |
| 59 | ) |
| 60 | prompt_json = ( |
| 61 | "Generate a JSON with the brand, model and car_type of " |
| 62 | "the most iconic car from the 90's" |
| 63 | ) |
| 64 | |
| 65 | # Structured outputs by Grammar |
| 66 | simplified_sql_grammar = """ |
| 67 | root ::= select_statement |
| 68 | select_statement ::= "SELECT " column " from " table " where " condition |
| 69 | column ::= "col_1 " | "col_2 " |
| 70 | table ::= "table_1 " | "table_2 " |
| 71 | condition ::= column "= " number |
| 72 | number ::= "1 " | "2 " |
| 73 | """ |
| 74 | structured_outputs_params_grammar = StructuredOutputsParams( |
| 75 | grammar=simplified_sql_grammar |
| 76 | ) |
| 77 | sampling_params_grammar = SamplingParams( |
| 78 | structured_outputs=structured_outputs_params_grammar, |
| 79 | max_tokens=MAX_TOKENS, |
| 80 | ) |
| 81 | prompt_grammar = ( |
| 82 | "Generate an SQL query to show the 'username' and 'email' from the 'users' table." |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | def format_output(title: str, output: str): |
| 87 | print(f"{'-' * 50}\n{title}: {output}\n{'-' * 50}") |
| 88 | |
| 89 | |
| 90 | def generate_output(prompt: str, sampling_params: SamplingParams, llm: LLM): |
| 91 | outputs = llm.generate(prompt, sampling_params=sampling_params) |
| 92 | return outputs[0].outputs[0].text |
| 93 | |
| 94 | |
| 95 | def main(): |
| 96 | llm = LLM(model="Qwen/Qwen2.5-3B-Instruct", max_model_len=100) |
| 97 | |
| 98 | choice_output = generate_output(prompt_choice, sampling_params_choice, llm) |
| 99 | format_output("Structured outputs by Choice", choice_output) |
| 100 | |
| 101 | regex_output = generate_output(prompt_regex, sampling_params_regex, llm) |
| 102 | format_output("Structured outputs by Regex", regex_output) |
| 103 | |
| 104 | json_output = generate_output(prompt_json, sampling_params_json, llm) |
| 105 | format_output("Structured outputs by JSON", json_output) |
| 106 | |
| 107 | grammar_output = generate_output(prompt_grammar, sampling_params_grammar, llm) |
| 108 | format_output("Structured outputs by Grammar", grammar_output) |
| 109 | |
| 110 | |
| 111 | if __name__ == "__main__": |
| 112 | main() |