WebLLM: use current tokenizer if not available

This commit is contained in:
Cohee
2025-01-07 20:01:59 +02:00
parent 5af7852465
commit 81841ca2a6

View File

@@ -3,6 +3,7 @@ import { extension_settings, openThirdPartyExtensionMenu } from '../extensions.j
import { oai_settings } from '../openai.js'; import { oai_settings } from '../openai.js';
import { SECRET_KEYS, secret_state } from '../secrets.js'; import { SECRET_KEYS, secret_state } from '../secrets.js';
import { textgen_types, textgenerationwebui_settings } from '../textgen-settings.js'; import { textgen_types, textgenerationwebui_settings } from '../textgen-settings.js';
import { getTokenCountAsync } from '../tokenizers.js';
import { createThumbnail, isValidUrl } from '../utils.js'; import { createThumbnail, isValidUrl } from '../utils.js';
/** /**
@@ -235,6 +236,7 @@ export async function generateWebLlmChatPrompt(messages, params = {}) {
/** /**
* Counts the number of tokens in the provided text using WebLLM's default model. * Counts the number of tokens in the provided text using WebLLM's default model.
* Fallbacks to the current model's tokenizer if WebLLM token count fails.
* @param {string} text Text to count tokens in * @param {string} text Text to count tokens in
* @returns {Promise<number>} Number of tokens in the text * @returns {Promise<number>} Number of tokens in the text
*/ */
@@ -243,9 +245,14 @@ export async function countWebLlmTokens(text) {
throw new Error('WebLLM extension is not installed.'); throw new Error('WebLLM extension is not installed.');
} }
const engine = SillyTavern.llm; try {
const response = await engine.countTokens(text); const engine = SillyTavern.llm;
return response; const response = await engine.countTokens(text);
return response;
} catch (error) {
// Fallback to using current model's tokenizer
return getTokenCountAsync(text);
}
} }
/** /**