mirror/OpenLLM
1
0
Fork 0
mirror of https://github.com/bentoml/OpenLLM.git synced 2026-01-06 06:29:21 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat(generation): add support for eos_token_id (#714)

Browse Source 
chore: add support for custom eos_token_id

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
This commit is contained in:
Aaron Pham
2023-11-21 02:01:36 -05:00
committed by GitHub
parent fde78a2c78
commit e70246ca5d
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
openllm-python/src/openllm/_llm.py
View File

@@ -454,6 +454,11 @@ class LLM(t.Generic[M, T], ReprMixin):
if stop_token_ids is None:
stop_token_ids = []
eos_token_id = attrs.get('eos_token_id', config['eos_token_id'])
if eos_token_id is not None:
if not isinstance(eos_token_id, list):
eos_token_id = [eos_token_id]
stop_token_ids.extend(eos_token_id)
if self.tokenizer.eos_token_id not in stop_token_ids:
stop_token_ids.append(self.tokenizer.eos_token_id)
if stop is None: