From 4267378b1f10b61058068f3847bfbf89f9a6c8a6 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Sun, 9 Apr 2023 10:06:53 +0200
Subject: [PATCH] fix validation error

---
 router/src/validation.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/router/src/validation.rs b/router/src/validation.rs
index 01b2de4b..d81a27e7 100644
--- a/router/src/validation.rs
+++ b/router/src/validation.rs
@@ -2,7 +2,6 @@ use crate::validation::ValidationError::{BestOfSampling, BestOfSeed, EmptyInput}
 /// Payload validation logic
 use crate::{GenerateParameters, GenerateRequest};
 use rand::{thread_rng, Rng};
-use std::cmp::max;
 use text_generation_client::{NextTokenChooserParameters, StoppingCriteriaParameters};
 use thiserror::Error;
 use tokenizers::tokenizer::Tokenizer;
@@ -115,7 +114,9 @@ impl Validation {
             // We make sure that truncate + max_new_tokens <= self.max_total_tokens
 
             // Validate MaxNewTokens
-            if (truncate + max_new_tokens) > self.max_total_tokens {
+            if (truncate.unwrap_or(self.max_input_length) as u32 + max_new_tokens)
+                > self.max_total_tokens as u32
+            {
                 return Err(ValidationError::MaxNewTokens(
                     self.max_total_tokens - self.max_input_length,
                     max_new_tokens,