diff --git a/public/script.js b/public/script.js index 2cf55464d..ca9acf810 100644 --- a/public/script.js +++ b/public/script.js @@ -581,6 +581,8 @@ function getTokenCount(str, padding = undefined) { return countTokensRemote('/tokenize_nerdstash', str, padding); case tokenizers.NERD2: return countTokensRemote('/tokenize_nerdstash_v2', str, padding); + case tokenizers.API: + return countTokensRemote('/tokenize_via_api', str, padding); default: console.warn("Unknown tokenizer type", tokenizerType); return Math.ceil(str.length / CHARACTERS_PER_TOKEN_RATIO) + padding; diff --git a/public/scripts/power-user.js b/public/scripts/power-user.js index f7dad23e5..e121a7a16 100644 --- a/public/scripts/power-user.js +++ b/public/scripts/power-user.js @@ -75,6 +75,7 @@ const tokenizers = { LLAMA: 3, NERD: 4, NERD2: 5, + API: 6, } const send_on_enter_options = { diff --git a/server.js b/server.js index 768b67d38..7be36b851 100644 --- a/server.js +++ b/server.js @@ -3401,6 +3401,27 @@ function createTokenizationHandler(getTokenizerFn) { app.post("/tokenize_llama", jsonParser, createTokenizationHandler(() => spp_llama)); app.post("/tokenize_nerdstash", jsonParser, createTokenizationHandler(() => spp_nerd)); app.post("/tokenize_nerdstash_v2", jsonParser, createTokenizationHandler(() => spp_nerd_v2)); +app.post("/tokenize_via_api", jsonParser, async function(request, response) { + if (!request.body) { + return response.sendStatus(400); + } + const text = request.body.text || ''; + + try { + const args = { + body: JSON.stringify({"prompt": text}), + headers: { "Content-Type": "application/json" } + }; + + const data = await postAsync(api_server + "/v1/token-count", args); + console.log(data); + return response.send({ count: data['results'][0]['tokens'] }); + } catch (error) { + console.log(error); + return response.send({ error: true }); + } +}); + // ** REST CLIENT ASYNC WRAPPERS **