"add support for vllm api stop parameter"

Former-commit-id: b9f21fa639b66db09c79404d885661c96bdf9395
This commit is contained in:
zhaonx
2024-04-30 17:17:09 +08:00
parent 3cef844079
commit 2d95127c33
4 changed files with 12 additions and 3 deletions

View File

@@ -1,5 +1,5 @@
from dataclasses import asdict, dataclass, field
from typing import Any, Dict
from typing import Any, Dict, Union, Optional, List
@dataclass
@@ -46,7 +46,10 @@ class GeneratingArguments:
default=1.0,
metadata={"help": "Exponential penalty to the length that is used with beam-based generation."},
)
stop: Union[Optional[str], List[str]] = field(
default=None,
metadata={"help": "List of strings or string that stop the generation when they are generated. The returned output will not contain the stop strings."},
)
def to_dict(self) -> Dict[str, Any]:
args = asdict(self)
if args.get("max_new_tokens", -1) > 0: