Merge pull request #354 from one-some/ui2-but-ui1-probs-fix

Fix probabilities on UI1

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)

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:

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