From 017f3de464682b7dfd96da7274adb907ccbcab77 Mon Sep 17 00:00:00 2001
From: Chirag Jain <jain.chirag925@gmail.com>
Date: Fri, 6 Oct 2023 18:22:12 +0530
Subject: [PATCH] Add ignore_eos_token to HTTP interface

---
 router/src/lib.rs        | 4 ++++
 router/src/validation.rs | 3 ++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/router/src/lib.rs b/router/src/lib.rs
index 560b8f74..3b7bbac4 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -138,6 +138,9 @@ pub(crate) struct GenerateParameters {
     #[serde(default)]
     #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 5)]
     pub top_n_tokens: Option<u32>,
+    #[serde(default)]
+    #[schema(default = "false")]
+    pub ignore_eos_token: bool
 }
 
 fn default_max_new_tokens() -> Option<u32> {
@@ -162,6 +165,7 @@ fn default_parameters() -> GenerateParameters {
         decoder_input_details: false,
         seed: None,
         top_n_tokens: None,
+        ignore_eos_token: false,
     }
 }
 
diff --git a/router/src/validation.rs b/router/src/validation.rs
index 9adedc5b..2d6b878a 100644
--- a/router/src/validation.rs
+++ b/router/src/validation.rs
@@ -156,6 +156,7 @@ impl Validation {
             watermark,
             decoder_input_details,
             top_n_tokens,
+            ignore_eos_token,
             ..
         } = request.parameters;
 
@@ -274,7 +275,7 @@ impl Validation {
         let stopping_parameters = StoppingCriteriaParameters {
             max_new_tokens,
             stop_sequences,
-            ignore_eos_token: false,
+            ignore_eos_token: ignore_eos_token,
         };
 
         metrics::histogram!("tgi_request_max_new_tokens", max_new_tokens as f64);