1
import {
2
    animation_duration,
3
    chat,
4
    cleanUpMessage,
5
    event_types,
6
    eventSource,
7
    Generate,
8
    getGeneratingApi,
9
    is_send_press,
10
    isStreamingEnabled,
11
    substituteParamsExtended,
12
} from '../script.js';
13
import { debounce, delay, getStringHash } from './utils.js';
14
import { decodeTextTokens, getTokenizerBestMatch } from './tokenizers.js';
15
import { power_user } from './power-user.js';
16
import { callGenericPopup, POPUP_TYPE } from './popup.js';
17
import { t } from './i18n.js';
18

19
const TINTS = 4;
20
const MAX_MESSAGE_LOGPROBS = 100;
21
const REROLL_BUTTON = $('#logprobsReroll');
22

23
/**
24
 * Tuple of a candidate token and its logarithm of probability of being chosen
25
 * @typedef {[string, number]} Candidate - (token, logprob)
26
 */
27

28
/**
29
 * @typedef {(Node|JQuery<Text>|JQuery<HTMLElement>)[]} NodeArray - Array of DOM nodes
30
 */
31

32
/**
33
 * Logprob data for a single message
34
 * @typedef {Object} MessageLogprobData
35
 * @property {number} created - timestamp of when the message was generated
36
 * @property {number} hash - hash of the message object
37
 * @property {number} messageId - ID of the source message
38
 * @property {number} swipeId - ID of the source swipe on the source message
39
 * @property {string} api - API used to generate the message
40
 * @property {TokenLogprobs[]} messageLogprobs Logprob data for each token, by
41
 * its index in the message
42
 * @property {string | null} continueFrom - the 'continue' prefix used to
43
 * generate the message, if any
44
 */
45

46
/**
47
 * Logprob data for a single token
48
 * @typedef {Object} TokenLogprobs
49
 * @property {string} token - A token generated by the model
50
 * @property {Candidate[]} topLogprobs - Array of top candidate tokens
51
 */
52

53
/**
54
 * State object for Token Probabilities
55
 * @typedef {Object} LogprobsState
56
 * @property {?TokenLogprobs} selectedTokenLogprobs Log probabilities for
57
 * currently-selected token.
58
 * @property {Map<number, MessageLogprobData>} messageLogprobs Log probabilities for
59
 * each message, keyed by message hash.
60
 */
61

62
/**
63
 * @type {LogprobsState} state
64
 */
65
const state = {
66
    selectedTokenLogprobs: null,
67
    messageLogprobs: new Map(),
68
};
69

70
/**
71
 * Renders the Token Probabilities UI and all subviews with the active message's
72
 * logprobs data. If the message has no token logprobs, a message is displayed.
73
 */
74
function renderAlternativeTokensView() {
75
    const view = $('#logprobs_generation_output');
76
    if (!view.is(':visible')) {
77
        return;
78
    }
79
    view.empty();
80
    state.selectedTokenLogprobs = null;
81
    renderTopLogprobs();
82

83
    const { messageLogprobs, continueFrom } = getActiveMessageLogprobData() || {};
84
    const usingSmoothStreaming = isStreamingEnabled() && power_user.smooth_streaming;
85
    if (!messageLogprobs?.length || usingSmoothStreaming) {
86
        const emptyState = $('<div></div>');
87
        const noTokensMsg = !power_user.request_token_probabilities
88
            ? '<span>Enable <b>Request token probabilities</b> in the User Settings menu to use this feature.</span>'
89
            : usingSmoothStreaming
90
                ? t`Token probabilities are not available when using Smooth Streaming.`
91
                : is_send_press
92
                    ? t`Generation in progress...`
93
                    : t`No token probabilities available for the current message.`;
94
        emptyState.html(noTokensMsg);
95
        emptyState.addClass('logprobs_empty_state');
96
        view.append(emptyState);
97
        return;
98
    }
99

100
    const prefix = continueFrom || '';
101
    const tokenSpans = [];
102
    REROLL_BUTTON.toggle(!!prefix);
103

104
    if (prefix) {
105
        REROLL_BUTTON.off('click').on('click', () => onPrefixClicked(prefix.length));
106

107
        let cumulativeOffset = 0;
108
        const words = prefix.split(/\s+/);
109
        const delimiters = prefix.match(/\s+/g) || []; // Capture the actual delimiters
110

111
        words.forEach((word, i) => {
112
            const span = $('<span></span>');
113
            span.text(`${word} `);
114

115
            span.addClass('logprobs_output_prefix');
116
            span.attr('title', t`Reroll from this point`);
117

118
            let offset = cumulativeOffset;
119
            span.on('click', () => onPrefixClicked(offset));
120
            addKeyboardProps(span);
121

122
            tokenSpans.push(span);
123
            tokenSpans.push(delimiters[i]?.includes('\n')
124
                ? document.createElement('br')
125
                : document.createTextNode(delimiters[i] || ' '),
126
            );
127

128
            cumulativeOffset += word.length + (delimiters[i]?.length || 0);
129
        });
130
    }
131

132
    messageLogprobs.forEach((tokenData, i) => {
133
        const { token } = tokenData;
134
        const span = $('<span></span>');
135
        const text = toVisibleWhitespace(token);
136
        span.text(text);
137
        span.addClass('logprobs_output_token');
138
        span.addClass('logprobs_tint_' + (i % TINTS));
139
        span.on('click', () => onSelectedTokenChanged(tokenData, span));
140
        addKeyboardProps(span);
141
        tokenSpans.push(...withVirtualWhitespace(token, span));
142
    });
143

144
    view.append(tokenSpans);
145

146
    // scroll past long prior context
147
    if (prefix) {
148
        const element = view.find('.logprobs_output_token').first();
149
        const scrollOffset = element.offset().top - element.parent().offset().top;
150
        element.parent().scrollTop(scrollOffset);
151
    }
152
}
153

154
function addKeyboardProps(element) {
155
    element.attr('role', 'button');
156
    element.attr('tabindex', '0');
157
    element.keydown(function (e) {
158
        if (e.key === 'Enter' || e.key === ' ') {
159
            element.click();
160
        }
161
    });
162
}
163

164
/**
165
 * renderTopLogprobs renders the top logprobs subview with the currently
166
 * selected token highlighted. If no token is selected, the subview is hidden.
167
 *
168
 * Callers:
169
 * - renderAlternativeTokensView, to render the entire view
170
 * - onSelectedTokenChanged, to update the view when a token is selected
171
 */
172
function renderTopLogprobs() {
173
    $('#logprobs_top_logprobs_hint').hide();
174
    const view = $('.logprobs_candidate_list');
175
    view.empty();
176

177
    if (!state.selectedTokenLogprobs) {
178
        return;
179
    }
180

181
    const { token: selectedToken, topLogprobs } = state.selectedTokenLogprobs;
182

183
    let sum = 0;
184
    const nodes = [];
185
    const candidates = topLogprobs
186
        .sort(([, logA], [, logB]) => logB - logA)
187
        .map(([text, log]) => {
188
            if (log <= 0) {
189
                const probability = Math.exp(log);
190
                sum += probability;
191
                return [text, probability, log];
192
            } else {
193
                return [text, log, null];
194
            }
195
        });
196
    candidates.push(['<others>', 1 - sum, 0]);
197

198
    let matched = false;
199
    for (const [token, probability, log] of candidates) {
200
        const container = $('<button class="flex-container flexFlowColumn logprobs_top_candidate"></button>');
201
        const tokenNormalized = String(token).replace(/^[▁Ġ]/g, ' ');
202

203
        if (token === selectedToken || tokenNormalized === selectedToken) {
204
            matched = true;
205
            container.addClass('selected');
206
        }
207

208
        const tokenText = $('<span></span>').text(`${toVisibleWhitespace(token.toString())}`);
209
        const percentText = $('<span></span>').text(`${(+probability * 100).toFixed(2)}%`);
210
        container.append(tokenText, percentText);
211
        if (log) {
212
            container.attr('title', `logarithm: ${log}`);
213
        }
214
        addKeyboardProps(container);
215
        if (token !== '<others>') {
216
            container.on('click', () => onAlternativeClicked(state.selectedTokenLogprobs, token.toString()));
217
        } else {
218
            container.prop('disabled', true);
219
        }
220
        nodes.push(container);
221
    }
222

223
    // Highlight the <others> node if the selected token was not included in the
224
    // top logprobs
225
    if (!matched) {
226
        nodes[nodes.length - 1].css('background-color', 'rgba(255, 0, 0, 0.1)');
227
    }
228

229
    view.append(nodes);
230
}
231

232
/**
233
 * User clicks on a token in the token output view. It updates the selected token state
234
 * and re-renders the top logprobs view, or deselects the token if it was already selected.
235
 * @param {TokenLogprobs} logprobs - logprob data for the selected token
236
 * @param {Node|JQuery} span - target span node that was clicked
237
 */
238
function onSelectedTokenChanged(logprobs, span) {
239
    $('.logprobs_output_token.selected').removeClass('selected');
240
    if (state.selectedTokenLogprobs === logprobs) {
241
        state.selectedTokenLogprobs = null;
242
    } else {
243
        state.selectedTokenLogprobs = logprobs;
244
        $(span).addClass('selected');
245
    }
246
    renderTopLogprobs();
247
}
248

249
/**
250
 * onAlternativeClicked is called when the user clicks on an alternative token
251
 * in the top logprobs view. It will create a new swipe message and prefill it
252
 * with all text up to the selected token, followed by the chosen alternative.
253
 * Then it requests a `continue` completion from the model with the new prompt.
254
 * @param {TokenLogprobs} tokenLogprobs - logprob data for selected alternative
255
 * @param {string} alternative - selected alternative token's text
256
 */
257
function onAlternativeClicked(tokenLogprobs, alternative) {
258
    if (!checkGenerateReady()) {
259
        return;
260
    }
261

262
    if (getGeneratingApi() === 'openai') {
263
        const title = t`Feature unavailable`;
264
        const message = t`Due to API limitations, rerolling a token is not supported with OpenAI. Try switching to a different API.`;
265
        const content = `<h3>${title}</h3><p>${message}</p>`;
266
        return callGenericPopup(content, POPUP_TYPE.TEXT);
267
    }
268

269
    const { messageLogprobs, continueFrom } = getActiveMessageLogprobData();
270
    const replaceIndex = messageLogprobs.findIndex(x => x === tokenLogprobs);
271

272
    const tokens = messageLogprobs.slice(0, replaceIndex + 1).map(({ token }) => token);
273
    tokens[replaceIndex] = String(alternative).replace(/^[▁Ġ]/g, ' ').replace(/Ċ/g, '\n');
274

275
    const prefix = continueFrom || '';
276
    const prompt = prefix + tokens.join('');
277
    addGeneration(prompt);
278
}
279

280
/**
281
 * User clicks on the reroll button in the token output view, or on a word in the
282
 * prefix. Retrieve the prefix for the current message and truncate it at the
283
 * offset for the selected word. Then request a `continue` completion from the
284
 * model with the new prompt.
285
 *
286
 * If no offset is provided, the entire prefix will be rerolled.
287
 *
288
 * @param {number} offset - index of the token in the prefix to reroll from
289
 * @returns {void}
290
 * @param offset
291
 */
292
function onPrefixClicked(offset = undefined) {
293
    if (!checkGenerateReady()) {
294
        return;
295
    }
296

297
    const { continueFrom } = getActiveMessageLogprobData() || {};
298
    const prefix = continueFrom ? continueFrom.substring(0, offset) : '';
299
    addGeneration(prefix);
300
}
301

302
function checkGenerateReady() {
303
    if (is_send_press) {
304
        toastr.warning('Please wait for the current generation to complete.');
305
        return false;
306
    }
307
    return true;
308
}
309

310
/**
311
 * Generates a new swipe as a continuation of the given prompt, when user selects
312
 * an alternative token or rerolls from a prefix.
313
 *
314
 * @param prompt
315
 */
316
function addGeneration(prompt) {
317
    const messageId = chat.length - 1;
318
    if (prompt && prompt.length > 0) {
319
        createSwipe(messageId, prompt);
320
        $('.swipe_right:last').trigger('click');
321
        void Generate('continue');
322
    } else {
323
        $('.swipe_right:last').trigger('click');
324
    }
325
}
326

327
/**
328
 * onToggleLogprobsPanel is called when the user performs an action that toggles
329
 * the logprobs view, such as clicking the Token Probabilities menu item or the
330
 * close button.
331
 */
332
function onToggleLogprobsPanel() {
333
    const logprobsViewer = $('#logprobsViewer');
334

335
    // largely copied from CFGScale toggle
336
    if (logprobsViewer.css('display') === 'none') {
337
        logprobsViewer.addClass('resizing');
338
        logprobsViewer.css('display', 'flex');
339
        logprobsViewer.css('opacity', 0.0);
340
        renderAlternativeTokensView();
341
        logprobsViewer.transition({
342
            opacity: 1.0,
343
            duration: animation_duration,
344
        }, async function () {
345
            await delay(50);
346
            logprobsViewer.removeClass('resizing');
347
        });
348
    } else {
349
        logprobsViewer.addClass('resizing');
350
        logprobsViewer.transition({
351
            opacity: 0.0,
352
            duration: animation_duration,
353
        },
354
        async function () {
355
            await delay(50);
356
            logprobsViewer.removeClass('resizing');
357
        });
358
        setTimeout(function () {
359
            logprobsViewer.hide();
360
        }, animation_duration);
361
    }
362
}
363

364
/**
365
 * Appends a new swipe to the target chat message with the given text.
366
 * @param {number} messageId - target chat message ID
367
 * @param {string} prompt - initial prompt text which will be continued
368
 */
369
function createSwipe(messageId, prompt) {
370
    // need to call `cleanUpMessage` on our new prompt, because we were working
371
    // with raw model output and our new prompt is missing trimming/macro replacements
372
    let cleanedPrompt = cleanUpMessage({
373
        getMessage: prompt,
374
        isImpersonate: false,
375
        isContinue: false,
376
        displayIncompleteSentences: true,
377
    });
378

379
    const msg = chat[messageId];
380

381
    const reasoningPrefix = substituteParamsExtended(power_user.reasoning.prefix);
382
    const reasoningSuffix = substituteParamsExtended(power_user.reasoning.suffix);
383
    const isReasoningAutoParsed = power_user.reasoning.auto_parse;
384
    const msgHasParsedReasoning = msg.extra?.reasoning?.length > 0;
385
    let shouldRerollReasoning = false;
386

387
    //if we have pre-existing reasoning and are currently autoparsing
388
    if (isReasoningAutoParsed && msgHasParsedReasoning) {
389
        console.debug('saw autoparse on with reasoning in message');
390
        //but the reroll prompt does not include the end of reasoning
391
        if (cleanedPrompt.includes(reasoningPrefix) && !cleanedPrompt.includes(reasoningSuffix)) {
392
            //we need to send the results to the reasoning block
393
            //this will involve the ReasoningHandler from reasoning.js
394
            console.debug('..with start tag but no end tag... reroll reasoning');
395
            shouldRerollReasoning = true;
396
        }
397

398
        let hasReasoningPrefix = cleanedPrompt.includes(reasoningPrefix);
399
        let hasReasoningSuffix = cleanedPrompt.includes(reasoningSuffix);
400

401
        //..with both the start and end think tags
402
        //OR
403
        //..with only the end think tag (implying prefilled think start)
404
        if (hasReasoningPrefix && hasReasoningSuffix) {
405
            //we need to send the results to the response block without reasoning attached
406
            console.debug('...incl. end tag...rerolling response');
407
            const endOfThink = cleanedPrompt.indexOf(reasoningSuffix) + reasoningSuffix.length;
408
            cleanedPrompt = cleanedPrompt.substring(endOfThink);
409
        }
410

411
        //if cleanedprompt includes the think prefix, but no suffix..
412
        if (hasReasoningPrefix && !hasReasoningSuffix) {
413
            console.debug('..no end tag...rerolling reasoning, so removing prefix');
414
            cleanedPrompt = cleanedPrompt.replace(reasoningPrefix, '');
415
        }
416
    }
417

418
    console.debug('cleanedPrompt: ', cleanedPrompt);
419

420
    /** @type {SwipeInfo} */
421
    const newSwipeInfo = {
422
        send_date: msg.send_date,
423
        gen_started: msg.gen_started,
424
        gen_finished: msg.gen_finished,
425
        extra: { ...structuredClone(msg.extra), from_logprobs: new Date().getTime() },
426
    };
427

428
    msg.swipes = msg.swipes || [];
429
    msg.swipe_info = msg.swipe_info || [];
430

431
    // Add our new swipe, then make sure the active swipe is the one just before
432
    // it. The call to `swipe_right` in addGeneration() will switch to it immediately.
433

434
    //if we determined that we need to reroll from reasoning
435
    if (shouldRerollReasoning) {
436
        //cleaned prompt goes into reasoning
437
        newSwipeInfo.extra.reasoning = cleanedPrompt;
438
        //mes_text becomes empty, causing the reasoning handler to parse the reasoning first
439
        msg.swipes.push('');
440
    } else {
441
        //otherwise just add the cleaned prompt to the message and continue
442
        msg.swipes.push(cleanedPrompt);
443
    }
444

445
    msg.swipe_info.push(newSwipeInfo);
446
    msg.swipe_id = Math.max(0, msg.swipes.length - 2);
447
}
448

449
/**
450
 * toVisibleWhitespace receives input text and replaces spaces with &middot; and
451
 * newlines with ↵.
452
 * @param {string} input
453
 * @returns {string}
454
 */
455
function toVisibleWhitespace(input) {
456
    return input.replace(/ /g, '·').replace(/[▁Ġ]/g, '·').replace(/[Ċ\n]/g, '↵');
457
}
458

459
/**
460
 * withVirtualWhitespace inserts line breaks and a zero-width space before and
461
 * after the span node if its token begins or ends with whitespace in order to
462
 * allow text to wrap despite whitespace characters being replaced with a dot.
463
 * @param {string} text - token text being evaluated for whitespace
464
 * @param {Node|JQuery} span - target span node to be wrapped
465
 * @returns {NodeArray} - array of nodes to be appended to the parent element
466
 */
467
function withVirtualWhitespace(text, span) {
468
    /** @type {NodeArray} */
469
    const result = [span];
470
    if (text.match(/^\s/)) {
471
        result.unshift(document.createTextNode('\u200b'));
472
    }
473
    if (text.match(/\s$/)) {
474
        result.push($(document.createTextNode('\u200b')));
475
    }
476
    if (text.match(/^[▁Ġ]/)) {
477
        result.unshift(document.createTextNode('\u200b'));
478
    }
479
    // line breaks are trickier. we don't currently handle consecutive line
480
    // breaks or line breaks occuring in between non-whitespace characters, but
481
    // tokenizers generally don't produce those anyway.
482

483
    // matches leading line break, at least one character, and trailing line break
484
    if (text.match(/^\n(?:.|\n)+\n$/)) {
485
        result.unshift($('<br>'));
486
        result.push($('<br>'));
487
    } else if (text.match(/^\n/)) {
488
        result.unshift($('<br>'));
489
    } else if (text.match(/\n$/)) {
490
        result.push($('<br>'));
491
    }
492
    return result;
493
}
494

495
/**
496
 * Receives the top logprobs for each token in a message and associates it with the active message.
497
 *
498
 * Ensure the active message has been updated and rendered before calling this function
499
 * or the logprobs data will be saved to the wrong message.
500
 *
501
 * Callers:
502
 * - Generate:onSuccess via saveLogprobsForActiveMessage, for non-streaming text completion
503
 * - StreamingProcessor:onFinishStreaming, for streaming text completion
504
 * - sendOpenAIRequest, for non-streaming chat completion
505
 *
506
 * @param {TokenLogprobs[]} logprobs - array of logprobs data for each token
507
 * @param {string | null} continueFrom  - for 'continue' generations, the prompt
508
 */
509
export function saveLogprobsForActiveMessage(logprobs, continueFrom) {
510
    if (!logprobs) {
511
        // non-streaming APIs could return null data
512
        return;
513
    }
514

515
    // NovelAI only returns token IDs in logprobs data; convert to text tokens in-place
516
    if (getGeneratingApi() === 'novel') {
517
        convertTokenIdLogprobsToText(logprobs);
518
    }
519

520
    const msgId = chat.length - 1;
521
    /** @type {MessageLogprobData} */
522
    const data = {
523
        created: new Date().getTime(),
524
        api: getGeneratingApi(),
525
        messageId: msgId,
526
        swipeId: chat[msgId].swipe_id,
527
        messageLogprobs: logprobs,
528
        continueFrom,
529
        hash: getMessageHash(chat[msgId]),
530
    };
531

532
    state.messageLogprobs.set(data.hash, data);
533

534
    // Clean up old logprobs data
535
    const oldLogprobs = Array.from(state.messageLogprobs.values())
536
        .sort((a, b) => b.created - a.created)
537
        .slice(MAX_MESSAGE_LOGPROBS);
538
    for (const oldData of oldLogprobs) {
539
        state.messageLogprobs.delete(oldData.hash);
540
    }
541
}
542

543
function getMessageHash(message) {
544
    // We don't use the swipe ID as a hash component because it's not stable,
545
    // deleting a swipe will change the ID of all subsequent swipes.
546
    const hashParams = {
547
        name: message.name,
548
        mid: chat.indexOf(message),
549
        text: message.mes,
550
    };
551
    return getStringHash(JSON.stringify(hashParams));
552
}
553

554
/**
555
 * getActiveMessageLogprobData returns the logprobs data for the active chat
556
 * message.
557
 * @returns {MessageLogprobData|null}
558
 */
559
function getActiveMessageLogprobData() {
560
    if (chat.length === 0) {
561
        return null;
562
    }
563

564
    const hash = getMessageHash(chat[chat.length - 1]);
565
    return state.messageLogprobs.get(hash) || null;
566
}
567

568

569
/**
570
 * convertLogprobTokenIdsToText replaces token IDs in logprobs data with text tokens,
571
 * for APIs that return token IDs instead of text tokens, to wit: NovelAI.
572
 *
573
 * @param {TokenLogprobs[]} input - logprobs data with numeric token IDs
574
 */
575
function convertTokenIdLogprobsToText(input) {
576
    const api = getGeneratingApi();
577
    if (api !== 'novel') {
578
        // should have been checked by the caller
579
        throw new Error('convertTokenIdLogprobsToText should only be called for NovelAI');
580
    }
581

582
    const tokenizerId = getTokenizerBestMatch(api);
583

584
    /** @type {any[]} Flatten unique token IDs across all logprobs */
585
    const tokenIds = Array.from(new Set(input.flatMap(logprobs =>
586
        logprobs.topLogprobs.map(([token]) => token).concat(logprobs.token),
587
    )));
588

589
    // Submit token IDs to tokenizer to get token text, then build ID->text map
590
    // noinspection JSCheckFunctionSignatures - mutates input in-place
591
    const { chunks } = decodeTextTokens(tokenizerId, tokenIds);
592
    const tokenIdText = new Map(tokenIds.map((id, i) => [id, chunks[i]]));
593

594
    // Fixup logprobs data with token text
595
    input.forEach(logprobs => {
596
        logprobs.token = tokenIdText.get(logprobs.token);
597
        logprobs.topLogprobs = logprobs.topLogprobs.map(([token, logprob]) =>
598
            [tokenIdText.get(token), logprob],
599
        );
600
    });
601
}
602

603
export function initLogprobs() {
604
    REROLL_BUTTON.hide();
605
    const debouncedRender = debounce(renderAlternativeTokensView);
606
    $('#logprobsViewerClose').on('click', onToggleLogprobsPanel);
607
    $('#option_toggle_logprobs').on('click', onToggleLogprobsPanel);
608
    eventSource.on(event_types.CHAT_CHANGED, debouncedRender);
609
    eventSource.on(event_types.CHARACTER_MESSAGE_RENDERED, debouncedRender);
610
    eventSource.on(event_types.IMPERSONATE_READY, debouncedRender);
611
    eventSource.on(event_types.MESSAGE_DELETED, debouncedRender);
612
    eventSource.on(event_types.MESSAGE_EDITED, debouncedRender);
613
    eventSource.on(event_types.MESSAGE_SWIPED, debouncedRender);
614
}