Merge pull request #3078 from ceruleandeep/chore/lintAndDocumentLogprobs Make dynamic reroll available without use of modifier key

61469ec999d4f16755887008b501b36a187bbca6

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed

Parents: 14d5c669b1252cd0833b4c914d204454db01cc49 b143b3c8208c83264e70fc2337fd0cce6378ef76

public/css/logprobs.css+8 -0 Viewed

     opacity: 0.5;
+}
+.logprobs_output_prefix:hover {
+    background-color: rgba(255, 0, 50, 0.4);
+}
+.logprobs_output_prefix:hover ~ .logprobs_output_prefix {
+    background-color: rgba(255, 0, 50, 0.4);
+}
 .logprobs_candidate_list {
     grid-row-start: 3;
     grid-row-end: 4;

public/index.html+4 -1 Viewed

                 </div>
             </div>
             <div class="logprobs_panel_content inline-drawer-content flex-container flexFlowColumn">
-                <small>
+                <small class="flex-container alignItemsCenter justifySpaceBetween">
                     <b data-i18n="Select a token to see alternatives considered by the AI.">Select a token to see alternatives considered by the AI.</b>
+                    <button id="logprobsReroll" class="menu_button" title="Reroll with the entire prefix" data-i18n="[title]Reroll with the entire prefix">
+                        <span class="fa-solid fa-redo logprobs_reroll"></span>
+                    </button>
                 </small>
                 <hr>
                 <div id="logprobs_generation_output"></div>

public/scripts/logprobs.js+129 -114 Viewed

 import {
     animation_duration,
-    callPopup,
     chat,
     cleanUpMessage,
     event_types,

 import { debounce, delay, getStringHash } from './utils.js';
 import { decodeTextTokens, getTokenizerBestMatch } from './tokenizers.js';
 import { power_user } from './power-user.js';
+import { callGenericPopup, POPUP_TYPE } from './popup.js';
+import { t } from './i18n.js';
 const TINTS = 4;
 const MAX_MESSAGE_LOGPROBS = 100;
+const REROLL_BUTTON = $('#logprobsReroll');
 /**
  * Tuple of a candidate token and its logarithm of probability of being chosen

  */
 /**
+ * @typedef {(Node|JQuery<Text>|JQuery<HTMLElement>)[]} NodeArray - Array of DOM nodes
+ */
+/**
  * Logprob data for a single message
  * @typedef {Object} MessageLogprobData
  * @property {number} created - timestamp of when the message was generated

  * @property {Candidate[]} topLogprobs - Array of top candidate tokens
  */
-let state = {
+/**
-    /** @type {TokenLogprobs | null} */
+ * State object for Token Probabilities
+ * @typedef {Object} LogprobsState
+ * @property {?TokenLogprobs} selectedTokenLogprobs Log probabilities for
+ * currently-selected token.
+ * @property {Map<number, MessageLogprobData>} messageLogprobs Log probabilities for
+ * each message, keyed by message hash.
+ */
+/**
+ * @type {LogprobsState} state
+ */
+const state = {
     selectedTokenLogprobs: null,
-    /** @type {Map<number, MessageLogprobData>} */
     messageLogprobs: new Map(),
 };
 /**
- * renderAlternativeTokensView renders the Token Probabilities UI and all
+ * Renders the Token Probabilities UI and all subviews with the active message's
- * subviews with the active message's logprobs data. If the message has no token
+ * logprobs data. If the message has no token logprobs, a message is displayed.
- * logprobs, a zero-state is rendered.
  */
 function renderAlternativeTokensView() {
     const view = $('#logprobs_generation_output');

     const usingSmoothStreaming = isStreamingEnabled() && power_user.smooth_streaming;
     if (!messageLogprobs?.length || usingSmoothStreaming) {
         const emptyState = $('<div></div>');
-        const noTokensMsg = usingSmoothStreaming
+        const noTokensMsg = !power_user.request_token_probabilities
-            ? 'Token probabilities are not available when using Smooth Streaming.'
+            ? '<span>Enable <b>Request token probabilities</b> in the User Settings menu to use this feature.</span>'
-            : 'No token probabilities available for the current message.';
+            : usingSmoothStreaming
-        const msg = power_user.request_token_probabilities
+                ? t`Token probabilities are not available when using Smooth Streaming.`
-            ? noTokensMsg
+                : is_send_press
-            : '<span>Enable <b>Request token probabilities</b> in the User Settings menu to use this feature.</span>';
+                    ? t`Generation in progress...`
-        emptyState.html(msg);
+                    : t`No token probabilities available for the current message.`;
+        emptyState.html(noTokensMsg);
         emptyState.addClass('logprobs_empty_state');
         view.append(emptyState);
         return;

     const prefix = continueFrom || '';
     const tokenSpans = [];
+    REROLL_BUTTON.toggle(!!prefix);
     if (prefix) {
-        const prefixSpan = $('<span></span>');
+        REROLL_BUTTON.off('click').on('click', () => onPrefixClicked(prefix.length));
-        prefixSpan.text(prefix);
-        prefixSpan.html(prefixSpan.html().replace(/\n/g, '<br>'));
+        let cumulativeOffset = 0;
-        prefixSpan.addClass('logprobs_output_prefix');
+        const words = prefix.split(/\s+/);
-        prefixSpan.attr('title', 'Select to reroll the last \'Continue\' generation.\nHold the CTRL key when clicking to reroll from before that word.');
+        const delimiters = prefix.match(/\s+/g) || []; // Capture the actual delimiters
-        prefixSpan.click(onPrefixClicked);
-        addKeyboardProps(prefixSpan);
+        words.forEach((word, i) => {
-        tokenSpans.push(...withVirtualWhitespace(prefix, prefixSpan));
+            const span = $('<span></span>');
+            span.text(`${word} `);
+            span.addClass('logprobs_output_prefix');
+            span.attr('title', t`Reroll from this point`);
+            let offset = cumulativeOffset;
+            span.on('click', () => onPrefixClicked(offset));
+            addKeyboardProps(span);
+            tokenSpans.push(span);
+            tokenSpans.push(delimiters[i]?.includes('\n')
+                ? document.createElement('br')
+                : document.createTextNode(delimiters[i] || ' '),
+            );
+            cumulativeOffset += word.length + (delimiters[i]?.length || 0);
+        });
+    }
     messageLogprobs.forEach((tokenData, i) => {

         span.text(text);
         span.addClass('logprobs_output_token');
         span.addClass('logprobs_tint_' + (i % TINTS));
-        span.click(() => onSelectedTokenChanged(tokenData, span));
+        span.on('click', () => onSelectedTokenChanged(tokenData, span));
         addKeyboardProps(span);
         tokenSpans.push(...withVirtualWhitespace(token, span));
     });

 /**
  * renderTopLogprobs renders the top logprobs subview with the currently
  * selected token highlighted. If no token is selected, the subview is hidden.
+ *
+ * Callers:
+ * - renderAlternativeTokensView, to render the entire view
+ * - onSelectedTokenChanged, to update the view when a token is selected
  */
 function renderTopLogprobs() {
     $('#logprobs_top_logprobs_hint').hide();

                 const probability = Math.exp(log);
                 sum += probability;
                 return [text, probability, log];
-            }
+            } else {
-            else {
                 return [text, log, null];
+            }
         });

@@ -167,15 +204,15 @@ function renderTopLogprobs() {
167	container.addClass('selected');	204	container.addClass('selected');
168	}	205	}
169		206
170	const tokenText = $('<span></span>').text(`${toVisibleWhitespace(token)}`);	207	const tokenText = $('<span></span>').text(`${toVisibleWhitespace(token.toString())}`);
171	const percentText = $('<span></span>').text(`${(probability * 100).toFixed(2)}%`);	208	const percentText = $('<span></span>').text(`${(+probability * 100).toFixed(2)}%`);
172	container.append(tokenText, percentText);	209	container.append(tokenText, percentText);
173	if (log) {	210	if (log) {
174	container.attr('title', `logarithm: ${log}`);	211	container.attr('title', `logarithm: ${log}`);
175	}	212	}
176	addKeyboardProps(container);	213	addKeyboardProps(container);
177	if (token !== '<others>') {	214	if (token !== '<others>') {
178	container.click(() => onAlternativeClicked(state.selectedTokenLogprobs, token));	215	container.on('click', () => onAlternativeClicked(state.selectedTokenLogprobs, token.toString()));
179	} else {	216	} else {
180	container.prop('disabled', true);	217	container.prop('disabled', true);
181	}	218	}

+}
 /**
- * onSelectedTokenChanged is called when the user clicks on a token in the
+ * User clicks on a token in the token output view. It updates the selected token state
- * token output view. It updates the selected token state and re-renders the
+ * and re-renders the top logprobs view, or deselects the token if it was already selected.
- * top logprobs view, or deselects the token if it was already selected.
  * @param {TokenLogprobs} logprobs - logprob data for the selected token
- * @param {Element} span - target span node that was clicked
+ * @param {Node|JQuery} span - target span node that was clicked
  */
 function onSelectedTokenChanged(logprobs, span) {
     $('.logprobs_output_token.selected').removeClass('selected');

@@ -223,7 +259,10 @@ function onAlternativeClicked(tokenLogprobs, alternative) {
223	}	259	}
224		260
225	if (getGeneratingApi() === 'openai') {	261	if (getGeneratingApi() === 'openai') {
226	return callPopup('<h3>Feature unavailable</h3><p>Due to API limitations, rerolling a token is not supported with OpenAI. Try switching to a different API.</p>', 'text');	262	const title = t`Feature unavailable`;
		263	const message = t`Due to API limitations, rerolling a token is not supported with OpenAI. Try switching to a different API.`;
		264	const content = `<h3>${title}</h3><p>${message}</p>`;
		265	return callGenericPopup(content, POPUP_TYPE.TEXT);
227	}	266	}
228		267
229	const { messageLogprobs, continueFrom } = getActiveMessageLogprobData();	268	const { messageLogprobs, continueFrom } = getActiveMessageLogprobData();

     const prefix = continueFrom || '';
     const prompt = prefix + tokens.join('');
-    const messageId = chat.length - 1;
+    addGeneration(prompt);
-    createSwipe(messageId, prompt);
-    $('.swipe_right:last').click(); // :see_no_evil:
-    Generate('continue').then(_ => void _);
+}
 /**
- * getTextBeforeClickedWord retrieves the portion of text within a span
+ * User clicks on the reroll button in the token output view, or on a word in the
- * that appears before the word clicked by the user. Using the x and y
+ * prefix. Retrieve the prefix for the current message and truncate it at the
- * coordinates from a PointerEvent, this function identifies the exact
+ * offset for the selected word. Then request a `continue` completion from the
- * word clicked and returns the text preceding it within the span.
+ * model with the new prompt.
- *
- * If the clicked position does not resolve to a valid word or text node,
- * the entire span text is returned as a fallback.
+ *
- * @param {PointerEvent} event - The click event containing the x and y coordinates.
+ * If no offset is provided, the entire prefix will be rerolled.
- * @param {string} spanText - The full text content of the span element.
- * @returns {string} The text before the clicked word, or the entire span text as fallback.
- */
-function getTextBeforeClickedWord(event, spanText) {
-    const x = event.clientX;
-    const y = event.clientY;
-    const range = document.caretRangeFromPoint(x, y);
-    if (range && range.startContainer.nodeType === Node.TEXT_NODE) {
-        const textNode = range.startContainer;
-        const offset = range.startOffset;
-        // Get the full text content of the text node
-        const text = textNode.nodeValue;
-        // Find the boundaries of the clicked word
-        const start = text.lastIndexOf(' ', offset - 1) + 1;
-        // Return the text before the clicked word
-        return text.slice(0, start);
-    }
-    // If we can't determine the exact word, return the full span text as a fallback
-    return spanText;
-}
-/**
- * onPrefixClicked is called when the user clicks on the carried-over prefix
- * in the token output view. It allows them to reroll the last 'continue'
- * completion with none of the output generated from it, in case they don't
- * like the results.
+ *
- * If the user holds the Ctrl key while clicking, only the portion of text
+ * @param {number} offset - index of the token in the prefix to reroll from
- * before the clicked word is retained as the prefix for rerolling
+ * @returns {void}
+ * @param offset
  */
-function onPrefixClicked() {
+function onPrefixClicked(offset = undefined) {
     if (!checkGenerateReady()) {
         return;
+    }
-    const { continueFrom } = getActiveMessageLogprobData();
+    const { continueFrom } = getActiveMessageLogprobData() || {};
-    const messageId = chat.length - 1;
+    const prefix = continueFrom ? continueFrom.substring(0, offset) : '';
+    addGeneration(prefix);
-    // Check if Ctrl key is pressed during the click
-    let prefix = continueFrom || '';
-    if (event.ctrlKey) {
-        // Ctrl is pressed - use the text before the clicked word
-        prefix = getTextBeforeClickedWord(event, continueFrom);
-    }
-    // Use the determined `prefix`
-    createSwipe(messageId, prefix);
-    $('.swipe_right:last').click();
-    Generate('continue').then(_ => void _);
+}
 function checkGenerateReady() {

     return true;
+}
+/**
+ * Generates a new swipe as a continuation of the given prompt, when user selects
+ * an alternative token or rerolls from a prefix.
+ *
+ * @param prompt
+ */
+function addGeneration(prompt) {
+    const messageId = chat.length - 1;
+    if (prompt && prompt.length > 0) {
+        createSwipe(messageId, prompt);
+        $('.swipe_right:last').trigger('click');
+        void Generate('continue');
+    } else {
+        $('.swipe_right:last').trigger('click');
+    }
+}
 /**
  * onToggleLogprobsPanel is called when the user performs an action that toggles

+}
 /**
- * createSwipe appends a new swipe to the target chat message with the given
+ * Appends a new swipe to the target chat message with the given text.
- * text.
  * @param {number} messageId - target chat message ID
  * @param {string} prompt - initial prompt text which will be continued
  */

  * after the span node if its token begins or ends with whitespace in order to
  * allow text to wrap despite whitespace characters being replaced with a dot.
  * @param {string} text - token text being evaluated for whitespace
- * @param {Element} span - target span node to be wrapped
+ * @param {Node|JQuery} span - target span node to be wrapped
- * @returns {Element[]} array of nodes to be appended to the DOM
+ * @returns {NodeArray} - array of nodes to be appended to the parent element
  */
 function withVirtualWhitespace(text, span) {
+    /** @type {NodeArray} */
     const result = [span];
     if (text.match(/^\s/)) {
         result.unshift(document.createTextNode('\u200b'));

+}
 /**
- * saveLogprobsForActiveMessage receives an array of TokenLogprobs objects
+ * Receives the top logprobs for each token in a message and associates it with the active message.
- * representing the top logprobs for each token in a message and associates it
+ *
- * with the active message.
+ * Ensure the active message has been updated and rendered before calling this function
+ * or the logprobs data will be saved to the wrong message.
+ *
+ * Callers:
+ * - Generate:onSuccess via saveLogprobsForActiveMessage, for non-streaming text completion
+ * - StreamingProcessor:onFinishStreaming, for streaming text completion
+ * - sendOpenAIRequest, for non-streaming chat completion
+ *
- * **Ensure the active message has been updated and rendered before calling
- * this function or the logprobs data will be saved to the wrong message.**
  * @param {TokenLogprobs[]} logprobs - array of logprobs data for each token
  * @param {string | null} continueFrom  - for 'continue' generations, the prompt
  */

         return;
+    }
-    convertTokenIdLogprobsToText(logprobs);
+    // NovelAI only returns token IDs in logprobs data; convert to text tokens in-place
+    if (getGeneratingApi() === 'novel') {
+        convertTokenIdLogprobsToText(logprobs);
+    }
     const msgId = chat.length - 1;
     /** @type {MessageLogprobData} */

     return state.messageLogprobs.get(hash) || null;
+}
 /**
- * convertLogprobTokenIdsToText mutates the given logprobs data's topLogprobs
+ * convertLogprobTokenIdsToText replaces token IDs in logprobs data with text tokens,
- * field keyed by token text instead of token ID. This is only necessary for
+ * for APIs that return token IDs instead of text tokens, to wit: NovelAI.
- * APIs which only return token IDs in their logprobs data; for others this
+ *
- * function is a no-op.
  * @param {TokenLogprobs[]} input - logprobs data with numeric token IDs
  */
 function convertTokenIdLogprobsToText(input) {
     const api = getGeneratingApi();
     if (api !== 'novel') {
-        return input;
+        // should have been checked by the caller
+        throw new Error('convertTokenIdLogprobsToText should only be called for NovelAI');
+    }
     const tokenizerId = getTokenizerBestMatch(api);

     )));
     // Submit token IDs to tokenizer to get token text, then build ID->text map
-    const { chunks } = decodeTextTokens(tokenizerId, tokenIds);
+    // noinspection JSCheckFunctionSignatures - mutates input in-place
+    const { chunks } = decodeTextTokens(tokenizerId, tokenIds.map(parseInt));
     const tokenIdText = new Map(tokenIds.map((id, i) => [id, chunks[i]]));
     // Fixup logprobs data with token text

+}
 export function initLogprobs() {
+    REROLL_BUTTON.hide();
     const debouncedRender = debounce(renderAlternativeTokensView);
-    $('#logprobsViewerClose').click(onToggleLogprobsPanel);
+    $('#logprobsViewerClose').on('click', onToggleLogprobsPanel);
-    $('#option_toggle_logprobs').click(onToggleLogprobsPanel);
+    $('#option_toggle_logprobs').on('click', onToggleLogprobsPanel);
     eventSource.on(event_types.CHAT_CHANGED, debouncedRender);
     eventSource.on(event_types.CHARACTER_MESSAGE_RENDERED, debouncedRender);
     eventSource.on(event_types.IMPERSONATE_READY, debouncedRender);