Fix updated tokenization via ooba API

This commit is contained in:
Cohee 2023-11-09 19:39:08 +02:00
parent 3780321ae6
commit 81fe9aa699
2 changed files with 22 additions and 13 deletions

View File

@ -363,6 +363,15 @@ function getTokenCacheObject() {
return tokenCache[String(chatId)]; return tokenCache[String(chatId)];
} }
function getRemoteTokenizationParams(str) {
return {
text: str,
api: main_api,
url: getAPIServerUrl(),
legacy_api: main_api === 'textgenerationwebui' && textgenerationwebui_settings.legacy_api && !isMancer(),
};
}
/** /**
* Counts token using the remote server API. * Counts token using the remote server API.
* @param {string} endpoint API endpoint. * @param {string} endpoint API endpoint.
@ -377,12 +386,7 @@ function countTokensRemote(endpoint, str, padding) {
async: false, async: false,
type: 'POST', type: 'POST',
url: endpoint, url: endpoint,
data: JSON.stringify({ data: JSON.stringify(getRemoteTokenizationParams(str)),
text: str,
api: main_api,
url: getAPIServerUrl(),
legacy_api: main_api === 'textgenerationwebui' && textgenerationwebui_settings.legacy_api && !isMancer() ,
}),
dataType: "json", dataType: "json",
contentType: "application/json", contentType: "application/json",
success: function (data) { success: function (data) {
@ -425,7 +429,7 @@ function getTextTokensRemote(endpoint, str, model = '') {
async: false, async: false,
type: 'POST', type: 'POST',
url: endpoint, url: endpoint,
data: JSON.stringify({ text: str }), data: JSON.stringify(getRemoteTokenizationParams(str)),
dataType: "json", dataType: "json",
contentType: "application/json", contentType: "application/json",
success: function (data) { success: function (data) {
@ -482,6 +486,8 @@ export function getTextTokens(tokenizerType, str) {
case tokenizers.OPENAI: case tokenizers.OPENAI:
const model = getTokenizerModel(); const model = getTokenizerModel();
return getTextTokensRemote('/api/tokenize/openai-encode', str, model); return getTextTokensRemote('/api/tokenize/openai-encode', str, model);
case tokenizers.API:
return getTextTokensRemote('/tokenize_via_api', str);
default: default:
console.warn("Calling getTextTokens with unsupported tokenizer type", tokenizerType); console.warn("Calling getTextTokens with unsupported tokenizer type", tokenizerType);
return []; return [];

View File

@ -3384,8 +3384,7 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) {
if (api == 'textgenerationwebui') { if (api == 'textgenerationwebui') {
const args = { const args = {
method: 'POST', method: 'POST',
body: JSON.stringify({ "prompt": text }), headers: { "Content-Type": "application/json" },
headers: { "Content-Type": "application/json" }
}; };
setAdditionalHeaders(request, args, null); setAdditionalHeaders(request, args, null);
@ -3395,9 +3394,10 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) {
if (legacyApi) { if (legacyApi) {
url += '/v1/token-count'; url += '/v1/token-count';
args.body = JSON.stringify({ "prompt": text});
} else { } else {
url += '/api/v1/token-count'; url += '/v1/internal/encode';
args.body = JSON.stringify({ "text": text });
} }
const result = await fetch(url, args); const result = await fetch(url, args);
@ -3408,7 +3408,10 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) {
} }
const data = await result.json(); const data = await result.json();
return response.send({ count: data['results'][0]['tokens'] }); const count = legacyApi ? data?.results[0]?.tokens : data?.length;
const ids = legacyApi ? [] : data?.tokens;
return response.send({ count, ids });
} }
else if (api == 'kobold') { else if (api == 'kobold') {
@ -3430,7 +3433,7 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) {
const data = await result.json(); const data = await result.json();
const count = data['value']; const count = data['value'];
return response.send({ count: count }); return response.send({ count: count, ids: [] });
} }
else { else {