wehub-resource-sync
77 lines
3.8 KiB
Python
77 lines
3.8 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional
|
|
|
|
from swift.infer_engine import RequestConfig
|
|
from swift.utils import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class GenerationArguments:
|
|
"""A dataclass that holds arguments for text generation.
|
|
|
|
Args:
|
|
max_new_tokens (Optional[int]): The maximum number of new tokens to generate. Defaults to None (unlimited).
|
|
temperature (Optional[float]): The sampling temperature. A higher temperature makes the output more random. To
|
|
disable randomness, you can set this to 0 or `top_k` to 1. Defaults to None, which means loading from
|
|
'generation_config.json'.
|
|
top_k (Optional[int]): The number of highest probability vocabulary tokens to keep for top-k-filtering.
|
|
Defaults to None (reads from 'generation_config.json').
|
|
top_p (Optional[float]): The cumulative probability for nucleus sampling. Filters the vocabulary to the
|
|
smallest set of tokens whose cumulative probability exceeds `top_p`. Defaults to None (reads from
|
|
'generation_config.json').
|
|
repetition_penalty (Optional[float]): The penalty applied to repeated tokens. A value of 1.0 means no penalty.
|
|
Defaults to None (reads from 'generation_config.json').
|
|
num_beams (Optional[int]): The number of beams to use for beam search. Defaults to 1.
|
|
stream (bool): Whether to enable streaming output. Defaults to None, which is `True` for interactive mode and
|
|
`False` for batch inference. Note: For ms-swift < 3.6, the default is `False`.
|
|
stop_words (List[str]): A list of extra stop words, in addition to the end-of-sequence token. Note: The
|
|
`eos_token` is removed from the output, while these stop words are preserved. Defaults to an empty list.
|
|
logprobs (bool): Whether to output log probabilities of the generated tokens. Defaults to False.
|
|
top_logprobs (Optional[int]): The number of top log probabilities to return for each token position. Requires
|
|
`logprobs` to be True. Defaults to None.
|
|
structured_outputs_regex (Optional[str]): A regular expression pattern for structured outputs (guided decoding).
|
|
When set, the model's generation is constrained to match the specified regex pattern. This is useful for
|
|
tasks requiring structured outputs like reasoning chains. Only effective when `infer_backend` is 'vllm'.
|
|
Defaults to None.
|
|
"""
|
|
|
|
# generation config
|
|
max_new_tokens: Optional[int] = None # Unlimited, constrained by max_model_len.
|
|
# If it is None, use the parameters from generation_config.
|
|
temperature: Optional[float] = None # Set to 0, which means do_sample is False.
|
|
top_k: Optional[int] = None
|
|
top_p: Optional[float] = None
|
|
repetition_penalty: Optional[float] = None
|
|
num_beams: int = 1
|
|
|
|
stream: Optional[bool] = None
|
|
stop_words: List[str] = field(default_factory=list)
|
|
logprobs: bool = False
|
|
top_logprobs: Optional[int] = None
|
|
# structured outputs (guided decoding), only effective for vllm backend
|
|
structured_outputs_regex: Optional[str] = None
|
|
|
|
def _init_stream(self):
|
|
if self.stream is None:
|
|
self.stream = False
|
|
|
|
def get_request_config(self):
|
|
if getattr(self, 'task_type') != 'causal_lm':
|
|
return
|
|
|
|
return RequestConfig(
|
|
max_tokens=self.max_new_tokens,
|
|
temperature=self.temperature,
|
|
top_p=self.top_p,
|
|
top_k=self.top_k,
|
|
num_beams=self.num_beams,
|
|
stop=self.stop_words,
|
|
stream=self.stream,
|
|
repetition_penalty=self.repetition_penalty,
|
|
logprobs=self.logprobs,
|
|
top_logprobs=self.top_logprobs,
|
|
structured_outputs_regex=self.structured_outputs_regex)
|