Unverified Commit 84e4e37d authored by Woosuk Kwon's avatar Woosuk Kwon Committed by GitHub
Browse files

[Minor] Fix type annotations (#1238)

parent a60b3530
...@@ -290,7 +290,7 @@ def _get_topk_logprobs( ...@@ -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: ...@@ -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: ...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment