Better stop token handling
This commit is contained in:
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user