From 7ff6c409f7f21b0ce1d31c740c07a1e9b1606a35 Mon Sep 17 00:00:00 2001 From: somebody Date: Wed, 21 Dec 2022 17:23:35 -0600 Subject: [PATCH] Fix probabilities on UI1 --- aiserver.py | 54 +++++++++++++++++++++---------------------- koboldai_settings.py | 10 ++++++-- static/application.js | 1 - 3 files changed, 35 insertions(+), 30 deletions(-) diff --git a/aiserver.py b/aiserver.py index 9b55ae67..bb79fe99 100644 --- a/aiserver.py +++ b/aiserver.py @@ -2297,34 +2297,34 @@ def patch_transformers(): if koboldai_vars.numseqs > 1 or not koboldai_vars.show_probs: return - if not koboldai_vars.show_probs: - return scores - - option_offset = 0 - if koboldai_vars.actions.action_count+1 in koboldai_vars.actions.actions: - for x in range(len(koboldai_vars.actions.actions[koboldai_vars.actions.action_count+1]['Options'])): - option = koboldai_vars.actions.actions[koboldai_vars.actions.action_count+1]['Options'][x] - if option['Pinned'] or option["Previous Selection"] or option["Edited"]: - option_offset = x+1 - batch_offset = int((koboldai_vars.generated_tkns-1) / koboldai_vars.genamt) if koboldai_vars.alt_multi_gen else 0 - for batch_index, batch in enumerate(scores): - probs = F.softmax(batch, dim = -1).cpu().numpy() - - token_prob_info = [] - for token_id, score in sorted(enumerate(probs), key=lambda x: x[1], reverse=True)[:8]: - token_prob_info.append({ - "tokenId": token_id, - "decoded": utils.decodenewlines(tokenizer.decode(token_id)), - "score": float(score), - }) - - - if koboldai_vars.numseqs == 1: - koboldai_vars.actions.set_probabilities(token_prob_info) - else: - koboldai_vars.actions.set_option_probabilities(token_prob_info, batch_index+option_offset+batch_offset) - + if not koboldai_vars.show_probs: return scores + + option_offset = 0 + if koboldai_vars.actions.action_count+1 in koboldai_vars.actions.actions: + for x in range(len(koboldai_vars.actions.actions[koboldai_vars.actions.action_count+1]['Options'])): + option = koboldai_vars.actions.actions[koboldai_vars.actions.action_count+1]['Options'][x] + if option['Pinned'] or option["Previous Selection"] or option["Edited"]: + option_offset = x+1 + batch_offset = int((koboldai_vars.generated_tkns-1) / koboldai_vars.genamt) if koboldai_vars.alt_multi_gen else 0 + for batch_index, batch in enumerate(scores): + probs = F.softmax(batch, dim = -1).cpu().numpy() + + token_prob_info = [] + for token_id, score in sorted(enumerate(probs), key=lambda x: x[1], reverse=True)[:8]: + token_prob_info.append({ + "tokenId": token_id, + "decoded": utils.decodenewlines(tokenizer.decode(token_id)), + "score": float(score), + }) + + + if koboldai_vars.numseqs == 1: + koboldai_vars.actions.set_probabilities(token_prob_info) + else: + koboldai_vars.actions.set_option_probabilities(token_prob_info, batch_index+option_offset+batch_offset) + + return scores def new_get_logits_processor(*args, **kwargs) -> LogitsProcessorList: processors = new_get_logits_processor.old_get_logits_processor(*args, **kwargs) diff --git a/koboldai_settings.py b/koboldai_settings.py index ef2f9a1f..a0588572 100644 --- a/koboldai_settings.py +++ b/koboldai_settings.py @@ -1355,6 +1355,7 @@ class KoboldStoryRegister(object): self.make_audio_queue = multiprocessing.Queue() self.make_audio_thread_slow = None self.make_audio_queue_slow = multiprocessing.Queue() + self.probability_buffer = None for item in sequence: self.append(item) @@ -1541,7 +1542,7 @@ class KoboldStoryRegister(object): if self.actions[action_id]["Selected Text"] != text: self.actions[action_id]["Selected Text"] = text self.actions[action_id]["Time"] = self.actions[action_id].get("Time", int(time.time())) - if 'Probabilities' in self.actions[action_id]: + if 'buffer' in self.actions[action_id]: if self.koboldai_vars.tokenizer is not None: tokens = self.koboldai_vars.tokenizer.encode(text) for token_num in range(len(self.actions[action_id]["Probabilities"])): @@ -1846,11 +1847,16 @@ class KoboldStoryRegister(object): if len(self.koboldai_vars.tokenizer.encode(self.actions[self.action_count+1]['Selected Text'])) != self.koboldai_vars.genamt: #ui1 if queue is not None: - queue.put(["from_server", {"cmd": "streamtoken", "data": [{'decoded': text_list[0]}]}, {"broadcast":True, "room":"UI_1"}]) + queue.put(["from_server", {"cmd": "streamtoken", "data": [{ + "decoded": text_list[0], + "probabilities": self.probability_buffer + }]}, {"broadcast":True, "room":"UI_1"}]) #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) def set_probabilities(self, probabilities, action_id=None): + self.probability_buffer = probabilities + if action_id is None: action_id = self.action_count+1 if action_id in self.actions: diff --git a/static/application.js b/static/application.js index b7cf1b4a..30daada1 100644 --- a/static/application.js +++ b/static/application.js @@ -2378,7 +2378,6 @@ $(document).ready(function(){ })(); $("body").addClass("connected"); } else if (msg.cmd == "streamtoken") { - console.log(msg); // Sometimes the stream_token messages will come in too late, after // we have recieved the full text. This leads to some stray tokens // appearing after the output. To combat this, we only allow tokens