Add API token ids from KoboldCpp

This commit is contained in:
Cohee
2023-12-14 01:28:18 +02:00
parent b957e3b875
commit cebd6e9e0f
2 changed files with 32 additions and 2 deletions

View File

@ -63,6 +63,7 @@ const TOKENIZER_URLS = {
},
[tokenizers.API_KOBOLD]: {
count: '/api/tokenizers/remote/kobold/count',
encode: '/api/tokenizers/remote/kobold/count',
},
[tokenizers.MISTRAL]: {
encode: '/api/tokenizers/mistral/encode',
@ -617,6 +618,32 @@ function getTextTokensFromTextgenAPI(str) {
return ids;
}
/**
* Calls the AI provider's tokenize API to encode a string to tokens.
* @param {string} str String to tokenize.
* @returns {number[]} Array of token ids.
*/
function getTextTokensFromKoboldAPI(str) {
let ids = [];
jQuery.ajax({
async: false,
type: 'POST',
url: TOKENIZER_URLS[tokenizers.API_KOBOLD].encode,
data: JSON.stringify({
text: str,
url: api_server,
}),
dataType: 'json',
contentType: 'application/json',
success: function (data) {
ids = data.ids;
},
});
return ids;
}
/**
* Calls the underlying tokenizer model to decode token ids to text.
* @param {string} endpoint API endpoint.
@ -650,6 +677,8 @@ export function getTextTokens(tokenizerType, str) {
return getTextTokens(currentRemoteTokenizerAPI(), str);
case tokenizers.API_TEXTGENERATIONWEBUI:
return getTextTokensFromTextgenAPI(str);
case tokenizers.API_KOBOLD:
return getTextTokensFromKoboldAPI(str);
default: {
const tokenizerEndpoints = TOKENIZER_URLS[tokenizerType];
if (!tokenizerEndpoints) {