From ecbc0a0f4a0cbe95aa8f7fc46949a5fd1e2d852c Mon Sep 17 00:00:00 2001
From: Chirag Jain <jain.chirag925@gmail.com>
Date: Mon, 21 Aug 2023 14:06:17 +0530
Subject: [PATCH] Add ignore_eos_token

---
 Dockerfile-compile-router | 36 ++++++++++++++++++++++++++++++++++++
 router/src/lib.rs         |  4 ++++
 router/src/validation.rs  |  3 ++-
 3 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 Dockerfile-compile-router

diff --git a/Dockerfile-compile-router b/Dockerfile-compile-router
new file mode 100644
index 00000000..b435e208
--- /dev/null
+++ b/Dockerfile-compile-router
@@ -0,0 +1,36 @@
+# Rust builder
+FROM --platform=linux/amd64 lukemathwalker/cargo-chef:latest-rust-1.69 AS chef
+WORKDIR /usr/src
+
+ARG CARGO_REGISTRIES_CRATES_IO_PROTOCOL=sparse
+
+FROM chef as planner
+COPY Cargo.toml Cargo.toml
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY proto proto
+COPY benchmark benchmark
+COPY router router
+COPY launcher launcher
+RUN cargo chef prepare --recipe-path recipe.json
+
+FROM --platform=linux/amd64 chef AS builder
+
+ARG GIT_SHA
+ARG DOCKER_LABEL
+
+RUN PROTOC_ZIP=protoc-21.12-linux-x86_64.zip && \
+curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v21.12/$PROTOC_ZIP && \
+unzip -o $PROTOC_ZIP -d /usr/local bin/protoc && \
+unzip -o $PROTOC_ZIP -d /usr/local 'include/*' && \
+rm -f $PROTOC_ZIP
+
+COPY --from=planner /usr/src/recipe.json recipe.json
+RUN cargo chef cook --release --recipe-path recipe.json
+
+COPY Cargo.toml Cargo.toml
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY proto proto
+COPY benchmark benchmark
+COPY router router
+COPY launcher launcher
+RUN cargo build --release
diff --git a/router/src/lib.rs b/router/src/lib.rs
index 7dff7a11..f2309d52 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -135,6 +135,9 @@ pub(crate) struct GenerateParameters {
         example = "null"
     )]
     pub seed: Option<u64>,
+    #[serde(default)]
+    #[schema(default = "false")]
+    pub ignore_eos_token: bool
 }
 
 fn default_max_new_tokens() -> u32 {
@@ -158,6 +161,7 @@ fn default_parameters() -> GenerateParameters {
         details: false,
         decoder_input_details: false,
         seed: None,
+        ignore_eos_token: false,
     }
 }
 
diff --git a/router/src/validation.rs b/router/src/validation.rs
index be835bf0..697b339a 100644
--- a/router/src/validation.rs
+++ b/router/src/validation.rs
@@ -142,6 +142,7 @@ impl Validation {
             seed,
             watermark,
             decoder_input_details,
+            ignore_eos_token,
             ..
         } = request.parameters;
 
@@ -251,7 +252,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);