[Test] Merge all constrained decoding tests. (#12633)
This commit is contained in:
Liangsheng Yin
224
python/sglang/test/kits/ebnf_constrained_kit.py
Normal file
224
python/sglang/test/kits/ebnf_constrained_kit.py
Normal file
@@ -0,0 +1,224 @@
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
class TestEBNFConstrainedMinxin:
|
||||
ebnf_grammar = 'root ::= "test"' # Default grammar
|
||||
|
||||
def _run_decode_ebnf(
|
||||
self,
|
||||
ebnf,
|
||||
expected_patterns,
|
||||
prompt,
|
||||
return_logprob=False,
|
||||
top_logprobs_num=0,
|
||||
n=1,
|
||||
):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0 if n == 1 else 0.5,
|
||||
"max_new_tokens": 128,
|
||||
"n": n,
|
||||
"ebnf": ebnf,
|
||||
},
|
||||
"stream": False,
|
||||
"return_logprob": return_logprob,
|
||||
"top_logprobs_num": top_logprobs_num,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
|
||||
ret = response.json()
|
||||
print(json.dumps(ret, indent=2))
|
||||
print("=" * 100)
|
||||
|
||||
if not isinstance(ret, list):
|
||||
self.fail(f"Expected response to be a list, but got {type(ret)}")
|
||||
|
||||
for item in ret:
|
||||
text = item.get("text", "").strip()
|
||||
if not text:
|
||||
self.fail("Generated text is empty.")
|
||||
|
||||
match = False
|
||||
for pattern in expected_patterns:
|
||||
if self.regex_match(text, pattern):
|
||||
match = True
|
||||
break
|
||||
if not match:
|
||||
self.fail(f"Text '{text}' does not match any of the allowed patterns.")
|
||||
|
||||
def regex_match(self, text, pattern):
|
||||
import re
|
||||
|
||||
return re.match(pattern, text) is not None
|
||||
|
||||
def test_ebnf_generate_email(self):
|
||||
self.__class__.ebnf_grammar = 'root ::= "user@example.com"'
|
||||
allowed_patterns = [r"^user@example\.com$"]
|
||||
prompt = "Generate an email address:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_greeting(self):
|
||||
self.__class__.ebnf_grammar = 'root ::= "Hello" | "Hi" | "Hey"'
|
||||
allowed_patterns = [r"^(Hello|Hi|Hey)$"]
|
||||
prompt = "Generate a greeting:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_number(self):
|
||||
self.__class__.ebnf_grammar = """
|
||||
root ::= digit digit digit
|
||||
digit ::= [0-9]
|
||||
"""
|
||||
allowed_patterns = [r"^\d{3}$"]
|
||||
prompt = "Generate a three-digit number:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_phone(self):
|
||||
self.__class__.ebnf_grammar = """
|
||||
root ::= "(" area ")" " " prefix "-" line
|
||||
area ::= [0-9] [0-9] [0-9]
|
||||
prefix ::= [0-9] [0-9] [0-9]
|
||||
line ::= [0-9] [0-9] [0-9] [0-9]
|
||||
"""
|
||||
allowed_patterns = [r"^\(\d{3}\) \d{3}-\d{4}$"]
|
||||
prompt = "Generate a phone number:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_date(self):
|
||||
self.__class__.ebnf_grammar = """
|
||||
root ::= year "-" month "-" day
|
||||
year ::= "2024"
|
||||
month ::= "01" | "02" | "03" | "04" | "05" | "06" | "07" | "08" | "09" | "10" | "11" | "12"
|
||||
day ::= "01" | "02" | "03" | "04" | "05" | "06" | "07" | "08" | "09" | "10" |
|
||||
"11" | "12" | "13" | "14" | "15" | "16" | "17" | "18" | "19" | "20" |
|
||||
"21" | "22" | "23" | "24" | "25" | "26" | "27" | "28" | "29" | "30" | "31"
|
||||
"""
|
||||
allowed_patterns = [r"^2024-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$"]
|
||||
prompt = "Generate a date in YYYY-MM-DD format:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_hex_color(self):
|
||||
self.__class__.ebnf_grammar = """
|
||||
root ::= "#" hex hex hex hex hex hex
|
||||
hex ::= [0-9] | [A-F]
|
||||
"""
|
||||
allowed_patterns = [r"^#[0-9A-F]{6}$"]
|
||||
prompt = "Generate a hex color code:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_complex_json(self):
|
||||
self.__class__.ebnf_grammar = """
|
||||
root ::= object
|
||||
object ::= "{" ws pair (ws "," ws pair)* ws "}"
|
||||
pair ::= "\\"name\\"" ws ":" ws value |
|
||||
"\\"age\\"" ws ":" ws number |
|
||||
"\\"city\\"" ws ":" ws string
|
||||
value ::= string | number
|
||||
string ::= "\\"" [a-zA-Z0-9 ]+ "\\""
|
||||
number ::= [1-9] [0-9]*
|
||||
ws ::= [ ]*
|
||||
"""
|
||||
allowed_patterns = [
|
||||
r'^{\s*"name"\s*:\s*"[a-zA-Z0-9 ]+"\s*,\s*"age"\s*:\s*[1-9][0-9]*\s*,\s*"city"\s*:\s*"[a-zA-Z0-9 ]+"\s*}$',
|
||||
]
|
||||
prompt = "Generate a simple JSON with name, age, and city:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_custom_log_format(self):
|
||||
self.__class__.ebnf_grammar = """
|
||||
root ::= logentry
|
||||
logentry ::= "[" datetime "] " level ": System.process - " message
|
||||
datetime ::= "2024-01-01T12:00:00Z"
|
||||
level ::= "INFO"
|
||||
message ::= "Operation " [a-z]+ " successfully"
|
||||
"""
|
||||
allowed_patterns = [
|
||||
r"^\[2024-01-01T12:00:00Z\] INFO: System\.process - Operation [a-z]+ successfully$"
|
||||
]
|
||||
prompt = "Generate a log entry:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_ebnf_generate_all_optional_function_params(self):
|
||||
"""Test function call with all optional parameters - verifies flexible ordering."""
|
||||
self.__class__.ebnf_grammar = """
|
||||
root ::= function_call
|
||||
function_call ::= call_config_service
|
||||
call_config_service ::= "{" "\\"name\\"" ":" "\\"config_service\\"" ", " "\\"arguments\\"" ":" arguments_config_service "}"
|
||||
arguments_config_service ::= "{" ( "\\"theme\\"" ":" ("\\"light\\"" | "\\"dark\\"") ( "," "\\"language\\"" ":" ("\\"en\\"" | "\\"es\\"" | "\\"fr\\"") )? ( "," "\\"notifications\\"" ":" ("true" | "false") )? | "\\"language\\"" ":" ("\\"en\\"" | "\\"es\\"" | "\\"fr\\"") ( "," "\\"notifications\\"" ":" ("true" | "false") )? | "\\"notifications\\"" ":" ("true" | "false") )? "}"
|
||||
"""
|
||||
# Test patterns that should match - flexible ordering of optional parameters
|
||||
allowed_patterns = [
|
||||
# Empty arguments
|
||||
r'^\{"name":"config_service",\s*"arguments":\{\}\}$',
|
||||
# Single optional parameters (any can appear first)
|
||||
r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)"\}\}$',
|
||||
r'^\{"name":"config_service",\s*"arguments":\{"language":"(en|es|fr)"\}\}$',
|
||||
r'^\{"name":"config_service",\s*"arguments":\{"notifications":(true|false)\}\}$',
|
||||
# Two optional parameters (in any order)
|
||||
r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)",\s*"language":"(en|es|fr)"\}\}$',
|
||||
r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)",\s*"notifications":(true|false)\}\}$',
|
||||
r'^\{"name":"config_service",\s*"arguments":\{"language":"(en|es|fr)",\s*"notifications":(true|false)\}\}$',
|
||||
# All three optional parameters
|
||||
r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)",\s*"language":"(en|es|fr)",\s*"notifications":(true|false)\}\}$',
|
||||
]
|
||||
prompt = "Configure the service with optional settings:"
|
||||
|
||||
self._run_decode_ebnf(
|
||||
ebnf=self.__class__.ebnf_grammar,
|
||||
expected_patterns=allowed_patterns,
|
||||
prompt=prompt,
|
||||
n=5,
|
||||
)
|
||||
97
python/sglang/test/kits/json_constrained_kit.py
Normal file
97
python/sglang/test/kits/json_constrained_kit.py
Normal file
@@ -0,0 +1,97 @@
|
||||
import json
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import openai
|
||||
import requests
|
||||
|
||||
|
||||
class TestJSONConstrainedMixin:
|
||||
json_schema = json.dumps(
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "pattern": "^[\\w]+$"},
|
||||
"population": {"type": "integer"},
|
||||
},
|
||||
"required": ["name", "population"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
)
|
||||
|
||||
def _run_decode_json(
|
||||
self, json_schema, return_logprob=False, top_logprobs_num=0, n=1
|
||||
):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0 if n == 1 else 0.5,
|
||||
"max_new_tokens": 128,
|
||||
"n": n,
|
||||
"stop_token_ids": [119690],
|
||||
"json_schema": json_schema,
|
||||
},
|
||||
"stream": False,
|
||||
"return_logprob": return_logprob,
|
||||
"top_logprobs_num": top_logprobs_num,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
ret = response.json()
|
||||
print(json.dumps(ret))
|
||||
print("=" * 100)
|
||||
|
||||
if not json_schema or json_schema == "INVALID":
|
||||
return
|
||||
|
||||
# Make sure the json output is valid
|
||||
try:
|
||||
js_obj = json.loads(ret["text"])
|
||||
except (TypeError, json.decoder.JSONDecodeError):
|
||||
raise
|
||||
|
||||
self.assertIsInstance(js_obj["name"], str)
|
||||
self.assertIsInstance(js_obj["population"], int)
|
||||
|
||||
def test_json_generate(self):
|
||||
self._run_decode_json(json_schema=self.json_schema)
|
||||
|
||||
def test_json_invalid(self):
|
||||
self._run_decode_json(json_schema="INVALID")
|
||||
|
||||
def test_json_openai(self):
|
||||
client = openai.Client(api_key="EMPTY", base_url=f"{self.base_url}/v1")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Introduce the capital of France. Return in a JSON format.",
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=128,
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "foo", "schema": json.loads(self.json_schema)},
|
||||
},
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
|
||||
try:
|
||||
js_obj = json.loads(text)
|
||||
except (TypeError, json.decoder.JSONDecodeError):
|
||||
print("JSONDecodeError", text)
|
||||
raise
|
||||
|
||||
self.assertIsInstance(js_obj["name"], str)
|
||||
self.assertIsInstance(js_obj["population"], int)
|
||||
|
||||
def test_mix_json_and_other(self):
|
||||
json_schemas = [None, None, self.json_schema, self.json_schema] * 10
|
||||
|
||||
with ThreadPoolExecutor(len(json_schemas)) as executor:
|
||||
list(executor.map(self._run_decode_json, json_schemas))
|
||||
130
python/sglang/test/kits/regex_constrained_kit.py
Normal file
130
python/sglang/test/kits/regex_constrained_kit.py
Normal file
@@ -0,0 +1,130 @@
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
class TestRegexConstrainedMixin:
|
||||
def _run_decode_regex(
|
||||
self,
|
||||
regex,
|
||||
prompt,
|
||||
return_logprob=False,
|
||||
top_logprobs_num=0,
|
||||
n=1,
|
||||
):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0 if n == 1 else 0.5,
|
||||
"max_new_tokens": 128,
|
||||
"n": n,
|
||||
"regex": regex,
|
||||
},
|
||||
"stream": False,
|
||||
"return_logprob": return_logprob,
|
||||
"top_logprobs_num": top_logprobs_num,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
|
||||
ret = response.json()
|
||||
print(json.dumps(ret, indent=2))
|
||||
print("=" * 100)
|
||||
|
||||
if not isinstance(ret, list):
|
||||
self.fail(f"Expected response to be a list, but got {type(ret)}")
|
||||
|
||||
for item in ret:
|
||||
text = item.get("text", "").strip()
|
||||
if not text:
|
||||
self.fail("Generated text is empty.")
|
||||
|
||||
if not self.regex_match(text, regex):
|
||||
self.fail(f"Text '{text}' does not match regex pattern.")
|
||||
|
||||
def regex_match(self, text, pattern):
|
||||
import re
|
||||
|
||||
return re.match(pattern, text) is not None
|
||||
|
||||
def test_regex_generate_email(self):
|
||||
pattern = r"^user@example\.com$"
|
||||
prompt = "Generate an email address:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_regex_generate_greeting(self):
|
||||
pattern = r"^(Hello|Hi|Hey)$"
|
||||
prompt = "Generate a greeting:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_regex_generate_number(self):
|
||||
pattern = r"^\d{3}$"
|
||||
prompt = "Generate a three-digit number:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_regex_generate_phone(self):
|
||||
pattern = r"^\(\d{3}\) \d{3}-\d{4}$"
|
||||
prompt = "Generate a phone number:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_regex_generate_date(self):
|
||||
pattern = r"^2024-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$"
|
||||
prompt = "Generate a date in YYYY-MM-DD format:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_regex_generate_hex_color(self):
|
||||
pattern = r"^#[0-9A-F]{6}$"
|
||||
prompt = "Generate a hex color code:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_regex_generate_complex_json(self):
|
||||
pattern = r'^\{\s*"name"\s*:\s*"[a-zA-Z0-9 ]+"\s*,\s*"age"\s*:\s*[1-9][0-9]*\s*,\s*"city"\s*:\s*"[a-zA-Z0-9 ]+"\s*\}$'
|
||||
prompt = "Generate a simple JSON with name, age, and city:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
|
||||
def test_regex_generate_custom_log_format(self):
|
||||
pattern = r"^\[2024-01-01T12:00:00Z\] INFO: System\.process - Operation [a-z]+ successfully$"
|
||||
prompt = "Generate a log entry:"
|
||||
|
||||
self._run_decode_regex(
|
||||
regex=pattern,
|
||||
prompt=prompt,
|
||||
n=3,
|
||||
)
|
||||
Reference in New Issue
Block a user