Better stop token handling

This commit is contained in:
2026-01-10 13:41:32 +01:00
parent b89cb9cfe1
commit 39e0c81ef1

View File

@@ -41,8 +41,14 @@ class GemmaLlmEngine(LlmEngine):
logits_processor=logits_processor_list
)
for output in stream:
if 'text' in output["choices"][0]:
yield output["choices"][0]['text']
choice = output["choices"][0]
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:
return len(self._format_messages(system, context, None))