1
import { main_api } from '../../../script.js';
2
import { getContext } from '../../extensions.js';
3
import { SlashCommand } from '../../slash-commands/SlashCommand.js';
4
import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';
5
import { getFriendlyTokenizerName, getTextTokens, getTokenCountAsync, tokenizers } from '../../tokenizers.js';
6
import { resetScrollHeight, debounce } from '../../utils.js';
7
import { debounce_timeout } from '../../constants.js';
8
import { POPUP_TYPE, callGenericPopup } from '../../popup.js';
9
import { renderExtensionTemplateAsync } from '../../extensions.js';
10
import { t } from '../../i18n.js';
11

12
async function doTokenCounter() {
13
    const { tokenizerName, tokenizerId } = getFriendlyTokenizerName(main_api);
14
    const html = await renderExtensionTemplateAsync('token-counter', 'window', { tokenizerName });
15

16
    const dialog = $(html);
17
    const countDebounced = debounce(async () => {
18
        const text = String($('#token_counter_textarea').val());
19
        const ids = main_api == 'openai' ? getTextTokens(tokenizers.OPENAI, text) : getTextTokens(tokenizerId, text);
20

21
        if (Array.isArray(ids) && ids.length > 0) {
22
            $('#token_counter_ids').text(`[${ids.join(', ')}]`);
23
            $('#token_counter_result').text(ids.length);
24

25
            if (Object.hasOwnProperty.call(ids, 'chunks')) {
26
                drawChunks(Object.getOwnPropertyDescriptor(ids, 'chunks').value, ids);
27
            }
28
        } else {
29
            const count = await getTokenCountAsync(text);
30
            $('#token_counter_ids').text('—');
31
            $('#token_counter_result').text(count);
32
            $('#tokenized_chunks_display').text('—');
33
        }
34

35
        if (!CSS.supports('field-sizing', 'content')) {
36
            await resetScrollHeight($('#token_counter_textarea'));
37
            await resetScrollHeight($('#token_counter_ids'));
38
        }
39
    }, debounce_timeout.relaxed);
40
    dialog.find('#token_counter_textarea').on('input', () => countDebounced());
41

42
    callGenericPopup(dialog, POPUP_TYPE.TEXT, '', { wide: true, large: true, allowVerticalScrolling: true });
43
}
44

45
/**
46
 * Draws the tokenized chunks in the UI
47
 * @param {string[]} chunks
48
 * @param {number[]} ids
49
 */
50
function drawChunks(chunks, ids) {
51
    const pastelRainbow = [
52
        //main_text_color,
53
        //italics_text_color,
54
        //quote_text_color,
55
        '#FFB3BA',
56
        '#FFDFBA',
57
        '#FFFFBA',
58
        '#BFFFBF',
59
        '#BAE1FF',
60
        '#FFBAF3',
61
    ];
62
    $('#tokenized_chunks_display').empty();
63

64
    for (let i = 0; i < chunks.length; i++) {
65
        let chunk = chunks[i].replace(/[▁Ġ]/g, ' '); // This is a leading space in sentencepiece. More info: Lower one eighth block (U+2581)
66

67
        // If <0xHEX>, decode it
68
        if (/^<0x[0-9A-F]+>$/i.test(chunk)) {
69
            const code = parseInt(chunk.substring(3, chunk.length - 1), 16);
70
            chunk = String.fromCodePoint(code);
71
        }
72

73
        // If newline - insert a line break
74
        if (chunk === '\n') {
75
            $('#tokenized_chunks_display').append('<br>');
76
            continue;
77
        }
78

79
        const color = pastelRainbow[i % pastelRainbow.length];
80
        const chunkHtml = $('<code></code>');
81
        chunkHtml.css('background-color', color);
82
        chunkHtml.text(chunk);
83
        chunkHtml.attr('title', ids[i]);
84
        $('#tokenized_chunks_display').append(chunkHtml);
85
    }
86
}
87

88
async function doCount() {
89
    // get all of the messages in the chat
90
    const context = getContext();
91
    const messages = context.chat.filter(x => x.mes && !x.is_system).map(x => x.mes);
92

93
    //concat all the messages into a single string
94
    const allMessages = messages.join(' ');
95

96
    console.debug('All messages:', allMessages);
97

98
    //toastr success with the token count of the chat
99
    const count = await getTokenCountAsync(allMessages);
100
    toastr.success(`Token count: ${count}`);
101
    return count;
102
}
103

104
export function init() {
105
    const buttonHtml = `
106
        <div id="token_counter" class="list-group-item flex-container flexGap5">
107
            <div class="fa-solid fa-1 extensionsMenuExtensionButton" /></div>` +
108
            t`Token Counter` +
109
        '</div>';
110
    $('#token_counter_wand_container').append(buttonHtml);
111
    $('#token_counter').on('click', doTokenCounter);
112
    SlashCommandParser.addCommandObject(SlashCommand.fromProps({
113
        name: 'count',
114
        callback: async () => String(await doCount()),
115
        returns: 'number of tokens',
116
        helpString: 'Counts the number of tokens in the current chat.',
117
    }));
118
}