rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-02-09 08:38:53 +01: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
Show Stats 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);
// Tokenizers
require('./src/endpoints/tokenizers').registerEndpoints(app, jsonParser);
app.use('/api/tokenizers', require('./src/endpoints/tokenizers').router);
// Preset management
require('./src/endpoints/presets').registerEndpoints(app, jsonParser);

54
src/endpoints/tokenizers.js
View File

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