File size: 2,170 Bytes
cb2428f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
# Copyright (c) Alibaba, Inc. and its affiliates.
from dataclasses import dataclass, field
from typing import List, Optional
from swift.utils import get_logger
logger = get_logger()
@dataclass
class GenerationArguments:
"""
GenerationArguments class is a dataclass that holds various arguments related to text generation.
Args:
max_new_tokens (Optional[int]): Maximum number of new tokens to generate. Default is None (unlimited).
temperature (Optional[float]): Sampling temperature. Default is None.
top_k (Optional[int]): Top-k sampling parameter. Default is None.
top_p (Optional[float]): Top-p (nucleus) sampling parameter. Default is None.
repetition_penalty (Optional[float]): Penalty for repeated tokens. Default is None.
num_beams (int): Number of beams for beam search. Default is 1.
stream (bool): Flag to indicate if streaming output should be enabled. Default is None.
stop_words (List[str]): List of stop words to end generation. Default is an empty list.
"""
# 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: bool = False
stop_words: List[str] = field(default_factory=list)
logprobs: bool = False
top_logprobs: Optional[int] = None
def get_request_config(self):
if getattr(self, 'task_type') != 'causal_lm':
return
from swift.llm import RequestConfig
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)
|