From ff4c67b99384776ce5bb290c861238d8ec382acd Mon Sep 17 00:00:00 2001
From: justpain02 <justpain02kor@gmail.com>
Date: Fri, 5 Jan 2024 23:57:59 +0900
Subject: [PATCH] Add selectable tokenizer supports on Ooba

---
 src/lib/Setting/Pages/OobaSettings.svelte |  2 ++
 src/ts/model/ooba.ts                      |  1 +
 src/ts/tokenizer.ts                       | 14 ++++++++++++--
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/lib/Setting/Pages/OobaSettings.svelte b/src/lib/Setting/Pages/OobaSettings.svelte
index 9fa6dfa9..5093291b 100644
--- a/src/lib/Setting/Pages/OobaSettings.svelte
+++ b/src/lib/Setting/Pages/OobaSettings.svelte
@@ -61,6 +61,8 @@
             <OptionalInput marginBottom={true} bind:value={$DataBase.reverseProxyOobaArgs.chat_instruct_command} />
         {/if}
     {/if}
+    <span class="text-textcolor">tokenizer</span>
+    <OptionalInput marginBottom={true} bind:value={$DataBase.reverseProxyOobaArgs.tokenizer} />
     <span class="text-textcolor">min_p</span>
     <OptionalInput marginBottom={true} bind:value={$DataBase.reverseProxyOobaArgs.min_p} numberMode />
     <span class="text-textcolor">top_k</span>
diff --git a/src/ts/model/ooba.ts b/src/ts/model/ooba.ts
index 9fb731cf..af28c241 100644
--- a/src/ts/model/ooba.ts
+++ b/src/ts/model/ooba.ts
@@ -11,6 +11,7 @@ export interface OobaChatCompletionRequestParams {
     greeting?: string
     chat_instruct_command?: string
     preset?: string; // The '?' denotes that the property is optional
+    tokenizer?: string;
     min_p?: number;
     top_k?: number;
     repetition_penalty?: number;
diff --git a/src/ts/tokenizer.ts b/src/ts/tokenizer.ts
index 8f26cef2..0fbedb09 100644
--- a/src/ts/tokenizer.ts
+++ b/src/ts/tokenizer.ts
@@ -24,10 +24,20 @@ async function encode(data:string):Promise<(number[]|Uint32Array|Int32Array)>{
     if(db.aiModel.startsWith('local_') ||
         db.aiModel === 'mancer' ||
         db.aiModel === 'textgen_webui' ||
-        (db.aiModel === 'reverse_proxy' && db.reverseProxyOobaMode ||
-        db.aiModel === 'ooba')){
+        (db.aiModel === 'reverse_proxy' && db.reverseProxyOobaMode)){
         return await tokenizeWebTokenizers(data, 'llama')
     }
+    if(db.aiModel === 'ooba'){
+        if(db.reverseProxyOobaArgs.tokenizer === 'mixtral' || db.reverseProxyOobaArgs.tokenizer === 'mistral'){
+            return await tokenizeWebTokenizers(data, 'mistral')
+        }
+        else if(db.reverseProxyOobaArgs.tokenizer === 'llama'){
+            return await tokenizeWebTokenizers(data, 'llama')
+        }
+        else{
+            return await tokenizeWebTokenizers(data, 'llama')
+        }
+    }
 
     return await tikJS(data)
 }