llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Better stop token handling

Browse Source
This commit is contained in:
Niels Geens
2026-01-10 13:41:32 +01:00
parent b89cb9cfe1
commit 39e0c81ef1
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
View File

@@ -41,8 +41,14 @@ class GemmaLlmEngine(LlmEngine):
logits_processor=logits_processor_list logits_processor=logits_processor_list
) )
for output in stream: for output in stream:
if 'text' in output["choices"][0]: choice = output["choices"][0]
yield output["choices"][0]['text'] if choice.get('finish_reason'):
break
if 'text' in choice:
text = choice['text']
if text == '':
break
yield text
def token_count(self, system: str, context: str) -> int: def token_count(self, system: str, context: str) -> int:
return len(self._format_messages(system, context, None)) return len(self._format_messages(system, context, None))