Allow setting top_p to 1.0 (mlc-ai#354)

abi · Apr 1, 2024 · d20e08c · d20e08c
1 parent e4e2903
commit d20e08c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/config.ts b/src/config.ts
@@ -121,8 +121,8 @@ export function postInitAndCheckGenerationConfigValues(config: GenerationConfig)
   if (_hasValue(config.shift_fill_factor) && config.shift_fill_factor! <= 0 || config.shift_fill_factor! > 1) {
     throw new Error("Make sure 0 < `shift_fill_factor` <= 1.");
   }
-  if (_hasValue(config.top_p) && config.top_p! <= 0 || config.top_p! >= 1) {
-    throw new Error("Make sure 0 < `top_p` < 1.");
+  if (_hasValue(config.top_p) && config.top_p! <= 0 || config.top_p! > 1) {
+    throw new Error("Make sure 0 < `top_p` <= 1.");
   }
   if (_hasValue(config.temperature) && config.temperature! < 0) {
     throw new Error("Make sure `temperature` >= 0.");

diff --git a/src/llm_chat.ts b/src/llm_chat.ts
@@ -581,7 +581,7 @@ export class LLMChatPipeline {
       if (_hasValue(genConfig.top_logprobs)) { top_logprobs = genConfig.top_logprobs!; }
     }
     // Check range validity
-    if (top_p <= 0 || top_p >= 1) { throw new Error("Make sure 0 < `top_p` < 1."); }
+    if (top_p <= 0 || top_p > 1) { throw new Error("Make sure 0 < `top_p` <= 1."); }
     if (temperature < 0) { throw new Error("Make sure `temperature` >= 0."); }
     if (repetition_penalty <= 0) { throw new Error("Make sure `repetition_penalty` > 0."); }
     if (frequency_penalty && (frequency_penalty < -2.0 || frequency_penalty > 2.0)) {