[Feature] Support new parameter - EBNF in xgrammar (#2526)

This commit is contained in:
Adarsh Shirawalmath
2024-12-26 18:42:41 +05:30
committed by GitHub
parent 08effbff35
commit acb340728c
8 changed files with 384 additions and 2 deletions

View File

@@ -517,6 +517,7 @@ def v1_generate_request(
"repetition_penalty": request.repetition_penalty,
"regex": request.regex,
"json_schema": request.json_schema,
"ebnf": request.ebnf,
"n": request.n,
"no_stop_trim": request.no_stop_trim,
"ignore_eos": request.ignore_eos,
@@ -692,6 +693,14 @@ def v1_generate_response(request, ret, tokenizer_manager, to_file=False):
async def v1_completions(tokenizer_manager, raw_request: Request):
request_json = await raw_request.json()
if "extra_body" in request_json:
extra = request_json["extra_body"]
if "ebnf" in extra:
request_json["ebnf"] = extra["ebnf"]
if "regex" in extra:
request_json["regex"] = extra["regex"]
# remove extra_body to avoid pydantic conflict
del request_json["extra_body"]
all_requests = [CompletionRequest(**request_json)]
adapted_request, request = v1_generate_request(all_requests)
@@ -936,6 +945,7 @@ def v1_chat_generate_request(
"frequency_penalty": request.frequency_penalty,
"repetition_penalty": request.repetition_penalty,
"regex": request.regex,
"ebnf": request.ebnf,
"n": request.n,
"no_stop_trim": request.no_stop_trim,
"ignore_eos": request.ignore_eos,
@@ -1108,6 +1118,15 @@ def v1_chat_generate_response(request, ret, to_file=False, cache_report=False):
async def v1_chat_completions(tokenizer_manager, raw_request: Request):
request_json = await raw_request.json()
if "extra_body" in request_json:
extra = request_json["extra_body"]
# For example, if 'ebnf' is given:
if "ebnf" in extra:
request_json["ebnf"] = extra["ebnf"]
if "regex" in extra:
request_json["regex"] = extra["regex"]
# remove extra_body to avoid pydantic conflict
del request_json["extra_body"]
all_requests = [ChatCompletionRequest(**request_json)]
adapted_request, request = v1_chat_generate_request(all_requests, tokenizer_manager)

View File

@@ -179,6 +179,7 @@ class CompletionRequest(BaseModel):
ignore_eos: bool = False
skip_special_tokens: bool = True
lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None
ebnf: Optional[str] = None
class CompletionResponseChoice(BaseModel):
@@ -288,6 +289,7 @@ class ChatCompletionRequest(BaseModel):
ignore_eos: bool = False
skip_special_tokens: bool = True
lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None
ebnf: Optional[str] = None
class ChatMessage(BaseModel):