rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Use Express router for tokenizers endpoint

Browse Source
This commit is contained in:
valadaptive
2023-12-04 13:00:13 -05:00
parent 414c9bd5fb
commit 4e073250a2
2 changed files with 161 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
server.js
View File

@ -3579,7 +3579,7 @@ async function fetchJSON(url, args = {}) {
require('./src/endpoints/openai').registerEndpoints(app, jsonParser, urlencodedParser); require('./src/endpoints/openai').registerEndpoints(app, jsonParser, urlencodedParser);
// Tokenizers // Tokenizers
require('./src/endpoints/tokenizers').registerEndpoints(app, jsonParser); app.use('/api/tokenizers', require('./src/endpoints/tokenizers').router);
// Preset management // Preset management
require('./src/endpoints/presets').registerEndpoints(app, jsonParser); require('./src/endpoints/presets').registerEndpoints(app, jsonParser);

54
src/endpoints/tokenizers.js
View File

@ -1,10 +1,12 @@
const fs = require('fs'); const fs = require('fs');
const path = require('path'); const path = require('path');
const express = require('express');
const { SentencePieceProcessor } = require('@agnai/sentencepiece-js'); const { SentencePieceProcessor } = require('@agnai/sentencepiece-js');
const tiktoken = require('@dqbd/tiktoken'); const tiktoken = require('@dqbd/tiktoken');
const { Tokenizer } = require('@agnai/web-tokenizers'); const { Tokenizer } = require('@agnai/web-tokenizers');
const { convertClaudePrompt } = require('../chat-completion'); const { convertClaudePrompt } = require('../chat-completion');
const { readSecret, SECRET_KEYS } = require('./secrets'); const { readSecret, SECRET_KEYS } = require('./secrets');
const { jsonParser } = require('../express-common');
/** /**
* @type {{[key: string]: import("@dqbd/tiktoken").Tiktoken}} Tokenizers cache * @type {{[key: string]: import("@dqbd/tiktoken").Tiktoken}} Tokenizers cache
@ -359,13 +361,9 @@ async function loadTokenizers() {
claude_tokenizer = await loadClaudeTokenizer('src/claude.json'); claude_tokenizer = await loadClaudeTokenizer('src/claude.json');
} }
/** const router = express.Router();
* Registers the tokenization endpoints.
* @param {import('express').Express} app Express app router.post('/ai21/count', jsonParser, async function (req, res) {
* @param {any} jsonParser JSON parser middleware
*/
function registerEndpoints(app, jsonParser) {
app.post('/api/tokenizers/ai21/count', jsonParser, async function (req, res) {
if (!req.body) return res.sendStatus(400); if (!req.body) return res.sendStatus(400);
const options = { const options = {
method: 'POST', method: 'POST',
@ -385,22 +383,22 @@ function registerEndpoints(app, jsonParser) {
console.error(err); console.error(err);
return res.send({ 'token_count': 0 }); return res.send({ 'token_count': 0 });
} }
}); });
app.post('/api/tokenizers/llama/encode', jsonParser, createSentencepieceEncodingHandler(spp_llama)); router.post('/llama/encode', jsonParser, createSentencepieceEncodingHandler(spp_llama));
app.post('/api/tokenizers/nerdstash/encode', jsonParser, createSentencepieceEncodingHandler(spp_nerd)); router.post('/nerdstash/encode', jsonParser, createSentencepieceEncodingHandler(spp_nerd));
app.post('/api/tokenizers/nerdstash_v2/encode', jsonParser, createSentencepieceEncodingHandler(spp_nerd_v2)); router.post('/nerdstash_v2/encode', jsonParser, createSentencepieceEncodingHandler(spp_nerd_v2));
app.post('/api/tokenizers/mistral/encode', jsonParser, createSentencepieceEncodingHandler(spp_mistral)); router.post('/mistral/encode', jsonParser, createSentencepieceEncodingHandler(spp_mistral));
app.post('/api/tokenizers/yi/encode', jsonParser, createSentencepieceEncodingHandler(spp_yi)); router.post('/yi/encode', jsonParser, createSentencepieceEncodingHandler(spp_yi));
app.post('/api/tokenizers/gpt2/encode', jsonParser, createTiktokenEncodingHandler('gpt2')); router.post('/gpt2/encode', jsonParser, createTiktokenEncodingHandler('gpt2'));
app.post('/api/tokenizers/llama/decode', jsonParser, createSentencepieceDecodingHandler(spp_llama)); router.post('/llama/decode', jsonParser, createSentencepieceDecodingHandler(spp_llama));
app.post('/api/tokenizers/nerdstash/decode', jsonParser, createSentencepieceDecodingHandler(spp_nerd)); router.post('/nerdstash/decode', jsonParser, createSentencepieceDecodingHandler(spp_nerd));
app.post('/api/tokenizers/nerdstash_v2/decode', jsonParser, createSentencepieceDecodingHandler(spp_nerd_v2)); router.post('/nerdstash_v2/decode', jsonParser, createSentencepieceDecodingHandler(spp_nerd_v2));
app.post('/api/tokenizers/mistral/decode', jsonParser, createSentencepieceDecodingHandler(spp_mistral)); router.post('/mistral/decode', jsonParser, createSentencepieceDecodingHandler(spp_mistral));
app.post('/api/tokenizers/yi/decode', jsonParser, createSentencepieceDecodingHandler(spp_yi)); router.post('/yi/decode', jsonParser, createSentencepieceDecodingHandler(spp_yi));
app.post('/api/tokenizers/gpt2/decode', jsonParser, createTiktokenDecodingHandler('gpt2')); router.post('/gpt2/decode', jsonParser, createTiktokenDecodingHandler('gpt2'));
app.post('/api/tokenizers/openai/encode', jsonParser, async function (req, res) { router.post('/openai/encode', jsonParser, async function (req, res) {
try { try {
const queryModel = String(req.query.model || ''); const queryModel = String(req.query.model || '');
@ -433,9 +431,9 @@ function registerEndpoints(app, jsonParser) {
console.log(error); console.log(error);
return res.send({ ids: [], count: 0, chunks: [] }); return res.send({ ids: [], count: 0, chunks: [] });
} }
}); });
app.post('/api/tokenizers/openai/decode', jsonParser, async function (req, res) { router.post('/openai/decode', jsonParser, async function (req, res) {
try { try {
const queryModel = String(req.query.model || ''); const queryModel = String(req.query.model || '');
@ -467,9 +465,9 @@ function registerEndpoints(app, jsonParser) {
console.log(error); console.log(error);
return res.send({ text: '' }); return res.send({ text: '' });
} }
}); });
app.post('/api/tokenizers/openai/count', jsonParser, async function (req, res) { router.post('/openai/count', jsonParser, async function (req, res) {
try { try {
if (!req.body) return res.sendStatus(400); if (!req.body) return res.sendStatus(400);
@ -534,8 +532,7 @@ function registerEndpoints(app, jsonParser) {
const num_tokens = Math.ceil(jsonBody.length / CHARS_PER_TOKEN); const num_tokens = Math.ceil(jsonBody.length / CHARS_PER_TOKEN);
res.send({ 'token_count': num_tokens }); res.send({ 'token_count': num_tokens });
} }
}); });
}
module.exports = { module.exports = {
TEXT_COMPLETION_MODELS, TEXT_COMPLETION_MODELS,
@ -543,8 +540,7 @@ module.exports = {
getTiktokenTokenizer, getTiktokenTokenizer,
countClaudeTokens, countClaudeTokens,
loadTokenizers, loadTokenizers,
registerEndpoints,
getSentencepiceTokenizer, getSentencepiceTokenizer,
sentencepieceTokenizers, sentencepieceTokenizers,
router,
}; };