1
import {
2
    abortStatusCheck,
3
    event_types,
4
    eventSource,
5
    getRequestHeaders,
6
    getStoppingStrings,
7
    resultCheckStatus,
8
    saveSettingsDebounced,
9
    setGenerationParamsFromPreset,
10
    setOnlineStatus,
11
    startStatusLoading,
12
} from '../script.js';
13
import { MAX_CONTEXT_DEFAULT, MAX_RESPONSE_DEFAULT, power_user } from './power-user.js';
14
import { getTextTokens, tokenizers } from './tokenizers.js';
15
import { getEventSourceStream } from './sse-stream.js';
16
import {
17
    getSortableDelay,
18
    getStringHash,
19
    onlyUnique,
20
} from './utils.js';
21
import { BIAS_CACHE, createNewLogitBiasEntry, displayLogitBias, getLogitBiasListResult } from './logit-bias.js';
22
import { SECRET_KEYS, secret_state, writeSecret } from './secrets.js';
23

24
const default_preamble = '[ Style: chat, complex, sensory, visceral ]';
25
const default_order = [1, 5, 0, 2, 3, 4];
26
const maximum_output_length = 150;
27
const default_presets = {
28
    'clio-v1': 'Talker-Chat-Clio',
29
    'kayra-v1': 'Carefree-Kayra',
30
    'llama-3-erato-v1': 'Erato-Dragonfruit',
31
};
32

33
export let novelai_settings;
34
export let novelai_setting_names;
35

36
export const nai_settings = {
37
    temperature: 1.5,
38
    repetition_penalty: 2.25,
39
    repetition_penalty_range: 2048,
40
    repetition_penalty_slope: 0.09,
41
    repetition_penalty_frequency: 0,
42
    repetition_penalty_presence: 0.005,
43
    tail_free_sampling: 0.975,
44
    top_k: 10,
45
    top_p: 0.75,
46
    top_a: 0.08,
47
    typical_p: 0.975,
48
    min_p: 0,
49
    math1_temp: 1,
50
    math1_quad: 0,
51
    math1_quad_entropy_scale: 0,
52
    min_length: 1,
53
    model_novel: 'clio-v1',
54
    preset_settings_novel: 'Talker-Chat-Clio',
55
    streaming_novel: false,
56
    preamble: default_preamble,
57
    prefix: '',
58
    banned_tokens: '',
59
    order: default_order,
60
    logit_bias: [],
61
    extensions: {},
62
};
63

64
const nai_tiers = {
65
    0: 'Paper',
66
    1: 'Tablet',
67
    2: 'Scroll',
68
    3: 'Opus',
69
};
70

71
const samplers = {
72
    temperature: 0,
73
    top_k: 1,
74
    top_p: 2,
75
    tfs: 3,
76
    top_a: 4,
77
    typical_p: 5,
78
    // removed samplers were here
79
    mirostat: 8,
80
    math1: 9,
81
    min_p: 10,
82
};
83

84
let novel_data = null;
85
let badWordsCache = {};
86
const BIAS_KEY = '#range_block_novel';
87

88
export function setNovelData(data) {
89
    novel_data = data;
90
}
91

92
export function getKayraMaxContextTokens() {
93
    switch (novel_data?.tier) {
94
        case 1:
95
            return 4096;
96
        case 2:
97
            return 8192;
98
        case 3:
99
            return 8192;
100
    }
101

102
    return null;
103
}
104

105
export function getNovelMaxResponseTokens() {
106
    switch (novel_data?.tier) {
107
        case 1:
108
            return 150;
109
        case 2:
110
            return 150;
111
        case 3:
112
            return 250;
113
    }
114

115
    return maximum_output_length;
116
}
117

118
export function convertNovelPreset(data) {
119
    if (!data || typeof data !== 'object' || data.presetVersion !== 3 || !data.parameters || typeof data.parameters !== 'object') {
120
        return data;
121
    }
122

123
    return {
124
        max_context: 8000,
125
        temperature: data.parameters.temperature,
126
        max_length: data.parameters.max_length,
127
        min_length: data.parameters.min_length,
128
        top_k: data.parameters.top_k,
129
        top_p: data.parameters.top_p,
130
        top_a: data.parameters.top_a,
131
        typical_p: data.parameters.typical_p,
132
        tail_free_sampling: data.parameters.tail_free_sampling,
133
        repetition_penalty: data.parameters.repetition_penalty,
134
        repetition_penalty_range: data.parameters.repetition_penalty_range,
135
        repetition_penalty_slope: data.parameters.repetition_penalty_slope,
136
        repetition_penalty_frequency: data.parameters.repetition_penalty_frequency,
137
        repetition_penalty_presence: data.parameters.repetition_penalty_presence,
138
        phrase_rep_pen: data.parameters.phrase_rep_pen,
139
        mirostat_lr: data.parameters.mirostat_lr,
140
        mirostat_tau: data.parameters.mirostat_tau,
141
        math1_temp: data.parameters.math1_temp,
142
        math1_quad: data.parameters.math1_quad,
143
        math1_quad_entropy_scale: data.parameters.math1_quad_entropy_scale,
144
        min_p: data.parameters.min_p,
145
        order: Array.isArray(data.parameters.order) ? data.parameters.order.filter(s => s.enabled && Object.keys(samplers).includes(s.id)).map(s => samplers[s.id]) : default_order,
146
        extensions: {},
147
    };
148
}
149

150
export function getNovelTier() {
151
    return nai_tiers[novel_data?.tier] ?? 'no_connection';
152
}
153

154
export function getNovelAnlas() {
155
    return novel_data?.trainingStepsLeft?.fixedTrainingStepsLeft ?? 0;
156
}
157

158
export function getNovelUnlimitedImageGeneration() {
159
    return novel_data?.perks?.unlimitedImageGeneration ?? false;
160
}
161

162
export async function loadNovelSubscriptionData() {
163
    const result = await fetch('/api/novelai/status', {
164
        method: 'POST',
165
        headers: getRequestHeaders(),
166
        signal: abortStatusCheck.signal,
167
    });
168

169
    if (result.ok) {
170
        const data = await result.json();
171
        setNovelData(data);
172
    }
173

174
    return result.ok;
175
}
176

177
export function loadNovelPreset(preset) {
178
    if (preset.genamt === undefined) {
179
        const needsUnlock = preset.max_context > MAX_CONTEXT_DEFAULT || preset.max_length > MAX_RESPONSE_DEFAULT;
180
        $('#amount_gen').val(preset.max_length).trigger('input');
181
        $('#max_context_unlocked').prop('checked', needsUnlock).trigger('change');
182
        $('#max_context').val(preset.max_context).trigger('input');
183
    } else {
184
        setGenerationParamsFromPreset(preset);
185
    }
186

187
    nai_settings.temperature = preset.temperature;
188
    nai_settings.repetition_penalty = preset.repetition_penalty;
189
    nai_settings.repetition_penalty_range = preset.repetition_penalty_range;
190
    nai_settings.repetition_penalty_slope = preset.repetition_penalty_slope;
191
    nai_settings.repetition_penalty_frequency = preset.repetition_penalty_frequency;
192
    nai_settings.repetition_penalty_presence = preset.repetition_penalty_presence;
193
    nai_settings.tail_free_sampling = preset.tail_free_sampling;
194
    nai_settings.top_k = preset.top_k;
195
    nai_settings.top_p = preset.top_p;
196
    nai_settings.top_a = preset.top_a;
197
    nai_settings.typical_p = preset.typical_p;
198
    nai_settings.min_length = preset.min_length;
199
    nai_settings.phrase_rep_pen = preset.phrase_rep_pen;
200
    nai_settings.mirostat_lr = preset.mirostat_lr;
201
    nai_settings.mirostat_tau = preset.mirostat_tau;
202
    nai_settings.prefix = preset.prefix;
203
    nai_settings.banned_tokens = preset.banned_tokens || '';
204
    nai_settings.order = preset.order || default_order;
205
    nai_settings.logit_bias = preset.logit_bias || [];
206
    nai_settings.preamble = preset.preamble || default_preamble;
207
    nai_settings.min_p = preset.min_p || 0;
208
    nai_settings.math1_temp = preset.math1_temp || 1;
209
    nai_settings.math1_quad = preset.math1_quad || 0;
210
    nai_settings.math1_quad_entropy_scale = preset.math1_quad_entropy_scale || 0;
211
    nai_settings.extensions = preset.extensions || {};
212
    loadNovelSettingsUi(nai_settings);
213
}
214

215
export function loadNovelSettings(data, settings) {
216
    novelai_setting_names = data.novelai_setting_names;
217
    novelai_settings = data.novelai_settings;
218
    novelai_settings.forEach(function (item, i, arr) {
219
        novelai_settings[i] = JSON.parse(item);
220
    });
221

222
    $('#settings_preset_novel').empty();
223
    const presetNames = {};
224
    novelai_setting_names.forEach(function (item, i, arr) {
225
        presetNames[item] = i;
226
        $('#settings_preset_novel').append(`<option value=${i}>${item}</option>`);
227
    });
228
    novelai_setting_names = presetNames;
229

230
    //load the rest of the Novel settings without any checks
231
    nai_settings.model_novel = settings.model_novel;
232
    $('#model_novel_select').val(nai_settings.model_novel);
233
    $(`#model_novel_select option[value=${nai_settings.model_novel}]`).prop('selected', true);
234

235
    if (settings.nai_preamble !== undefined) {
236
        nai_settings.preamble = settings.nai_preamble;
237
        delete settings.nai_preamble;
238
    }
239
    nai_settings.preset_settings_novel = settings.preset_settings_novel;
240
    nai_settings.temperature = settings.temperature;
241
    nai_settings.repetition_penalty = settings.repetition_penalty;
242
    nai_settings.repetition_penalty_range = settings.repetition_penalty_range;
243
    nai_settings.repetition_penalty_slope = settings.repetition_penalty_slope;
244
    nai_settings.repetition_penalty_frequency = settings.repetition_penalty_frequency;
245
    nai_settings.repetition_penalty_presence = settings.repetition_penalty_presence;
246
    nai_settings.tail_free_sampling = settings.tail_free_sampling;
247
    nai_settings.top_k = settings.top_k;
248
    nai_settings.top_p = settings.top_p;
249
    nai_settings.top_a = settings.top_a;
250
    nai_settings.typical_p = settings.typical_p;
251
    nai_settings.min_length = settings.min_length;
252
    nai_settings.phrase_rep_pen = settings.phrase_rep_pen;
253
    nai_settings.mirostat_lr = settings.mirostat_lr;
254
    nai_settings.mirostat_tau = settings.mirostat_tau;
255
    nai_settings.streaming_novel = !!settings.streaming_novel;
256
    nai_settings.preamble = settings.preamble || default_preamble;
257
    nai_settings.prefix = settings.prefix;
258
    nai_settings.banned_tokens = settings.banned_tokens || '';
259
    nai_settings.order = settings.order || default_order;
260
    nai_settings.logit_bias = settings.logit_bias || [];
261
    nai_settings.min_p = settings.min_p || 0;
262
    nai_settings.math1_temp = settings.math1_temp || 1;
263
    nai_settings.math1_quad = settings.math1_quad || 0;
264
    nai_settings.math1_quad_entropy_scale = settings.math1_quad_entropy_scale || 0;
265
    nai_settings.extensions = settings.extensions || {};
266
    loadNovelSettingsUi(nai_settings);
267
}
268

269
function loadNovelSettingsUi(ui_settings) {
270
    $('#temp_novel').val(ui_settings.temperature);
271
    $('#temp_counter_novel').val(Number(ui_settings.temperature).toFixed(2));
272
    $('#rep_pen_novel').val(ui_settings.repetition_penalty);
273
    $('#rep_pen_counter_novel').val(Number(ui_settings.repetition_penalty).toFixed(3));
274
    $('#rep_pen_size_novel').val(ui_settings.repetition_penalty_range);
275
    $('#rep_pen_size_counter_novel').val(Number(ui_settings.repetition_penalty_range).toFixed(0));
276
    $('#rep_pen_slope_novel').val(ui_settings.repetition_penalty_slope);
277
    $('#rep_pen_slope_counter_novel').val(Number(`${ui_settings.repetition_penalty_slope}`).toFixed(2));
278
    $('#rep_pen_freq_novel').val(ui_settings.repetition_penalty_frequency);
279
    $('#rep_pen_freq_counter_novel').val(Number(ui_settings.repetition_penalty_frequency).toFixed(3));
280
    $('#rep_pen_presence_novel').val(ui_settings.repetition_penalty_presence);
281
    $('#rep_pen_presence_counter_novel').val(Number(ui_settings.repetition_penalty_presence).toFixed(3));
282
    $('#tail_free_sampling_novel').val(ui_settings.tail_free_sampling);
283
    $('#tail_free_sampling_counter_novel').val(Number(ui_settings.tail_free_sampling).toFixed(3));
284
    $('#top_k_novel').val(ui_settings.top_k);
285
    $('#top_k_counter_novel').val(Number(ui_settings.top_k).toFixed(0));
286
    $('#top_p_novel').val(ui_settings.top_p);
287
    $('#top_p_counter_novel').val(Number(ui_settings.top_p).toFixed(3));
288
    $('#top_a_novel').val(ui_settings.top_a);
289
    $('#top_a_counter_novel').val(Number(ui_settings.top_a).toFixed(3));
290
    $('#typical_p_novel').val(ui_settings.typical_p);
291
    $('#typical_p_counter_novel').val(Number(ui_settings.typical_p).toFixed(3));
292
    $('#phrase_rep_pen_novel').val(ui_settings.phrase_rep_pen || 'off');
293
    $('#mirostat_lr_novel').val(ui_settings.mirostat_lr);
294
    $('#mirostat_lr_counter_novel').val(Number(ui_settings.mirostat_lr).toFixed(2));
295
    $('#mirostat_tau_novel').val(ui_settings.mirostat_tau);
296
    $('#mirostat_tau_counter_novel').val(Number(ui_settings.mirostat_tau).toFixed(2));
297
    $('#min_length_novel').val(ui_settings.min_length);
298
    $('#min_length_counter_novel').val(Number(ui_settings.min_length).toFixed(0));
299
    $('#nai_preamble_textarea').val(ui_settings.preamble);
300
    $('#nai_prefix').val(ui_settings.prefix || 'vanilla');
301
    $('#nai_banned_tokens').val(ui_settings.banned_tokens || '');
302
    $('#min_p_novel').val(ui_settings.min_p);
303
    $('#min_p_counter_novel').val(Number(ui_settings.min_p).toFixed(3));
304
    $('#math1_temp_novel').val(ui_settings.math1_temp);
305
    $('#math1_temp_counter_novel').val(Number(ui_settings.math1_temp).toFixed(2));
306
    $('#math1_quad_novel').val(ui_settings.math1_quad);
307
    $('#math1_quad_counter_novel').val(Number(ui_settings.math1_quad).toFixed(2));
308
    $('#math1_quad_entropy_scale_novel').val(ui_settings.math1_quad_entropy_scale);
309
    $('#math1_quad_entropy_scale_counter_novel').val(Number(ui_settings.math1_quad_entropy_scale).toFixed(2));
310
    $(`#settings_preset_novel option[value=${novelai_setting_names[nai_settings.preset_settings_novel]}]`).prop('selected', true);
311

312
    $('#streaming_novel').prop('checked', ui_settings.streaming_novel);
313
    sortItemsByOrder(ui_settings.order);
314
    displayLogitBias(ui_settings.logit_bias, BIAS_KEY);
315
}
316

317
const sliders = [
318
    {
319
        sliderId: '#temp_novel',
320
        counterId: '#temp_counter_novel',
321
        format: (val) => Number(val).toFixed(2),
322
        setValue: (val) => { nai_settings.temperature = Number(val); },
323
    },
324
    {
325
        sliderId: '#rep_pen_novel',
326
        counterId: '#rep_pen_counter_novel',
327
        format: (val) => Number(val).toFixed(3),
328
        setValue: (val) => { nai_settings.repetition_penalty = Number(val); },
329
    },
330
    {
331
        sliderId: '#rep_pen_size_novel',
332
        counterId: '#rep_pen_size_counter_novel',
333
        format: (val) => `${val}`,
334
        setValue: (val) => { nai_settings.repetition_penalty_range = Number(val); },
335
    },
336
    {
337
        sliderId: '#rep_pen_slope_novel',
338
        counterId: '#rep_pen_slope_counter_novel',
339
        format: (val) => `${val}`,
340
        setValue: (val) => { nai_settings.repetition_penalty_slope = Number(val); },
341
    },
342
    {
343
        sliderId: '#rep_pen_freq_novel',
344
        counterId: '#rep_pen_freq_counter_novel',
345
        format: (val) => Number(val).toFixed(2),
346
        setValue: (val) => { nai_settings.repetition_penalty_frequency = Number(val); },
347
    },
348
    {
349
        sliderId: '#rep_pen_presence_novel',
350
        counterId: '#rep_pen_presence_counter_novel',
351
        format: (val) => `${val}`,
352
        setValue: (val) => { nai_settings.repetition_penalty_presence = Number(val); },
353
    },
354
    {
355
        sliderId: '#tail_free_sampling_novel',
356
        counterId: '#tail_free_sampling_counter_novel',
357
        format: (val) => `${val}`,
358
        setValue: (val) => { nai_settings.tail_free_sampling = Number(val); },
359
    },
360
    {
361
        sliderId: '#top_k_novel',
362
        counterId: '#top_k_counter_novel',
363
        format: (val) => `${val}`,
364
        setValue: (val) => { nai_settings.top_k = Number(val); },
365
    },
366
    {
367
        sliderId: '#top_p_novel',
368
        counterId: '#top_p_counter_novel',
369
        format: (val) => Number(val).toFixed(3),
370
        setValue: (val) => { nai_settings.top_p = Number(val); },
371
    },
372
    {
373
        sliderId: '#top_a_novel',
374
        counterId: '#top_a_counter_novel',
375
        format: (val) => Number(val).toFixed(2),
376
        setValue: (val) => { nai_settings.top_a = Number(val); },
377
    },
378
    {
379
        sliderId: '#typical_p_novel',
380
        counterId: '#typical_p_counter_novel',
381
        format: (val) => Number(val).toFixed(3),
382
        setValue: (val) => { nai_settings.typical_p = Number(val); },
383
    },
384
    {
385
        sliderId: '#mirostat_tau_novel',
386
        counterId: '#mirostat_tau_counter_novel',
387
        format: (val) => Number(val).toFixed(2),
388
        setValue: (val) => { nai_settings.mirostat_tau = Number(val); },
389
    },
390
    {
391
        sliderId: '#mirostat_lr_novel',
392
        counterId: '#mirostat_lr_counter_novel',
393
        format: (val) => Number(val).toFixed(2),
394
        setValue: (val) => { nai_settings.mirostat_lr = Number(val); },
395
    },
396
    {
397
        sliderId: '#min_length_novel',
398
        counterId: '#min_length_counter_novel',
399
        format: (val) => `${val}`,
400
        setValue: (val) => { nai_settings.min_length = Number(val); },
401
    },
402
    {
403
        sliderId: '#nai_banned_tokens',
404
        counterId: '#nai_banned_tokens_counter',
405
        format: (val) => val,
406
        setValue: (val) => { nai_settings.banned_tokens = val; },
407
    },
408
    {
409
        sliderId: '#min_p_novel',
410
        counterId: '#min_p_counter_novel',
411
        format: (val) => Number(val).toFixed(3),
412
        setValue: (val) => { nai_settings.min_p = Number(val); },
413
    },
414
    {
415
        sliderId: '#math1_temp_novel',
416
        counterId: '#math1_temp_counter_novel',
417
        format: (val) => Number(val).toFixed(2),
418
        setValue: (val) => { nai_settings.math1_temp = Number(val); },
419
    },
420
    {
421
        sliderId: '#math1_quad_novel',
422
        counterId: '#math1_quad_counter_novel',
423
        format: (val) => Number(val).toFixed(2),
424
        setValue: (val) => { nai_settings.math1_quad = Number(val); },
425
    },
426
    {
427
        sliderId: '#math1_quad_entropy_scale_novel',
428
        counterId: '#math1_quad_entropy_scale_counter_novel',
429
        format: (val) => Number(val).toFixed(2),
430
        setValue: (val) => { nai_settings.math1_quad_entropy_scale = Number(val); },
431
    },
432
];
433

434
function getBadWordIds(banned_tokens, tokenizerType) {
435
    if (tokenizerType === tokenizers.NONE) {
436
        return [];
437
    }
438

439
    const cacheKey = `${getStringHash(banned_tokens)}-${tokenizerType}`;
440

441
    if (cacheKey in badWordsCache && Array.isArray(badWordsCache[cacheKey])) {
442
        console.debug(`Bad words ids cache hit for "${banned_tokens}"`, badWordsCache[cacheKey]);
443
        return badWordsCache[cacheKey];
444
    }
445

446
    const result = [];
447
    const sequence = banned_tokens.split('\n');
448

449
    for (let token of sequence) {
450
        const trimmed = token.trim();
451

452
        // Skip empty lines
453
        if (trimmed.length === 0) {
454
            continue;
455
        }
456

457
        // Verbatim text
458
        if (trimmed.startsWith('{') && trimmed.endsWith('}')) {
459
            const tokens = getTextTokens(tokenizerType, trimmed.slice(1, -1));
460
            result.push(tokens);
461
        } else if (trimmed.startsWith('[') && trimmed.endsWith(']')) {
462
            // Raw token ids, JSON serialized
463
            try {
464
                const tokens = JSON.parse(trimmed);
465

466
                if (Array.isArray(tokens) && tokens.every(t => Number.isInteger(t))) {
467
                    result.push(tokens);
468
                } else {
469
                    throw new Error('Not an array of integers');
470
                }
471
            } catch (err) {
472
                console.log(`Failed to parse bad word token list: ${trimmed}`, err);
473
            }
474
        } else {
475
            // Apply permutations
476
            const permutations = getBadWordPermutations(trimmed).map(t => getTextTokens(tokenizerType, t));
477
            result.push(...permutations);
478
        }
479
    }
480

481
    // Cache the result
482
    console.debug(`Bad words ids for "${banned_tokens}"`, result);
483
    badWordsCache[cacheKey] = result;
484

485
    return result;
486
}
487

488
function getBadWordPermutations(text) {
489
    const result = [];
490

491
    // Original text
492
    result.push(text);
493
    // Original text + leading space
494
    result.push(` ${text}`);
495
    // First letter capitalized
496
    result.push(text[0].toUpperCase() + text.slice(1));
497
    // Ditto + leading space
498
    result.push(` ${text[0].toUpperCase() + text.slice(1)}`);
499
    // First letter lower cased
500
    result.push(text[0].toLowerCase() + text.slice(1));
501
    // Ditto + leading space
502
    result.push(` ${text[0].toLowerCase() + text.slice(1)}`);
503
    // Original all upper cased
504
    result.push(text.toUpperCase());
505
    // Ditto + leading space
506
    result.push(` ${text.toUpperCase()}`);
507
    // Original all lower cased
508
    result.push(text.toLowerCase());
509
    // Ditto + leading space
510
    result.push(` ${text.toLowerCase()}`);
511

512
    return result.filter(onlyUnique);
513
}
514

515
export function getNovelGenerationData(finalPrompt, settings, maxLength, isImpersonate, isContinue, _cfgValues, type) {
516
    console.debug('NovelAI generation data for', type);
517
    const isKayra = nai_settings.model_novel.includes('kayra');
518
    const isErato = nai_settings.model_novel.includes('erato');
519

520
    const tokenizerType = getTokenizerTypeForModel(nai_settings.model_novel);
521
    const stoppingStrings = getStoppingStrings(isImpersonate, isContinue);
522

523
    // Llama 3 tokenizer, huh?
524
    if (isErato) {
525
        const additionalStopStrings = [];
526
        for (const stoppingString of stoppingStrings) {
527
            if (stoppingString.startsWith('\n')) {
528
                additionalStopStrings.push('.' + stoppingString);
529
                additionalStopStrings.push('!' + stoppingString);
530
                additionalStopStrings.push('?' + stoppingString);
531
                additionalStopStrings.push('*' + stoppingString);
532
                additionalStopStrings.push('"' + stoppingString);
533
                additionalStopStrings.push('_' + stoppingString);
534
                additionalStopStrings.push('...' + stoppingString);
535
                additionalStopStrings.push('."' + stoppingString);
536
                additionalStopStrings.push('?"' + stoppingString);
537
                additionalStopStrings.push('!"' + stoppingString);
538
                additionalStopStrings.push('.*' + stoppingString);
539
                additionalStopStrings.push(')' + stoppingString);
540
            }
541
        }
542
        stoppingStrings.push(...additionalStopStrings);
543
    }
544

545
    const MAX_STOP_SEQUENCES = 1024;
546
    const stopSequences = (tokenizerType !== tokenizers.NONE)
547
        ? stoppingStrings.slice(0, MAX_STOP_SEQUENCES).map(t => getTextTokens(tokenizerType, t))
548
        : undefined;
549

550
    const badWordIds = (tokenizerType !== tokenizers.NONE)
551
        ? getBadWordIds(nai_settings.banned_tokens, tokenizerType)
552
        : undefined;
553

554
    const prefix = selectPrefix(nai_settings.prefix, finalPrompt);
555

556
    let logitBias = [];
557
    if (tokenizerType !== tokenizers.NONE && Array.isArray(nai_settings.logit_bias) && nai_settings.logit_bias.length) {
558
        logitBias = BIAS_CACHE.get(BIAS_KEY) || calculateLogitBias();
559
        BIAS_CACHE.set(BIAS_KEY, logitBias);
560
    }
561

562
    if (power_user.console_log_prompts) {
563
        console.log(finalPrompt);
564
    }
565

566

567
    if (isErato) {
568
        finalPrompt = '<|startoftext|><|reserved_special_token81|>' + finalPrompt;
569
    }
570

571
    const adjustedMaxLength = (isKayra || isErato) ? getNovelMaxResponseTokens() : maximum_output_length;
572

573
    return {
574
        'input': finalPrompt,
575
        'model': nai_settings.model_novel,
576
        'use_string': true,
577
        'temperature': Number(nai_settings.temperature),
578
        'max_length': maxLength < adjustedMaxLength ? maxLength : adjustedMaxLength,
579
        'min_length': Number(nai_settings.min_length),
580
        'tail_free_sampling': Number(nai_settings.tail_free_sampling),
581
        'repetition_penalty': Number(nai_settings.repetition_penalty),
582
        'repetition_penalty_range': Number(nai_settings.repetition_penalty_range),
583
        'repetition_penalty_slope': Number(nai_settings.repetition_penalty_slope),
584
        'repetition_penalty_frequency': Number(nai_settings.repetition_penalty_frequency),
585
        'repetition_penalty_presence': Number(nai_settings.repetition_penalty_presence),
586
        'top_a': Number(nai_settings.top_a),
587
        'top_p': Number(nai_settings.top_p),
588
        'top_k': Number(nai_settings.top_k),
589
        'min_p': Number(nai_settings.min_p),
590
        'math1_temp': Number(nai_settings.math1_temp),
591
        'math1_quad': Number(nai_settings.math1_quad),
592
        'math1_quad_entropy_scale': Number(nai_settings.math1_quad_entropy_scale),
593
        'typical_p': Number(nai_settings.typical_p),
594
        'mirostat_lr': Number(nai_settings.mirostat_lr),
595
        'mirostat_tau': Number(nai_settings.mirostat_tau),
596
        'phrase_rep_pen': nai_settings.phrase_rep_pen,
597
        'stop_sequences': stopSequences,
598
        'bad_words_ids': badWordIds,
599
        'logit_bias_exp': logitBias,
600
        'generate_until_sentence': true,
601
        'use_cache': false,
602
        'return_full_text': false,
603
        'prefix': prefix,
604
        'order': nai_settings.order || settings.order || default_order,
605
        'num_logprobs': power_user.request_token_probabilities ? 10 : undefined,
606
    };
607
}
608

609
// Check if the prefix needs to be overridden to use instruct mode
610
function selectPrefix(selected_prefix, finalPrompt) {
611
    let useInstruct = false;
612
    const clio = nai_settings.model_novel.includes('clio');
613
    const kayra = nai_settings.model_novel.includes('kayra');
614
    const erato = nai_settings.model_novel.includes('erato');
615
    const isNewModel = clio || kayra || erato;
616

617
    if (isNewModel) {
618
        // NovelAI claims they scan backwards 1000 characters (not tokens!) to look for instruct brackets. That's really short.
619
        const tail = finalPrompt.slice(-1500);
620
        useInstruct = tail.includes('}');
621
        return useInstruct ? 'special_instruct' : selected_prefix;
622
    }
623

624
    return 'vanilla';
625
}
626

627
function getTokenizerTypeForModel(model) {
628
    if (model.includes('clio')) {
629
        return tokenizers.NERD;
630
    }
631
    if (model.includes('kayra')) {
632
        return tokenizers.NERD2;
633
    }
634
    if (model.includes('erato')) {
635
        return tokenizers.LLAMA3;
636
    }
637
    return tokenizers.NONE;
638
}
639

640
// Sort the samplers by the order array
641
function sortItemsByOrder(orderArray) {
642
    console.debug('Preset samplers order: ' + orderArray);
643
    const $draggableItems = $('#novel_order');
644

645
    // Sort the items by the order array
646
    for (let i = 0; i < orderArray.length; i++) {
647
        const index = orderArray[i];
648
        const $item = $draggableItems.find(`[data-id="${index}"]`).detach();
649
        $draggableItems.append($item);
650
    }
651

652
    // Update the disabled class for each sampler
653
    $draggableItems.children().each(function () {
654
        const isEnabled = orderArray.includes(parseInt($(this).data('id')));
655
        $(this).toggleClass('disabled', !isEnabled);
656

657
        // If the sampler is disabled, move it to the bottom of the list
658
        if (!isEnabled) {
659
            const item = $(this).detach();
660
            $draggableItems.append(item);
661
        }
662
    });
663
}
664

665
function saveSamplingOrder() {
666
    const order = [];
667
    $('#novel_order').children().each(function () {
668
        const isEnabled = !$(this).hasClass('disabled');
669
        if (isEnabled) {
670
            order.push($(this).data('id'));
671
        }
672
    });
673
    nai_settings.order = order;
674
    console.log('Samplers reordered:', nai_settings.order);
675
    saveSettingsDebounced();
676
}
677

678
/**
679
 * Calculates logit bias for Novel AI
680
 * @returns {object[]} Array of logit bias objects
681
 */
682
function calculateLogitBias() {
683
    const biasPreset = nai_settings.logit_bias;
684

685
    if (!Array.isArray(biasPreset) || biasPreset.length === 0) {
686
        return [];
687
    }
688

689
    const tokenizerType = getTokenizerTypeForModel(nai_settings.model_novel);
690

691
    /**
692
     * Creates a bias object for Novel AI
693
     * @param {number} bias Bias value
694
     * @param {number[]} sequence Sequence of token ids
695
     */
696
    function getBiasObject(bias, sequence) {
697
        return {
698
            bias: bias,
699
            ensure_sequence_finish: false,
700
            generate_once: false,
701
            sequence: sequence,
702
        };
703
    }
704

705
    const result = getLogitBiasListResult(biasPreset, tokenizerType, getBiasObject);
706
    return result;
707
}
708

709
/**
710
 * Transforms instruction into compatible format for Novel AI if Novel AI instruct format not already detected.
711
 * 1. Instruction must begin and end with curly braces followed and preceded by a space.
712
 * 2. Instruction must not contain square brackets as it serves different purpose in NAI.
713
 * @param {string} prompt Original instruction prompt
714
 * @returns Processed prompt
715
 */
716
export function adjustNovelInstructionPrompt(prompt) {
717
    const stripedPrompt = prompt.replace(/[[\]]/g, '').trim();
718
    if (!stripedPrompt.includes('{ ')) {
719
        return `{ ${stripedPrompt} }`;
720
    }
721
    return stripedPrompt;
722
}
723

724
function tryParseStreamingError(response, decoded) {
725
    try {
726
        const data = JSON.parse(decoded);
727

728
        if (!data) {
729
            return;
730
        }
731

732
        if (data.message || data.error) {
733
            toastr.error(data.message || data.error?.message || response.statusText, 'NovelAI API');
734
            throw new Error(data);
735
        }
736
    } catch {
737
        // No JSON. Do nothing.
738
    }
739
}
740

741
export async function generateNovelWithStreaming(generate_data, signal) {
742
    generate_data.streaming = nai_settings.streaming_novel;
743

744
    const response = await fetch('/api/novelai/generate', {
745
        headers: getRequestHeaders(),
746
        body: JSON.stringify(generate_data),
747
        method: 'POST',
748
        signal: signal,
749
    });
750
    if (!response.ok) {
751
        tryParseStreamingError(response, await response.text());
752
        throw new Error(`Got response status ${response.status}`);
753
    }
754
    const eventStream = getEventSourceStream();
755
    response.body.pipeThrough(eventStream);
756
    const reader = eventStream.readable.getReader();
757

758
    return async function* streamData() {
759
        let text = '';
760
        while (true) {
761
            const { done, value } = await reader.read();
762
            if (done) return;
763

764
            const data = JSON.parse(value.data);
765

766
            if (data.token) {
767
                text += data.token;
768
            }
769

770
            yield { text, swipes: [], logprobs: parseNovelAILogprobs(data.logprobs), toolCalls: [], state: {} };
771
        }
772
    };
773
}
774

775
/**
776
 * A single token's ID.
777
 * @typedef {[number]} TokenIdEntry
778
 */
779
/**
780
 * A single token's log probabilities. The first element is before repetition
781
 * penalties and samplers are applied, the second is after.
782
 * @typedef {[number, number]} LogprobsEntry
783
 */
784
/**
785
 * Combination of token ID and its corresponding log probabilities.
786
 * @typedef {[TokenIdEntry, LogprobsEntry]} TokenLogprobTuple
787
 */
788
/**
789
 * Represents all logprob data for a single token, including its
790
 * before, after, and the ultimately selected token.
791
 * @typedef {Object} NAITokenLogprobs
792
 * @property {TokenLogprobTuple[]} chosen - always length 1
793
 * @property {TokenLogprobTuple[]} before - always `top_logprobs` length
794
 * @property {TokenLogprobTuple[]} after - maybe less than `top_logprobs` length
795
 */
796
/**
797
 * parseNovelAILogprobs converts a logprobs object returned from the NovelAI API
798
 * for a single token into a TokenLogprobs object used by the Token Probabilities
799
 * feature.
800
 * @param {NAITokenLogprobs} data - NAI logprobs object for one token
801
 * @returns {import('./logprobs.js').TokenLogprobs | null} converted logprobs
802
 */
803
export function parseNovelAILogprobs(data) {
804
    if (!data) {
805
        return null;
806
    }
807
    const befores = data.before.map(([[tokenId], [before, _]]) => [tokenId, before]);
808
    const afters = data.after.map(([[tokenId], [_, after]]) => [tokenId, after]);
809

810
    // Find any tokens in `befores` that are missing from `afters`. Then add
811
    // them with a logprob of -Infinity (0% probability)
812
    const notInAfter = befores
813
        .filter(([id]) => !afters.some(([aid]) => aid === id))
814
        .map(([id]) => [id, -Infinity]);
815
    const merged = afters.concat(notInAfter);
816

817
    // Add the chosen token to `merged` if it's not already there. This can
818
    // happen if the chosen token was not among the top 10 most likely ones.
819
    // eslint-disable-next-line no-unused-vars
820
    const [[chosenId], [_, chosenAfter]] = data.chosen[0];
821
    if (!merged.some(([id]) => id === chosenId)) {
822
        merged.push([chosenId, chosenAfter]);
823
    }
824

825
    // nb: returned logprobs are provided alongside token IDs, not decoded text.
826
    // We don't want to send an API call for every streaming tick to decode the
827
    // text so we will use the IDs instead and bulk decode them in
828
    // StreamingProcessor. JSDoc typechecking may complain about this, but it's
829
    // intentional.
830
    // @ts-ignore
831
    return { token: chosenId, topLogprobs: merged };
832
}
833

834
$('#nai_preamble_textarea').on('input', function () {
835
    nai_settings.preamble = String($('#nai_preamble_textarea').val());
836
    saveSettingsDebounced();
837
});
838

839
$('#nai_preamble_restore').on('click', function () {
840
    nai_settings.preamble = default_preamble;
841
    $('#nai_preamble_textarea').val(nai_settings.preamble);
842
    saveSettingsDebounced();
843
});
844

845
export async function getStatusNovel() {
846
    try {
847
        const result = await loadNovelSubscriptionData();
848

849
        if (!result) {
850
            throw new Error('Could not load subscription data');
851
        }
852

853
        setOnlineStatus(getNovelTier());
854
    } catch {
855
        setOnlineStatus('no_connection');
856
    }
857

858
    return resultCheckStatus();
859
}
860

861
export function initNovelAISettings() {
862
    sliders.forEach(slider => {
863
        $(document).on('input', slider.sliderId, function () {
864
            const value = $(this).val();
865
            const formattedValue = slider.format(value);
866
            slider.setValue(value);
867
            $(slider.counterId).val(formattedValue);
868
            saveSettingsDebounced();
869
        });
870
    });
871

872
    $('#api_button_novel').on('click', async function (e) {
873
        e.stopPropagation();
874
        const api_key_novel = String($('#api_key_novel').val()).trim();
875

876
        if (api_key_novel.length) {
877
            await writeSecret(SECRET_KEYS.NOVEL, api_key_novel);
878
        }
879

880
        if (!secret_state[SECRET_KEYS.NOVEL]) {
881
            console.log('No secret key saved for NovelAI');
882
            return;
883
        }
884

885
        startStatusLoading();
886
        await getStatusNovel();
887
    });
888

889
    $('#settings_preset_novel').on('change', async function () {
890
        nai_settings.preset_settings_novel = $('#settings_preset_novel').find(':selected').text();
891
        const preset = novelai_settings[novelai_setting_names[nai_settings.preset_settings_novel]];
892
        loadNovelPreset(preset);
893
        saveSettingsDebounced();
894
        await eventSource.emit(event_types.PRESET_CHANGED, { apiId: 'novel', name: nai_settings.preset_settings_novel });
895
    });
896

897
    $('#streaming_novel').on('input', function () {
898
        const value = !!$(this).prop('checked');
899
        nai_settings.streaming_novel = value;
900
        saveSettingsDebounced();
901
    });
902

903
    $('#model_novel_select').on('change', function () {
904
        nai_settings.model_novel = String($('#model_novel_select').find(':selected').val());
905
        saveSettingsDebounced();
906

907
        // Update the selected preset to something appropriate
908
        const default_preset = default_presets[nai_settings.model_novel];
909
        $('#settings_preset_novel').val(novelai_setting_names[default_preset]);
910
        $(`#settings_preset_novel option[value=${novelai_setting_names[default_preset]}]`).attr('selected', 'true');
911
        $('#settings_preset_novel').trigger('change');
912
    });
913

914
    $('#nai_prefix').on('change', function () {
915
        nai_settings.prefix = String($('#nai_prefix').find(':selected').val());
916
        saveSettingsDebounced();
917
    });
918

919
    $('#phrase_rep_pen_novel').on('change', function () {
920
        nai_settings.phrase_rep_pen = String($('#phrase_rep_pen_novel').find(':selected').val());
921
        saveSettingsDebounced();
922
    });
923

924
    $('#novel_order').sortable({
925
        delay: getSortableDelay(),
926
        stop: saveSamplingOrder,
927
    });
928

929
    $('#novel_order .toggle_button').on('click', function () {
930
        const $item = $(this).closest('[data-id]');
931
        const isEnabled = !$item.hasClass('disabled');
932
        $item.toggleClass('disabled', isEnabled);
933
        console.log('Sampler toggled:', $item.data('id'), !isEnabled);
934
        saveSamplingOrder();
935
    });
936

937
    $('#novelai_logit_bias_new_entry').on('click', () => createNewLogitBiasEntry(nai_settings.logit_bias, BIAS_KEY));
938
}