From f82da93318b5b2703a57b88d98acec79af38f052 Mon Sep 17 00:00:00 2001
From: Karol Damaszke <karol.damaszke@intel.com>
Date: Mon, 6 May 2024 09:55:58 +0200
Subject: [PATCH] Fix input length validation (#135)

Co-authored-by: Karol Damaszke <kdamaszke@habana.ai>
---
 router/src/validation.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/router/src/validation.rs b/router/src/validation.rs
index 4f3036e8..012a1b73 100644
--- a/router/src/validation.rs
+++ b/router/src/validation.rs
@@ -6,7 +6,7 @@ use crate::validation::ValidationError::{BestOfSampling, BestOfSeed, EmptyInput}
 use crate::{GenerateParameters, GenerateRequest, GrammarType};
 use jsonschema::{Draft, JSONSchema};
 use rand::{thread_rng, Rng};
-use std::env;
+use std::{cmp, env};
 use serde_json::Value;
 use std::io::Cursor;
 use text_generation_client::{
@@ -131,7 +131,10 @@ impl Validation {
             let input_length = if self.skip_tokenizer_in_tgi {
                 inputs.chars().filter(|&c| c == ',').count() + 1
             } else {
-                truncate.unwrap_or(self.max_input_length)
+                cmp::max(
+                    encoding.len(),
+                    truncate.unwrap_or(self.max_input_length)
+                )
             };
 
             // Get total tokens