Skip to content

GitLab

Commit 46ce1356 authored by Woosuk Kwon's avatar Woosuk Kwon
Browse files

Add max_num_steps to SamplingParams

parent b39f149a
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
cacheflow/sampling_params.py
View file @ 46ce1356
...@@ -10,6 +10,7 @@ class SamplingParams: ...@@ -10,6 +10,7 @@ class SamplingParams:
top_p: float = 1.0, top_p: float = 1.0,
use_beam_search: bool = False, use_beam_search: bool = False,
stop_token_ids: Set[int] = [], stop_token_ids: Set[int] = [],
max_num_steps: int = 16, # From OpenAI API.
max_context_len: Optional[int] = None, max_context_len: Optional[int] = None,
) -> None: ) -> None:
assert n >= 1 assert n >= 1
...@@ -23,6 +24,7 @@ class SamplingParams: ...@@ -23,6 +24,7 @@ class SamplingParams:
# Zero temperature means greedy decoding. # Zero temperature means greedy decoding.
assert n == 1 assert n == 1
assert top_p == 1.0 assert top_p == 1.0
assert max_num_steps >= 1
assert max_context_len is None or max_context_len >= 0 assert max_context_len is None or max_context_len >= 0
self.n = n self.n = n
...@@ -30,4 +32,5 @@ class SamplingParams: ...@@ -30,4 +32,5 @@ class SamplingParams:
self.top_p = top_p self.top_p = top_p
self.use_beam_search = use_beam_search self.use_beam_search = use_beam_search
self.stop_token_ids = stop_token_ids self.stop_token_ids = stop_token_ids
self.max_num_steps = max_num_steps
self.max_context_len = max_context_len self.max_context_len = max_context_len
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