From 7f718e09be933bb6207799fcd5d75de963b2397a Mon Sep 17 00:00:00 2001
From: Cohee1207 <titan0599@gmail.com>
Date: Sat, 6 May 2023 15:30:15 +0300
Subject: [PATCH] #228 Don't use a selected tokenizer for parallel prompt
 building of OAI prompts

---
 public/script.js         | 16 +++++++++++++++-
 public/scripts/openai.js |  5 +++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/public/script.js b/public/script.js
index da24c169b..e9537087e 100644
--- a/public/script.js
+++ b/public/script.js
@@ -73,6 +73,7 @@ import {
     oai_settings,
     is_get_status_openai,
     openai_messages_count,
+    getTokenCountOpenAI,
 } from "./scripts/openai.js";
 
 import {
@@ -368,7 +369,20 @@ $(document).ajaxError(function myErrorHandler(_, xhr) {
 });
 
 function getTokenCount(str, padding = 0) {
-    switch (power_user.tokenizer) {
+    let tokenizerType = power_user.tokenizer; 
+
+    if (main_api === 'openai') {
+        // For main prompt building
+        if (padding == power_user.token_padding) {
+            tokenizerType = tokenizers.NONE;
+        // For extensions and WI
+        } else {
+            return getTokenCountOpenAI(str);
+        }
+        
+    }
+
+    switch (tokenizerType) {
         case tokenizers.NONE:
             return Math.ceil(str.length / CHARACTERS_PER_TOKEN_RATIO) + padding;
         case tokenizers.GPT3:
diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 102c10984..8dac60050 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -127,6 +127,11 @@ const oai_settings = {
 let openai_setting_names;
 let openai_settings;
 
+export function getTokenCountOpenAI(text) {
+    const message = { role: 'system', content: text };
+    return countTokens(message, true);
+}
+
 function validateReverseProxy() {
     if (!oai_settings.reverse_proxy) {
         return;