EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Support stop_token_ids in sglang API (#1092)

Browse Source
This commit is contained in:
Liangsheng Yin
2024-08-14 17:31:39 -07:00
committed by GitHub
parent 1c2b5f5240
commit 73cf6834f2
6 changed files with 34 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/api.py
View File

@@ -62,6 +62,7 @@ def gen(
name: Optional[str] = None,
max_tokens: Optional[int] = None,
stop: Optional[Union[str, List[str]]] = None,
stop_token_ids: Optional[List[int]] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
@@ -98,6 +99,7 @@ def gen(
name,
max_tokens,
stop,
stop_token_ids,
temperature,
top_p,
top_k,
@@ -117,6 +119,7 @@ def gen_int(
name: Optional[str] = None,
max_tokens: Optional[int] = None,
stop: Optional[Union[str, List[str]]] = None,
stop_token_ids: Optional[List[int]] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
@@ -132,6 +135,7 @@ def gen_int(
name,
max_tokens,
stop,
stop_token_ids,
temperature,
top_p,
top_k,
@@ -151,6 +155,7 @@ def gen_string(
name: Optional[str] = None,
max_tokens: Optional[int] = None,
stop: Optional[Union[str, List[str]]] = None,
stop_token_ids: Optional[List[int]] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
@@ -166,6 +171,7 @@ def gen_string(
name,
max_tokens,
stop,
stop_token_ids,
temperature,
top_p,
top_k,

6
python/sglang/lang/interpreter.py
View File

@@ -20,7 +20,6 @@ from sglang.lang.ir import (
SglConstantText,
SglExpr,
SglExprList,
SglFunction,
SglGen,
SglImage,
SglRoleBegin,
@@ -181,8 +180,10 @@ class StreamExecutor:
num_api_spec_tokens=None,
use_thread=True,
):
from sglang.lang.backend.base_backend import BaseBackend
self.sid = uuid.uuid4().hex
self.backend = backend
self.backend: BaseBackend = backend
self.arguments: Dict[str, Any] = arguments
self.default_sampling_para = default_sampling_para
self.stream = stream
@@ -658,6 +659,7 @@ class StreamExecutor:
for item in [
"max_new_tokens",
"stop",
"stop_token_ids",
"temperature",
"top_p",
"top_k",

11
python/sglang/lang/ir.py
View File

@@ -18,6 +18,7 @@ REGEX_STR = r"\"[\w\d\s]*\"" # bugs with regex r"\".*\"" in interegular pkg
class SglSamplingParams:
max_new_tokens: int = 128
stop: Union[str, List[str]] = ()
stop_token_ids: Optional[List[int]] = ()
temperature: float = 1.0
top_p: float = 1.0
top_k: int = -1 # -1 means disable
@@ -37,6 +38,7 @@ class SglSamplingParams:
return SglSamplingParams(
self.max_new_tokens,
self.stop,
self.stop_token_ids,
self.temperature,
self.top_p,
self.top_k,
@@ -108,6 +110,7 @@ class SglSamplingParams:
return {
"max_new_tokens": self.max_new_tokens,
"stop": self.stop,
"stop_token_ids": self.stop_token_ids,
"temperature": self.temperature,
"top_p": self.top_p,
"top_k": self.top_k,
@@ -141,7 +144,8 @@ class SglFunction:
self,
*args,
max_new_tokens: int = 128,
stop: Union[str, List[str]] = (),
stop: Union[str, List[str]] = [],
stop_token_ids: Optional[List[int]] = [],
temperature: float = 1.0,
top_p: float = 1.0,
top_k: int = -1,
@@ -161,6 +165,7 @@ class SglFunction:
default_sampling_para = SglSamplingParams(
max_new_tokens=max_new_tokens,
stop=stop,
stop_token_ids=stop_token_ids,
temperature=temperature,
top_p=top_p,
top_k=top_k,
@@ -181,6 +186,7 @@ class SglFunction:
*,
max_new_tokens: int = 128,
stop: Union[str, List[str]] = (),
stop_token_ids: Optional[List[int]] = [],
temperature: float = 1.0,
top_p: float = 1.0,
top_k: int = -1,
@@ -218,6 +224,7 @@ class SglFunction:
default_sampling_para = SglSamplingParams(
max_new_tokens=max_new_tokens,
stop=stop,
stop_token_ids=stop_token_ids,
temperature=temperature,
top_p=top_p,
top_k=top_k,
@@ -397,6 +404,7 @@ class SglGen(SglExpr):
name: Optional[str] = None,
max_new_tokens: Optional[int] = None,
stop: Optional[Union[str, List[str]]] = None,
stop_token_ids: Optional[List[int]] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
@@ -416,6 +424,7 @@ class SglGen(SglExpr):
self.sampling_params = SglSamplingParams(
max_new_tokens=max_new_tokens,
stop=stop,
stop_token_ids=stop_token_ids,
temperature=temperature,
top_p=top_p,
top_k=top_k,

10
python/sglang/srt/managers/schedule_batch.py
View File

@@ -235,10 +235,12 @@ class Req:
return
last_token_id = self.output_ids[-1]
if self.tokenizer is None:
matched_eos = last_token_id in self.sampling_params.stop_token_ids
else:
matched_eos = last_token_id == self.tokenizer.eos_token_id
matched_eos = last_token_id in self.sampling_params.stop_token_ids
if self.tokenizer is not None:
matched_eos |= last_token_id == self.tokenizer.eos_token_id
if matched_eos and not self.sampling_params.ignore_eos:
self.finished_reason = FINISH_MATCHED_TOKEN(matched=last_token_id)
return

11
python/sglang/test/test_programs.py
View File

@@ -106,13 +106,16 @@ def test_decode_json_regex():
from sglang.lang.ir import REGEX_FLOAT, REGEX_INT, REGEX_STR
s += "Generate a JSON object to describe the basic city information of Paris.\n"
s += "Here are the JSON object:\n"
# NOTE: we recommend using dtype gen or whole regex string to control the output
with s.var_scope("json_output"):
s += "{\n"
s += ' "name": ' + sgl.gen(regex=REGEX_STR + ",") + "\n"
s += ' "population": ' + sgl.gen(regex=REGEX_INT + ",") + "\n"
s += ' "area": ' + sgl.gen(regex=REGEX_INT + ",") + "\n"
s += ' "latitude": ' + sgl.gen(regex=REGEX_FLOAT) + "\n"
s += ' "name": ' + sgl.gen(regex=REGEX_STR) + ",\n"
s += ' "population": ' + sgl.gen(regex=REGEX_INT, stop=[" ", "\n"]) + ",\n"
s += ' "area": ' + sgl.gen(regex=REGEX_INT, stop=[" ", "\n"]) + ",\n"
s += ' "latitude": ' + sgl.gen(regex=REGEX_FLOAT, stop=[" ", "\n"]) + "\n"
s += "}"
ret = decode_json.run(temperature=0.0)