Fix for token streaming not visually showing sentence trimming correctly

This commit is contained in:
ebolam
2022-09-01 09:27:55 -04:00
parent 5ee620db30
commit 08fac1157a
2 changed files with 7 additions and 4 deletions

View File

@@ -1828,7 +1828,7 @@ def patch_transformers():
#koboldai_vars.actions.stream_token(tokenizer_text, batch=batch)
if koboldai_vars.output_streaming:
koboldai_vars.actions.stream_tokens([utils.decodenewlines(tokenizer.decode(x[-1])) for x in input_ids])
koboldai_vars.actions.stream_tokens([utils.decodenewlines(tokenizer.decode(x[-1])) for x in input_ids], koboldai_vars.genamt)
#if len(input_ids) > 1:
# koboldai_vars.actions.clear_unused_options()
# koboldai_vars.actions.append_options([utils.decodenewlines(tokenizer.decode(x[-1])) for x in input_ids])

View File

@@ -1053,7 +1053,7 @@ class KoboldStoryRegister(object):
process_variable_changes(self.socketio, "story", 'actions', {"id": key, 'action': self.actions[key]}, None)
ignore = self.koboldai_vars.calc_ai_text()
def stream_tokens(self, text_list):
def stream_tokens(self, text_list, max_tokens):
if len(text_list) > 1:
if self.action_count+1 in self.actions:
for i in range(len(text_list)):
@@ -1070,7 +1070,10 @@ class KoboldStoryRegister(object):
for i in range(len(text_list)):
self.actions[self.action_count+1]['Options'].append({"text": text_list[i], "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": [], "stream_id": i})
process_variable_changes(self.socketio, "actions", "Options", {"id": self.action_count+1, "options": self.actions[self.action_count+1]["Options"]}, {"id": self.action_count+1, "options": None})
#We need to see if this is the last token being streamed. If so due to the rely it will come in AFTER the actual trimmed final text overwriting it in the UI
if self.tokenizer is not None:
if len(self.tokenizer.encode(self.actions[self.action_count+1]["Options"][0]['text'])) != max_tokens:
#process_variable_changes(self.socketio, "actions", "Options", {"id": self.action_count+1, "options": self.actions[self.action_count+1]["Options"]}, {"id": self.action_count+1, "options": None})
process_variable_changes(self.socketio, "story", 'actions', {"id": self.action_count+1, 'action': self.actions[self.action_count+1]}, None)
else:
#We're streaming single options so our output is our selected