From 9287ff18decaab914f3a3f09711ab142c3d2a830 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sat, 24 Feb 2024 14:50:06 +0200
Subject: [PATCH] Fix for non-streaming

---
 public/script.js                   | 6 ++++++
 public/scripts/textgen-settings.js | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/public/script.js b/public/script.js
index a17dd748c..9f9f82a5c 100644
--- a/public/script.js
+++ b/public/script.js
@@ -18,6 +18,7 @@ import {
     textgen_types,
     getTextGenServer,
     validateTextGenUrl,
+    parseTextgenLogprobs,
 } from './scripts/textgen-settings.js';
 
 const { MANCER, TOGETHERAI, OOBA, APHRODITE, OLLAMA } = textgen_types;
@@ -4478,6 +4479,11 @@ function parseAndSaveLogprobs(data, continueFrom) {
             // `sendOpenAIRequest`. `data` for these APIs is just a string with
             // the text of the generated message, logprobs are not included.
             return;
+        case 'textgenerationwebui':
+            if (textgen_settings.type === textgen_types.LLAMACPP) {
+                logprobs = data?.completion_probabilities?.map(x => parseTextgenLogprobs(x.content, [x])) || null;
+            }
+            break;
         default:
             return;
     }
diff --git a/public/scripts/textgen-settings.js b/public/scripts/textgen-settings.js
index 793081665..910a03615 100644
--- a/public/scripts/textgen-settings.js
+++ b/public/scripts/textgen-settings.js
@@ -764,7 +764,7 @@ async function generateTextGenWithStreaming(generate_data, signal) {
  * @param {Object} logprobs - logprobs object returned from the API
  * @returns {import('logprobs.js').TokenLogprobs | null} - converted logprobs
  */
-function parseTextgenLogprobs(token, logprobs) {
+export function parseTextgenLogprobs(token, logprobs) {
     if (!logprobs) {
         return null;
     }