From 62faddac8d1272507b678c76ebc7c150487bf062 Mon Sep 17 00:00:00 2001
From: kingbri <bdashore3@proton.me>
Date: Sat, 11 May 2024 00:58:29 -0400
Subject: [PATCH] Textgen: Add banned_strings

TabbyAPI supports the ability to ban the presence of strings during
a generation. Add this support in SillyTavern by handling lines
enclosed in quotes as a special case.

Signed-off-by: kingbri <bdashore3@proton.me>
---
 public/scripts/textgen-settings.js | 31 ++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/public/scripts/textgen-settings.js b/public/scripts/textgen-settings.js
index 588a52fb5..ed1065935 100644
--- a/public/scripts/textgen-settings.js
+++ b/public/scripts/textgen-settings.js
@@ -328,15 +328,20 @@ function getTokenizerForTokenIds() {
 }
 
 /**
- * @returns {string} String with comma-separated banned token IDs
+ * @typedef {{banned_tokens: string, banned_strings: string[]}} TokenBanResult
+ * @returns {TokenBanResult} String with comma-separated banned token IDs
  */
 function getCustomTokenBans() {
     if (!settings.banned_tokens && !textgenerationwebui_banned_in_macros.length) {
-        return '';
+        return {
+            banned_tokens: '',
+            banned_strings: [],
+        };
     }
 
     const tokenizer = getTokenizerForTokenIds();
-    const result = [];
+    const banned_tokens = [];
+    const banned_strings = [];
     const sequences = settings.banned_tokens
         .split('\n')
         .concat(textgenerationwebui_banned_in_macros)
@@ -358,24 +363,31 @@ function getCustomTokenBans() {
                 const tokens = JSON.parse(line);
 
                 if (Array.isArray(tokens) && tokens.every(t => Number.isInteger(t))) {
-                    result.push(...tokens);
+                    banned_tokens.push(...tokens);
                 } else {
                     throw new Error('Not an array of integers');
                 }
             } catch (err) {
                 console.log(`Failed to parse bad word token list: ${line}`, err);
             }
+        } else if (line.startsWith('"') && line.endsWith('"')) {
+            // Remove the enclosing quotes
+
+            banned_strings.push(line.slice(1, -1))
         } else {
             try {
                 const tokens = getTextTokens(tokenizer, line);
-                result.push(...tokens);
+                banned_tokens.push(...tokens);
             } catch {
                 console.log(`Could not tokenize raw text: ${line}`);
             }
         }
     }
 
-    return result.filter(onlyUnique).map(x => String(x)).join(',');
+    return {
+        banned_tokens: banned_tokens.filter(onlyUnique).map(x => String(x)).join(','),
+        banned_strings: banned_strings,
+    };
 }
 
 /**
@@ -987,6 +999,8 @@ export function isJsonSchemaSupported() {
 
 export function getTextGenGenerationData(finalPrompt, maxTokens, isImpersonate, isContinue, cfgValues, type) {
     const canMultiSwipe = !isContinue && !isImpersonate && type !== 'quiet';
+    const {banned_tokens, banned_strings} = getCustomTokenBans();
+
     let params = {
         'prompt': finalPrompt,
         'model': getTextGenModel(),
@@ -1033,8 +1047,9 @@ export function getTextGenGenerationData(finalPrompt, maxTokens, isImpersonate,
         'mirostat_tau': settings.mirostat_tau,
         'mirostat_eta': settings.mirostat_eta,
         'custom_token_bans': [APHRODITE, MANCER].includes(settings.type) ?
-            toIntArray(getCustomTokenBans()) :
-            getCustomTokenBans(),
+            toIntArray(banned_tokens) :
+            banned_tokens,
+        'banned_strings': banned_strings,
         'api_type': settings.type,
         'api_server': getTextGenServer(),
         'legacy_api': settings.legacy_api && (settings.type === OOBA || settings.type === APHRODITE),