fix: usize overflow issue in ctranslate2 max_length truncation

2023-09-12 20:56:35 +08:00 · 2023-09-12 20:56:35 +08:00 · 1e7ecce697
parent 3af32c8c7c
commit 1e7ecce697
1 changed files with 6 additions and 2 deletions
--- a/crates/ctranslate2-bindings/src/lib.rs
+++ b/crates/ctranslate2-bindings/src/lib.rs
@ -136,8 +136,12 @@ impl TextGeneration for CTranslate2Engine {
 }

 fn truncate_tokens(tokens: &[String], max_length: usize) -> &[String] {
-    let start = std::cmp::max(tokens.len() - max_length, 0);
-    &tokens[start..]
+    if max_length < tokens.len() {
+        let start = tokens.len() - max_length;
+        &tokens[start..]
+    } else {
+        tokens
+    }
 }

 fn inference_callback(