From 8ae5d4c7d6912bf9470a4b9b5330ef6c3b8dc0ae Mon Sep 17 00:00:00 2001 From: Sun Choi Date: Wed, 16 Oct 2024 02:57:36 -0700 Subject: [PATCH] Ignore EOS for benchmark by using TEXT_GENERATION_SERVER_IGNORE_EOS_TOKEN (#234) --- server/text_generation_server/utils/tokens.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/server/text_generation_server/utils/tokens.py b/server/text_generation_server/utils/tokens.py index 267ea068..d2dca8b9 100644 --- a/server/text_generation_server/utils/tokens.py +++ b/server/text_generation_server/utils/tokens.py @@ -1,5 +1,6 @@ # Copyright (C) 2024 Habana Labs, Ltd. an Intel Company. +import os import re from typing import List, Optional, Tuple, Set, Union @@ -164,7 +165,11 @@ class StoppingCriteria: self.max_new_tokens = max_new_tokens self.current_tokens = 0 self.current_output = "" - self.ignore_eos_token = ignore_eos_token + + if os.getenv("TEXT_GENERATION_SERVER_IGNORE_EOS_TOKEN", "false") == "true": + self.ignore_eos_token = True + else: + self.ignore_eos_token = ignore_eos_token def __call__(self, last_token: int, last_output: str) -> Tuple[bool, Optional[str]]: self.current_tokens += 1