huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-24 16:32:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Ignore EOS for benchmark by using TEXT_GENERATION_SERVER_IGNORE_EOS_TOKEN (#234)

Browse Source
This commit is contained in:
Sun Choi 2024-10-16 02:57:36 -07:00 committed by GitHub
parent d07e7f4f62
commit 8ae5d4c7d6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
server/text_generation_server/utils/tokens.py
View File

@ -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