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

Add a token counting endpoint for vLLM

Browse Source
This commit is contained in:
sasha0552
2024-06-20 01:13:32 +00:00
committed by GitHub
parent 584d0e6222
commit 6c87dfe10d
2 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/endpoints/tokenizers.js
View File

@ -784,6 +784,7 @@ router.post('/remote/textgenerationwebui/encode', jsonParser, async function (re
const text = String(request.body.text) || '';
const baseUrl = String(request.body.url);
const legacyApi = Boolean(request.body.legacy_api);
const vllmModel = String(request.body.vllm_model) || '';
try {
const args = {
@ -814,7 +815,9 @@ router.post('/remote/textgenerationwebui/encode', jsonParser, async function (re
args.body = JSON.stringify({ 'content': text });
break;
case TEXTGEN_TYPES.VLLM:
return response.send({ error: true });
url += '/tokenize';
args.body = JSON.stringify({ 'model': vllmModel, 'prompt': text });
break;
case TEXTGEN_TYPES.APHRODITE:
url += '/v1/tokenize';
args.body = JSON.stringify({ 'prompt': text });
@ -834,7 +837,7 @@ router.post('/remote/textgenerationwebui/encode', jsonParser, async function (re
}
const data = await result.json();
const count = legacyApi ? data?.results[0]?.tokens : (data?.length ?? data?.value ?? data?.tokens?.length);
const count = legacyApi ? data?.results[0]?.tokens : (data?.length ?? data?.count ?? data?.value ?? data?.tokens?.length);
const ids = legacyApi ? [] : (data?.tokens ?? data?.ids ?? []);
return response.send({ count, ids });