[Minor] Fix type annotations (#1238)
This commit is contained in:
parent
a60b353005
commit
84e4e37d14
@ -290,7 +290,7 @@ def _get_topk_logprobs(
|
|||||||
def _build_sequence_outputs(
|
def _build_sequence_outputs(
|
||||||
parent_ids: List[int],
|
parent_ids: List[int],
|
||||||
next_token_ids: List[int],
|
next_token_ids: List[int],
|
||||||
selected_token_logprobs: torch.Tensor,
|
selected_token_logprobs: List[float],
|
||||||
parent_seq_ids: List[int],
|
parent_seq_ids: List[int],
|
||||||
parent_logprobs: torch.Tensor,
|
parent_logprobs: torch.Tensor,
|
||||||
num_output_logprobs: Optional[int],
|
num_output_logprobs: Optional[int],
|
||||||
|
|||||||
@ -61,7 +61,6 @@ class SamplingParams:
|
|||||||
max_tokens: Maximum number of tokens to generate per output sequence.
|
max_tokens: Maximum number of tokens to generate per output sequence.
|
||||||
logprobs: Number of log probabilities to return per output token.
|
logprobs: Number of log probabilities to return per output token.
|
||||||
skip_special_tokens: Whether to skip special tokens in the output.
|
skip_special_tokens: Whether to skip special tokens in the output.
|
||||||
Defaults to true.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@ -76,8 +75,8 @@ class SamplingParams:
|
|||||||
use_beam_search: bool = False,
|
use_beam_search: bool = False,
|
||||||
length_penalty: float = 1.0,
|
length_penalty: float = 1.0,
|
||||||
early_stopping: Union[bool, str] = False,
|
early_stopping: Union[bool, str] = False,
|
||||||
stop: Union[None, str, List[str]] = None,
|
stop: Optional[Union[str, List[str]]] = None,
|
||||||
stop_token_ids: List[int] = None,
|
stop_token_ids: Optional[List[int]] = None,
|
||||||
ignore_eos: bool = False,
|
ignore_eos: bool = False,
|
||||||
max_tokens: int = 16,
|
max_tokens: int = 16,
|
||||||
logprobs: Optional[int] = None,
|
logprobs: Optional[int] = None,
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user