From 722801bb5002e43b848f577c1af9972b7157fe14 Mon Sep 17 00:00:00 2001
From: Mike Weldon <mweldon@yahoo.com>
Date: Wed, 16 Aug 2023 22:14:04 -0700
Subject: [PATCH] Add logit biases for NovelAI

---
 public/index.html              |  47 ++++++-
 public/script.js               |   2 +-
 public/scripts/nai-settings.js | 218 ++++++++++++++++++++++++++++++++-
 public/scripts/openai.js       |   2 +-
 public/style.css               |  35 ++++++
 server.js                      |   8 +-
 6 files changed, 305 insertions(+), 7 deletions(-)

diff --git a/public/index.html b/public/index.html
index 233135e62..aae8629a6 100644
--- a/public/index.html
+++ b/public/index.html
@@ -857,7 +857,7 @@
                                                 <small>3</small>
                                             </div>
                                             <div data-id="4">
-                                                <span data-i18n="Typical Sampling">Typical Sampling</span>
+                                                <span data-i18n="Typical Sampling">Typical P Sampling</span>
                                                 <small>4</small>
                                             </div>
                                             <div data-id="5">
@@ -901,6 +901,40 @@
                                         <textarea id="nai_banned_tokens" class="text_pole textarea_compact" name="nai_banned_tokens" rows="2" placeholder=""></textarea>
                                     </div>
                                 </div>
+
+                                <div class="range-block-title openai_restorable" data-i18n="Logit Bias">
+                                    Logit Bias
+                                </div>
+                                <div class="toggle-description justifyLeft" data-i18n="Helps to ban or reenforce the usage of certain words">
+                                    Helps to ban or reinforce the usage of certain tokens.
+                                </div>
+                                <div class="openai_logit_bias_preset_form">
+                                    <select id="novelai_logit_bias_preset">
+                                    </select>
+                                    <i title="New preset" id="novelai_logit_bias_new_preset" class="menu_button fa-solid fa-plus"
+                                        data-i18n="[title]New preset"></i>
+                                    <i title="Import preset" id="novelai_logit_bias_import_preset" class="menu_button fa-solid fa-upload"
+                                        data-i18n="[title]Import preset"></i>
+                                    <i title="Export preset" id="novelai_logit_bias_export_preset" class="menu_button fa-solid fa-download"
+                                        data-i18n="[title]Export preset"></i>
+                                    <i title="Delete preset" id="novelai_logit_bias_delete_preset" class="menu_button fa-solid fa-trash-can"
+                                        data-i18n="[title]Delete preset"></i>
+                                    <input id="novelai_logit_bias_import_file" type="file" accept=".json" hidden />
+                                </div>
+                                <div class="inline-drawer wide100p">
+                                    <div class="inline-drawer-toggle inline-drawer-header">
+                                        <b data-i18n="View / Edit bias preset">View / Edit bias preset</b>
+                                        <div class="fa-solid fa-circle-chevron-down inline-drawer-icon down"></div>
+                                    </div>
+                                    <div class="inline-drawer-content">
+                                        <div id="novelai_logit_bias_new_entry" class="menu_button wide100p flex-container justifyCenter"
+                                            data-i18n="Add bias entry">
+                                            Add bias entry
+                                        </div>
+                                        <div class="novelai_logit_bias_list"></div>
+                                    </div>
+                                </div>
+
                                 <hr>
                                 <div class="range-block">
                                     <div class="range-block-title justifyLeft" data-i18n="CFG Scale">
@@ -3783,6 +3817,17 @@
             </div>
         </div>
     </div>
+
+    <div id="novelai_logit_bias_template" class="template_element">
+        <div class="novelai_logit_bias_form">
+            <input class="novelai_logit_bias_text text_pole" data-i18n="[placeholder]Type here..." placeholder="type here..." />
+            <input class="novelai_logit_bias_value text_pole" type="number" min="-2" value="0" max="2" step="0.01" />
+            <i class="menu_button fa-solid fa-xmark novelai_logit_bias_remove"></i>
+                </form>
+            </div>
+        </div>
+    </div>
+    
     <div id="completion_prompt_manager_popup" class="drawer-content" style="display:none;">
         <div id="completion_prompt_manager_popup_inspect">
             <h3>Inspect</h3>
diff --git a/public/script.js b/public/script.js
index 08d2d2c7d..39015042c 100644
--- a/public/script.js
+++ b/public/script.js
@@ -2899,7 +2899,7 @@ async function Generate(type, { automatic_trigger, force_name2, resolve, reject,
             }
             else if (main_api == 'novel') {
                 const this_settings = novelai_settings[novelai_setting_names[nai_settings.preset_settings_novel]];
-                generate_data = getNovelGenerationData(finalPromt, this_settings, this_amount_gen, isImpersonate);
+                generate_data = await getNovelGenerationData(finalPromt, this_settings, this_amount_gen, isImpersonate);
             }
             else if (main_api == 'openai') {
                 let [prompt, counts] = prepareOpenAIMessages({
diff --git a/public/scripts/nai-settings.js b/public/scripts/nai-settings.js
index c93c6366c..270ff2aa9 100644
--- a/public/scripts/nai-settings.js
+++ b/public/scripts/nai-settings.js
@@ -1,4 +1,5 @@
 import {
+    callPopup,
     getRequestHeaders,
     getStoppingStrings,
     getTextTokens,
@@ -9,7 +10,11 @@ import {
 } from "../script.js";
 import { getCfg } from "./extensions/cfg/util.js";
 import { MAX_CONTEXT_DEFAULT, tokenizers } from "./power-user.js";
-import { getStringHash } from "./utils.js";
+import { 
+    download,
+    parseJsonFile,
+    getStringHash 
+} from "./utils.js";
 
 export {
     nai_settings,
@@ -27,6 +32,10 @@ const default_presets = {
     "clio-v1": "Talker-Chat-Clio",
     "kayra-v1": "Carefree-Kayra"
 }
+const default_bias = 'Default (none)';
+const default_bias_presets = {
+    [default_bias]: []
+};
 
 const nai_settings = {
     temperature: 1.5,
@@ -49,6 +58,8 @@ const nai_settings = {
     cfg_uc: '',
     banned_tokens: '',
     order: default_order,
+    bias_preset_selected: default_bias,
+    bias_presets: default_bias_presets,
 };
 
 const nai_tiers = {
@@ -60,6 +71,7 @@ const nai_tiers = {
 
 let novel_data = null;
 let badWordsCache = {};
+let biasCache = undefined;
 
 export function setNovelData(data) {
     novel_data = data;
@@ -147,6 +159,8 @@ function loadNovelSettings(settings) {
     nai_settings.cfg_uc = settings.cfg_uc || '';
     nai_settings.banned_tokens = settings.banned_tokens || '';
     nai_settings.order = settings.order || default_order;
+    nai_settings.bias_preset_selected = settings.bias_preset_selected ?? default_bias;
+    nai_settings.bias_presets = settings.bias_presets ?? default_bias_presets;
     loadNovelSettingsUi(nai_settings);
 }
 
@@ -189,6 +203,16 @@ function loadNovelSettingsUi(ui_settings) {
 
     $("#streaming_novel").prop('checked', ui_settings.streaming_novel);
     sortItemsByOrder(ui_settings.order);
+
+    $('#novelai_logit_bias_preset').empty();
+    for (const preset of Object.keys(ui_settings.bias_presets)) {
+        const option = document.createElement('option');
+        option.innerText = preset;
+        option.value = preset;
+        option.selected = preset === ui_settings.bias_preset_selected;
+        $('#novelai_logit_bias_preset').append(option);
+    }
+    $('#novelai_logit_bias_preset').trigger('change');
 }
 
 const sliders = [
@@ -381,7 +405,7 @@ function getBadWordPermutations(text) {
     return result;
 }
 
-export function getNovelGenerationData(finalPrompt, this_settings, this_amount_gen, isImpersonate) {
+export async function getNovelGenerationData(finalPrompt, this_settings, this_amount_gen, isImpersonate) {
     const clio = nai_settings.model_novel.includes('clio');
     const kayra = nai_settings.model_novel.includes('kayra');
 
@@ -398,6 +422,15 @@ export function getNovelGenerationData(finalPrompt, this_settings, this_amount_g
     const prefix = selectPrefix(nai_settings.prefix, finalPrompt);
     const cfgSettings = getCfg();
 
+    let logitBias = [];
+    if (tokenizerType !== tokenizers.NONE
+        && nai_settings.bias_preset_selected
+        && Array.isArray(nai_settings.bias_presets[nai_settings.bias_preset_selected])
+        && nai_settings.bias_presets[nai_settings.bias_preset_selected].length) {
+        logitBias = biasCache || await calculateLogitBias();
+        biasCache = logitBias;
+    }
+
     return {
         "input": finalPrompt,
         "model": nai_settings.model_novel,
@@ -422,6 +455,7 @@ export function getNovelGenerationData(finalPrompt, this_settings, this_amount_g
         "phrase_rep_pen": nai_settings.phrase_rep_pen,
         "stop_sequences": stopSequences,
         "bad_words_ids": badWordIds,
+        "logit_bias_exp": logitBias,
         "generate_until_sentence": true,
         "use_cache": false,
         "use_string": true,
@@ -487,6 +521,178 @@ function saveSamplingOrder() {
     saveSettingsDebounced();
 }
 
+function onLogitBiasPresetChange() {
+    const value = $('#novelai_logit_bias_preset').find(':selected').val();
+    const preset = nai_settings.bias_presets[value];
+
+    if (!Array.isArray(preset)) {
+        console.error('Preset not found');
+        return;
+    }
+
+    nai_settings.bias_preset_selected = value;
+    $('.novelai_logit_bias_list').empty();
+
+    for (const entry of preset) {
+        if (entry) {
+            createLogitBiasListItem(entry);
+        }
+    }
+
+    biasCache = undefined;
+    saveSettingsDebounced();
+}
+
+function createNewLogitBiasEntry() {
+    const entry = { text: '', value: 0 };
+    nai_settings.bias_presets[nai_settings.bias_preset_selected].push(entry);
+    biasCache = undefined;
+    createLogitBiasListItem(entry);
+    saveSettingsDebounced();
+}
+
+function createLogitBiasListItem(entry) {
+    const id = nai_settings.bias_presets[nai_settings.bias_preset_selected].indexOf(entry);
+    const template = $('#novelai_logit_bias_template .novelai_logit_bias_form').clone();
+    template.data('id', id);
+    template.find('.novelai_logit_bias_text').val(entry.text).on('input', function () {
+        nai_settings.bias_presets[nai_settings.bias_preset_selected][id].text = $(this).val();
+        biasCache = undefined;
+        saveSettingsDebounced();
+    });
+    template.find('.novelai_logit_bias_value').val(entry.value).on('input', function () {
+        nai_settings.bias_presets[nai_settings.bias_preset_selected][id].value = Number($(this).val());
+        biasCache = undefined;
+        saveSettingsDebounced();
+    });
+    template.find('.novelai_logit_bias_remove').on('click', function () {
+        $(this).closest('.novelai_logit_bias_form').remove();
+        nai_settings.bias_presets[nai_settings.bias_preset_selected].splice(id, 1);
+        biasCache = undefined;
+        saveSettingsDebounced();
+    });
+    $('.novelai_logit_bias_list').prepend(template);
+}
+
+async function createNewLogitBiasPreset() {
+    const name = await callPopup('Preset name:', 'input');
+
+    if (!name) {
+        return;
+    }
+
+    if (name in nai_settings.bias_presets) {
+        toastr.error('Preset name should be unique.');
+        return;
+    }
+
+    nai_settings.bias_preset_selected = name;
+    nai_settings.bias_presets[name] = [];
+
+    addLogitBiasPresetOption(name);
+    saveSettingsDebounced();
+}
+
+function addLogitBiasPresetOption(name) {
+    const option = document.createElement('option');
+    option.innerText = name;
+    option.value = name;
+    option.selected = true;
+
+    $('#novelai_logit_bias_preset').append(option);
+    $('#novelai_logit_bias_preset').trigger('change');
+}
+
+async function onLogitBiasPresetDeleteClick() {
+    const value = await callPopup('Delete the preset?', 'confirm');
+
+    if (!value) {
+        return;
+    }
+
+    $(`#novelai_logit_bias_preset option[value="${nai_settings.bias_preset_selected}"]`).remove();
+    delete nai_settings.bias_presets[nai_settings.bias_preset_selected];
+    nai_settings.bias_preset_selected = null;
+
+    if (Object.keys(nai_settings.bias_presets).length) {
+        nai_settings.bias_preset_selected = Object.keys(nai_settings.bias_presets)[0];
+        $(`#novelai_logit_bias_preset option[value="${nai_settings.bias_preset_selected}"]`).attr('selected', true);
+        $('#novelai_logit_bias_preset').trigger('change');
+    }
+
+    biasCache = undefined;
+    saveSettingsDebounced();
+}
+
+async function onLogitBiasPresetImportFileChange(e) {
+    const file = e.target.files[0];
+
+    if (!file || file.type !== "application/json") {
+        return;
+    }
+
+    const name = file.name.replace(/\.[^/.]+$/, "");
+    const importedFile = await parseJsonFile(file);
+    e.target.value = '';
+
+    if (name in nai_settings.bias_presets) {
+        toastr.error('Preset name should be unique.');
+        return;
+    }
+
+    if (!Array.isArray(importedFile)) {
+        toastr.error('Invalid logit bias preset file.');
+        return;
+    }
+
+    for (const entry of importedFile) {
+        if (typeof entry == 'object') {
+            if (entry.hasOwnProperty('text') && entry.hasOwnProperty('value')) {
+                continue;
+            }
+        }
+
+        callPopup('Invalid logit bias preset file.', 'text');
+        return;
+    }
+
+    nai_settings.bias_presets[name] = importedFile;
+    nai_settings.bias_preset_selected = name;
+
+    addLogitBiasPresetOption(name);
+    saveSettingsDebounced();
+}
+
+function onLogitBiasPresetImportClick() {
+    $('#novelai_logit_bias_import_file').trigger('click');
+}
+
+function onLogitBiasPresetExportClick() {
+    if (!nai_settings.bias_preset_selected || Object.keys(nai_settings.bias_presets).length === 0) {
+        return;
+    }
+
+    const presetJsonString = JSON.stringify(nai_settings.bias_presets[nai_settings.bias_preset_selected], null, 4);
+    download(presetJsonString, nai_settings.bias_preset_selected, 'application/json');
+}
+
+async function calculateLogitBias() {
+    const bias_preset = nai_settings.bias_presets[nai_settings.bias_preset_selected];
+
+    const clio = nai_settings.model_novel.includes('clio');
+    const kayra = nai_settings.model_novel.includes('kayra');
+    const tokenizerType = kayra ? tokenizers.NERD2 : (clio ? tokenizers.NERD : tokenizers.NONE);
+
+    return bias_preset.map(bias => {
+        return {
+            bias: bias.value,
+            ensure_sequence_finish: false,
+            generate_once: false,
+            sequence: getTextTokens(tokenizerType, bias.text)
+        }
+    });
+}
+
 export async function generateNovelWithStreaming(generate_data, signal) {
     const response = await fetch('/generate_novelai', {
         headers: getRequestHeaders(),
@@ -590,4 +796,12 @@ $(document).ready(function () {
         console.log('Sampler toggled:', $item.data('id'), !isEnabled);
         saveSamplingOrder();
     });
+
+    $("#novelai_logit_bias_preset").on("change", onLogitBiasPresetChange);
+    $("#novelai_logit_bias_new_preset").on("click", createNewLogitBiasPreset);
+    $("#novelai_logit_bias_new_entry").on("click", createNewLogitBiasEntry);
+    $("#novelai_logit_bias_import_file").on("input", onLogitBiasPresetImportFileChange);
+    $("#novelai_logit_bias_import_preset").on("click", onLogitBiasPresetImportClick);
+    $("#novelai_logit_bias_export_preset").on("click", onLogitBiasPresetExportClick);
+    $("#novelai_logit_bias_delete_preset").on("click", onLogitBiasPresetDeleteClick);
 });
diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 71b83b9f2..79072e56c 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -2179,7 +2179,7 @@ function createLogitBiasListItem(entry) {
     });
     template.find('.openai_logit_bias_remove').on('click', function () {
         $(this).closest('.openai_logit_bias_form').remove();
-        oai_settings.bias_presets[oai_settings.bias_preset_selected][id] = undefined;
+        oai_settings.bias_presets[oai_settings.bias_preset_selected].splice(id, 1);
         biasCache = undefined;
         saveSettingsDebounced();
     });
diff --git a/public/style.css b/public/style.css
index 9764133c9..809d8601b 100644
--- a/public/style.css
+++ b/public/style.css
@@ -4590,6 +4590,41 @@ toolcool-color-picker {
     height: 100%;
 }
 
+.novelai_logit_bias_form {
+    display: flex;
+    flex-direction: row;
+    column-gap: 10px;
+    align-items: center;
+}
+
+.novelai_logit_bias_text,
+.novelai_logit_bias_value {
+    flex: 1;
+}
+
+.novelai_logit_bias_list {
+    display: flex;
+    flex-direction: column;
+    gap: 10px;
+}
+
+.novelai_logit_bias_list:empty {
+    width: 100%;
+    height: 100%;
+}
+
+.novelai_logit_bias_list:empty::before {
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    content: "No items";
+    font-weight: bolder;
+    width: 100%;
+    height: 100%;
+    opacity: 0.8;
+    min-height: 3rem;
+}
+
 /*used to fix smallness of certain FontAwesome glyph which break button squareness*/
 /*currently used on: CharList Import*/
 
diff --git a/server.js b/server.js
index 0a0695d79..b7334c781 100644
--- a/server.js
+++ b/server.js
@@ -1864,6 +1864,11 @@ app.post("/generate_novelai", jsonParser, async function (request, response_gene
         }
     }
 
+    // Add default biases for dinkus and asterism
+    const logit_bias_exp = isNewModel
+        ? request.body.logit_bias_exp.concat(novelai.logitBiasExp)
+        : null;
+
     const data = {
         "input": request.body.input,
         "model": request.body.model,
@@ -1889,9 +1894,8 @@ app.post("/generate_novelai", jsonParser, async function (request, response_gene
             "cfg_uc": request.body.cfg_uc,
             "phrase_rep_pen": request.body.phrase_rep_pen,
             "stop_sequences": request.body.stop_sequences,
-            //"stop_sequences": {{187}},
             "bad_words_ids": badWordsList,
-            "logit_bias_exp": isNewModel ? novelai.logitBiasExp : null,
+            "logit_bias_exp": logit_bias_exp,
             //generate_until_sentence = true;
             "use_cache": request.body.use_cache,
             "use_string": true,