#1569 Add logit bias for text completions

public/index.html

@@ -1383,6 +1383,21 @@
                                             <textarea id="banned_tokens_textgenerationwebui" class="text_pole textarea_compact" name="banned_tokens_textgenerationwebui" rows="3" placeholder="Example:&#10;some text&#10;[42, 69, 1337]"></textarea>
                                         </div>
                                     </div>
+                                    <div class="range-block wide100p">
+                                        <div class="range-block-title title_restorable">
+                                            <span data-i18n="Logit Bias">Logit Bias</span>
+                                            <div id="textgen_logit_bias_new_entry" class="menu_button menu_button_icon">
+                                                <i class="fa-xs fa-solid fa-plus"></i>
+                                                <small data-i18n="Add">Add</small>
+                                            </div>
+                                        </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="flex-container flexFlowColumn wide100p">
+                                            <div class="logit_bias_list"></div>
+                                        </div>
+                                    </div>
                                     <div data-newbie-hidden data-forAphro=False class="wide100p">
                                         <hr  class="width100p">
                                         <h4 data-i18n="CFG" class="textAlignCenter">CFG

public/scripts/textgen-settings.js

@@ -9,6 +9,7 @@ import {
     setOnlineStatus,
     substituteParams,
 } from '../script.js';
+import { BIAS_CACHE, createNewLogitBiasEntry, displayLogitBias, getLogitBiasListResult } from './logit-bias.js';
 
 import {
     power_user,
@@ -35,6 +36,7 @@ export const textgen_types = {
 };
 
 const { MANCER, APHRODITE, TOGETHERAI } = textgen_types;
+const BIAS_KEY = '#textgenerationwebui_api-settings';
 
 // Maybe let it be configurable in the future?
 // (7 days later) The future has come.
@@ -94,6 +96,7 @@ const settings = {
     togetherai_model: 'Gryphe/MythoMax-L2-13b',
     legacy_api: false,
     sampler_order: KOBOLDCPP_ORDER,
+    logit_bias: [],
     n: 1,
 };
 
@@ -147,6 +150,7 @@ const setting_names = [
     //'prompt_log_probs_aphrodite'
     'sampler_order',
     'n',
+    'logit_bias',
 ];
 
 async function selectPreset(name) {
@@ -162,6 +166,7 @@ async function selectPreset(name) {
         setSettingByName(name, value, true);
     }
     setGenerationParamsFromPreset(preset);
+    displayLogitBias(preset.logit_bias, BIAS_KEY);
     saveSettingsDebounced();
 }
 
@@ -243,6 +248,42 @@ function getCustomTokenBans() {
     return result.filter(onlyUnique).map(x => String(x)).join(',');
 }
 
+/**
+ * Calculates logit bias object from the logit bias list.
+ * @returns {object} Logit bias object
+ */
+function calculateLogitBias() {
+    if (!Array.isArray(settings.logit_bias) || settings.logit_bias.length === 0) {
+        return {};
+    }
+
+    const tokenizer = SENTENCEPIECE_TOKENIZERS.includes(power_user.tokenizer) ? power_user.tokenizer : tokenizers.LLAMA;
+    const result = {};
+
+    /**
+     * Adds bias to the logit bias object.
+     * @param {number} bias
+     * @param {number[]} sequence
+     * @returns {object} Accumulated logit bias object
+     */
+    function addBias(bias, sequence) {
+        if (sequence.length === 0) {
+            return;
+        }
+
+        for (const logit of sequence) {
+            const key = String(logit);
+            result[key] = bias;
+        }
+
+        return result;
+    }
+
+    getLogitBiasListResult(settings.logit_bias, tokenizer, addBias);
+
+    return result;
+}
+
 function loadTextGenSettings(data, loadedSettings) {
     textgenerationwebui_presets = convertPresets(data.textgenerationwebui_presets);
     textgenerationwebui_preset_names = data.textgenerationwebui_preset_names ?? [];
@@ -270,6 +311,7 @@ function loadTextGenSettings(data, loadedSettings) {
 
     $('#textgen_type').val(settings.type);
     showTypeSpecificControls(settings.type);
+    displayLogitBias(settings.logit_bias, BIAS_KEY);
     //this is needed because showTypeSpecificControls() does not handle NOT declarations
     if (settings.type === textgen_types.APHRODITE) {
         $('[data-forAphro=False]').each(function () {
@@ -415,6 +457,8 @@ jQuery(function () {
             saveSettingsDebounced();
         });
     }
+
+    $('#textgen_logit_bias_new_entry').on('click', () => createNewLogitBiasEntry(settings.logit_bias, BIAS_KEY));
 });
 
 function showTypeSpecificControls(type) {
@@ -440,6 +484,11 @@ function setSettingByName(setting, value, trigger) {
         return;
     }
 
+    if ('logit_bias' === setting) {
+        settings.logit_bias = Array.isArray(value) ? value : [];
+        return;
+    }
+
     const isCheckbox = $(`#${setting}_textgenerationwebui`).attr('type') == 'checkbox';
     const isText = $(`#${setting}_textgenerationwebui`).attr('type') == 'text' || $(`#${setting}_textgenerationwebui`).is('textarea');
     if (isCheckbox) {
@@ -642,6 +691,12 @@ export function getTextGenGenerationData(finalPrompt, maxTokens, isImpersonate,
         APIflags = Object.assign(APIflags, aphroditeExclusionFlags);
     }
 
+    if (Array.isArray(settings.logit_bias) && settings.logit_bias.length) {
+        const logitBias = BIAS_CACHE.get(BIAS_KEY) || calculateLogitBias();
+        BIAS_CACHE.set(BIAS_KEY, logitBias);
+        APIflags.logit_bias = logitBias;
+    }
+
     return APIflags;
 }
 

public/style.css

@@ -3480,7 +3480,7 @@ a {
     width: 100%;
     height: 100%;
     opacity: 0.8;
-    min-height: 2.5rem;
+    min-height: 2.5em;
 }
 
 .openai_logit_bias_preset_form {