From 3e82a7d4399d6824afa4a60fa4a0965fd8ee980b Mon Sep 17 00:00:00 2001
From: based <phrog@riseup.net>
Date: Thu, 14 Dec 2023 16:31:08 +1000
Subject: [PATCH] tokenizer changes and fixes. + a toggle

---
 public/index.html            |  8 ++++++++
 public/scripts/openai.js     | 11 +++++++++++
 public/scripts/tokenizers.js |  2 +-
 src/chat-completion.js       |  1 -
 src/endpoints/tokenizers.js  |  9 ++++-----
 5 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/public/index.html b/public/index.html
index 474d605de..67bb1b96d 100644
--- a/public/index.html
+++ b/public/index.html
@@ -1516,6 +1516,14 @@
                                             <span data-i18n="Use the appropriate tokenizer for Jurassic models, which is more efficient than GPT's.">Use the appropriate tokenizer for Jurassic models, which is more efficient than GPT's.</span>
                                         </div>
                                     </div>
+                                    <div class="range-block" data-source="makersuite">
+                                        <label for="use_google_tokenizer" title="Use Google Tokenizer" class="checkbox_label widthFreeExpand">
+                                            <input id="use_google_tokenizer" type="checkbox" /><span data-i18n="Use Google Tokenizer">Use Google Tokenizer</span>
+                                        </label>
+                                        <div class="toggle-description justifyLeft">
+                                            <span data-i18n="Use the appropriate tokenizer for Google models via their API. Slower prompt processing, but offers much more accurate token counting.">Use the appropriate tokenizer for Google models via their API. Slower prompt processing, but offers much more accurate token counting.</span>
+                                        </div>
+                                    </div>
                                     <div data-newbie-hidden class="range-block" data-source="claude">
                                         <label for="exclude_assistant" title="Exclude Assistant suffix" class="checkbox_label widthFreeExpand">
                                             <input id="exclude_assistant" type="checkbox" /><span data-i18n="Exclude Assistant suffix">Exclude Assistant suffix</span>
diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 015dd66a6..a71363cf4 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -223,6 +223,7 @@ const default_settings = {
     proxy_password: '',
     assistant_prefill: '',
     use_ai21_tokenizer: false,
+    use_google_tokenizer: false,
     exclude_assistant: false,
     use_alt_scale: false,
     squash_system_messages: false,
@@ -277,6 +278,7 @@ const oai_settings = {
     proxy_password: '',
     assistant_prefill: '',
     use_ai21_tokenizer: false,
+    use_google_tokenizer: false,
     exclude_assistant: false,
     use_alt_scale: false,
     squash_system_messages: false,
@@ -2313,6 +2315,7 @@ function loadOpenAISettings(data, settings) {
     if (settings.names_in_completion !== undefined) oai_settings.names_in_completion = !!settings.names_in_completion;
     if (settings.openai_model !== undefined) oai_settings.openai_model = settings.openai_model;
     if (settings.use_ai21_tokenizer !== undefined) { oai_settings.use_ai21_tokenizer = !!settings.use_ai21_tokenizer; oai_settings.use_ai21_tokenizer ? ai21_max = 8191 : ai21_max = 9200; }
+    if (settings.use_google_tokenizer !== undefined) oai_settings.use_google_tokenizer = !!settings.use_google_tokenizer;
     if (settings.exclude_assistant !== undefined) oai_settings.exclude_assistant = !!settings.exclude_assistant;
     if (settings.use_alt_scale !== undefined) { oai_settings.use_alt_scale = !!settings.use_alt_scale; updateScaleForm(); }
     $('#stream_toggle').prop('checked', oai_settings.stream_openai);
@@ -2345,6 +2348,7 @@ function loadOpenAISettings(data, settings) {
     $('#openai_show_external_models').prop('checked', oai_settings.show_external_models);
     $('#openai_external_category').toggle(oai_settings.show_external_models);
     $('#use_ai21_tokenizer').prop('checked', oai_settings.use_ai21_tokenizer);
+    $('#use_google_tokenizer').prop('checked', oai_settings.use_google_tokenizer);
     $('#exclude_assistant').prop('checked', oai_settings.exclude_assistant);
     $('#scale-alt').prop('checked', oai_settings.use_alt_scale);
     $('#openrouter_use_fallback').prop('checked', oai_settings.openrouter_use_fallback);
@@ -2537,6 +2541,7 @@ async function saveOpenAIPreset(name, settings, triggerUi = true) {
         show_external_models: settings.show_external_models,
         assistant_prefill: settings.assistant_prefill,
         use_ai21_tokenizer: settings.use_ai21_tokenizer,
+        use_google_tokenizer: settings.use_google_tokenizer,
         exclude_assistant: settings.exclude_assistant,
         use_alt_scale: settings.use_alt_scale,
         squash_system_messages: settings.squash_system_messages,
@@ -2898,6 +2903,7 @@ function onSettingsPresetChange() {
         proxy_password: ['#openai_proxy_password', 'proxy_password', false],
         assistant_prefill: ['#claude_assistant_prefill', 'assistant_prefill', false],
         use_ai21_tokenizer: ['#use_ai21_tokenizer', 'use_ai21_tokenizer', true],
+        use_google_tokenizer: ['#use_google_tokenizer', 'use_google_tokenizer', true],
         exclude_assistant: ['#exclude_assistant', 'exclude_assistant', true],
         use_alt_scale: ['#use_alt_scale', 'use_alt_scale', true],
         squash_system_messages: ['#squash_system_messages', 'squash_system_messages', true],
@@ -3506,6 +3512,11 @@ $(document).ready(async function () {
         saveSettingsDebounced();
     });
 
+    $('#use_google_tokenizer').on('change', function () {
+        oai_settings.use_google_tokenizer = !!$('#use_google_tokenizer').prop('checked');
+        saveSettingsDebounced();
+    });
+
     $('#exclude_assistant').on('change', function () {
         oai_settings.exclude_assistant = !!$('#exclude_assistant').prop('checked');
         $('#claude_assistant_prefill_block').toggle(!oai_settings.exclude_assistant);
diff --git a/public/scripts/tokenizers.js b/public/scripts/tokenizers.js
index 7600a0909..83842d73d 100644
--- a/public/scripts/tokenizers.js
+++ b/public/scripts/tokenizers.js
@@ -393,7 +393,7 @@ export function getTokenizerModel() {
  */
 export function countTokensOpenAI(messages, full = false) {
     const shouldTokenizeAI21 = oai_settings.chat_completion_source === chat_completion_sources.AI21 && oai_settings.use_ai21_tokenizer;
-    const shouldTokenizeGoogle = oai_settings.chat_completion_source === chat_completion_sources.MAKERSUITE;
+    const shouldTokenizeGoogle = oai_settings.chat_completion_source === chat_completion_sources.MAKERSUITE && oai_settings.use_google_tokenizer;
     let tokenizerEndpoint = '';
     if(shouldTokenizeAI21) {
         tokenizerEndpoint = '/api/tokenizers/ai21/count';
diff --git a/src/chat-completion.js b/src/chat-completion.js
index d1f97f8a3..cbed76abf 100644
--- a/src/chat-completion.js
+++ b/src/chat-completion.js
@@ -76,7 +76,6 @@ function convertGooglePrompt(messages) {
     const contents = [];
     let lastRole = '';
     let currentText = '';
-
     messages.forEach((message, index) => {
         const role = message.role === 'assistant' ? 'model' : 'user';
         if (lastRole === role) {
diff --git a/src/endpoints/tokenizers.js b/src/endpoints/tokenizers.js
index 096b0f093..3f7636795 100644
--- a/src/endpoints/tokenizers.js
+++ b/src/endpoints/tokenizers.js
@@ -4,7 +4,7 @@ const express = require('express');
 const { SentencePieceProcessor } = require('@agnai/sentencepiece-js');
 const tiktoken = require('@dqbd/tiktoken');
 const { Tokenizer } = require('@agnai/web-tokenizers');
-const { convertClaudePrompt } = require('../chat-completion');
+const { convertClaudePrompt, convertGooglePrompt } = require('../chat-completion');
 const { readSecret, SECRET_KEYS } = require('./secrets');
 const { TEXTGEN_TYPES } = require('../constants');
 const { jsonParser } = require('../express-common');
@@ -395,13 +395,12 @@ router.post('/google/count', jsonParser, async function (req, res) {
             accept: 'application/json',
             'content-type': 'application/json',
         },
-        body: JSON.stringify({ prompt: { text: req.body[0].content } }),
+        body: JSON.stringify({ contents: convertGooglePrompt(req.body) }),
     };
     try {
-        const response = await fetch(`https://generativelanguage.googleapis.com/v1beta/models/${req.query.model}:countTextTokens?key=${readSecret(SECRET_KEYS.MAKERSUITE)}`, options);
+        const response = await fetch(`https://generativelanguage.googleapis.com/v1beta/models/${req.query.model}:countTokens?key=${readSecret(SECRET_KEYS.MAKERSUITE)}`, options);
         const data = await response.json();
-        console.log(data)
-        return res.send({ 'token_count': data?.tokenCount || 0 });
+        return res.send({ 'token_count': data?.totalTokens || 0 });
     } catch (err) {
         console.error(err);
         return res.send({ 'token_count': 0 });