Merge pull request #2712 from SillyTavern/ai21-jamba New AI21 Jamba + tokenizer

3700396b145a4576836e8842b4dc57d673deb420

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed
Parents: ff834efde3bd173036ebef30790d4fc1c4ba0183 662e29911c6a07352051ae69f40556c8edeb6046
10 files changed, +192 -266Ignore whitespace
default/content/presets/openai/Default.json+1 -4
@@ -8,7 +8,7 @@
88 "openrouter_force_instruct": false,
99 "openrouter_group_models": false,
1010 "openrouter_sort_models": "alphabetically",
1111 "ai21_model": "j2jamba-ultra1.5-large",
1212 "mistralai_model": "mistral-large-latest",
1313 "custom_model": "",
1414 "custom_url": "",
@@ -19,7 +19,6 @@
1919 "temperature": 1,
2020 "frequency_penalty": 0,
2121 "presence_penalty": 0,
22- "count_penalty": 0,
2322 "top_p": 1,
2423 "top_k": 0,
2524 "top_a": 0,
@@ -233,8 +232,6 @@
233232 "assistant_prefill": "",
234233 "assistant_impersonation": "",
235234 "human_sysprompt_message": "Let's get started. Please generate your response based on the information and instructions provided above.",
236- "use_ai21_tokenizer": false,
237- "use_google_tokenizer": false,
238235 "claude_use_sysprompt": false,
239236 "use_alt_scale": false,
240237 "squash_system_messages": false,
default/content/settings.json+2 -4
@@ -452,7 +452,6 @@
452452 "temp_openai": 1.0,
453453 "freq_pen_openai": 0,
454454 "pres_pen_openai": 0,
455- "count_pen": 0,
456455 "top_p_openai": 1,
457456 "top_k_openai": 0,
458457 "stream_openai": true,
@@ -614,7 +613,7 @@
614613 "wi_format": "{0}",
615614 "openai_model": "gpt-4-turbo",
616615 "claude_model": "claude-3-5-sonnet-20240620",
617616 "ai21_model": "j2jamba-ultra1.5-large",
618617 "windowai_model": "",
619618 "openrouter_model": "OR_Website",
620619 "jailbreak_system": true,
@@ -625,7 +624,6 @@
625624 "show_external_models": false,
626625 "proxy_password": "",
627626 "assistant_prefill": "",
628627 "assistant_impersonation": "",
629- "use_ai21_tokenizer": false
630628 }
631629}
public/index.html+10 -36
@@ -436,7 +436,7 @@
436436 </div>
437437 </div>
438438 </div>
439439 <div data-newbie-hidden class="range-block" data-source="openai,openrouter,ai21,custom,cohere,perplexity,groq">
440440 <div class="range-block-title" data-i18n="Frequency Penalty">
441441 Frequency Penalty
442442 </div>
@@ -449,7 +449,7 @@
449449 </div>
450450 </div>
451451 </div>
452452 <div data-newbie-hidden class="range-block" data-source="openai,openrouter,ai21,custom,cohere,perplexity,groq">
453453 <div class="range-block-title" data-i18n="Presence Penalty">
454454 Presence Penalty
455455 </div>
@@ -462,20 +462,7 @@
462462 </div>
463463 </div>
464464 </div>
465465 <div data-newbie-hidden class="range-block" data-source="ai21claude,openrouter,makersuite,cohere,perplexity">
466- <div class="range-block-title" data-i18n="Count Penalty">
467- Count Penalty
468- </div>
469- <div class="range-block-range-and-counter">
470- <div class="range-block-range">
471- <input type="range" id="count_pen" name="volume" min="0" max="1" step="0.01">
472- </div>
473- <div class="range-block-counter">
474- <input type="number" min="0" max="1" step="0.01" data-for="count_pen" id="count_pen_counter">
475- </div>
476- </div>
477- </div>
478- <div data-newbie-hidden class="range-block" data-source="claude,openrouter,ai21,makersuite,cohere,perplexity">
479466 <div class="range-block-title" data-i18n="Top K">
480467 Top K
481468 </div>
@@ -1791,22 +1778,6 @@
17911778 </select>
17921779 </div>
17931780 </div>
1794- <div class="range-block" data-source="ai21">
1795- <label for="use_ai21_tokenizer" title="Use AI21 Tokenizer" data-i18n="[title]Use AI21 Tokenizer" class="checkbox_label widthFreeExpand">
1796- <input id="use_ai21_tokenizer" type="checkbox" /><span data-i18n="Use AI21 Tokenizer">Use AI21 Tokenizer</span>
1797- </label>
1798- <div class="toggle-description justifyLeft">
1799- <span data-i18n="Use the appropriate tokenizer for Jurassic models, which is more efficient than GPT's.">Use the appropriate tokenizer for Jurassic models, which is more efficient than GPT's.</span>
1800- </div>
1801- </div>
1802- <div class="range-block" data-source="makersuite">
1803- <label for="use_google_tokenizer" title="Use Google Tokenizer" data-i18n="[title]Use Google Tokenizer" class="checkbox_label widthFreeExpand">
1804- <input id="use_google_tokenizer" type="checkbox" /><span data-i18n="Use Google Tokenizer">Use Google Tokenizer</span>
1805- </label>
1806- <div class="toggle-description justifyLeft">
1807- <span data-i18n="Use the appropriate tokenizer for Google models via their API. Slower prompt processing, but offers much more accurate token counting.">Use the appropriate tokenizer for Google models via their API. Slower prompt processing, but offers much more accurate token counting.</span>
1808- </div>
1809- </div>
18101781 <div class="range-block" data-source="makersuite">
18111782 <label for="use_makersuite_sysprompt" class="checkbox_label widthFreeExpand">
18121783 <input id="use_makersuite_sysprompt" type="checkbox" />
@@ -2791,10 +2762,12 @@
27912762 <div>
27922763 <h4 data-i18n="AI21 Model">AI21 Model</h4>
27932764 <select id="model_ai21_select">
27942765 <optgroup label="LatestJamba 1.5">
27952766 <option value="j2jamba-ultra1.5-mini">j2jamba-ultra1.5-mini</option>
27962767 <option value="j2jamba-mid1.5-large">j2jamba-mid1.5-large</option>
2797- <option value="j2-light">j2-light</option>
2768+ </optgroup>
2769+ <optgroup label="Jamba-Instruct (Deprecated)">
2770+ <option value="jamba-instruct-preview">jamba-instruct-preview</option>
27982771 </optgroup>
27992772 </select>
28002773 </div>
@@ -3450,6 +3423,7 @@
34503423 <option value="3">Llama 1/2</option>
34513424 <option value="12">Llama 3</option>
34523425 <option value="13">Gemma / Gemini</option>
3426+ <option value="14">Jamba</option>
34533427 <option value="4">NerdStash (NovelAI Clio)</option>
34543428 <option value="5">NerdStash v2 (NovelAI Kayra)</option>
34553429 <option value="7">Mistral</option>
public/script.js+1 -1
@@ -2826,7 +2826,7 @@ export function getCharacterCardFields() {
28262826}
28272827
28282828export function isStreamingEnabled() {
28292829 const noStreamSources = [chat_completion_sources.SCALE, chat_completion_sources.AI21];
28302830 return ((main_api == 'openai' && oai_settings.stream_openai && !noStreamSources.includes(oai_settings.chat_completion_source) && !(oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE && oai_settings.google_model.includes('bison')))
28312831 || (main_api == 'kobold' && kai_settings.streaming_kobold && kai_flags.can_use_streaming)
28322832 || (main_api == 'novel' && nai_settings.streaming_novel)
public/scripts/openai.js+17 -101
@@ -132,13 +132,9 @@ const max_2mil = 2000 * 1000;
132132const scale_max = 8191;
133133const claude_max = 9000; // We have a proper tokenizer, so theoretically could be larger (up to 9k)
134134const claude_100k_max = 99000;
135-let ai21_max = 9200; //can easily fit 9k gpt tokens because j2's tokenizer is efficient af
136135const unlocked_max = max_200k;
137136const oai_max_temp = 2.0;
138137const claude_max_temp = 1.0; //same as j2
139-const j2_max_topk = 10.0;
140-const j2_max_freq = 5.0;
141-const j2_max_pres = 5.0;
142138const openrouter_website_model = 'OR_Website';
143139const openai_max_stop_strings = 4;
144140
@@ -218,25 +214,11 @@ const sensitiveFields = [
218214 'custom_include_headers',
219215];
220216
221-function getPrefixMap() {
222- return selected_group ? {
223- assistant: '',
224- user: '',
225- system: 'OOC: ',
226- }
227- : {
228- assistant: '{{char}}:',
229- user: '{{user}}:',
230- system: '',
231- };
232-}
233-
234217const default_settings = {
235218 preset_settings_openai: 'Default',
236219 temp_openai: 1.0,
237220 freq_pen_openai: 0,
238221 pres_pen_openai: 0,
239- count_pen: 0.0,
240222 top_p_openai: 1.0,
241223 top_k_openai: 0,
242224 min_p_openai: 0,
@@ -264,7 +246,7 @@ const default_settings = {
264246 openai_model: 'gpt-4-turbo',
265247 claude_model: 'claude-3-5-sonnet-20240620',
266248 google_model: 'gemini-1.5-pro',
267249 ai21_model: 'j2jamba-ultra1.5-large',
268250 mistralai_model: 'mistral-large-latest',
269251 cohere_model: 'command-r-plus',
270252 perplexity_model: 'llama-3.1-70b-instruct',
@@ -294,8 +276,6 @@ const default_settings = {
294276 assistant_prefill: '',
295277 assistant_impersonation: '',
296278 human_sysprompt_message: default_claude_human_sysprompt_message,
297- use_ai21_tokenizer: false,
298- use_google_tokenizer: false,
299279 claude_use_sysprompt: false,
300280 use_makersuite_sysprompt: true,
301281 use_alt_scale: false,
@@ -317,7 +297,6 @@ const oai_settings = {
317297 temp_openai: 1.0,
318298 freq_pen_openai: 0,
319299 pres_pen_openai: 0,
320- count_pen: 0.0,
321300 top_p_openai: 1.0,
322301 top_k_openai: 0,
323302 min_p_openai: 0,
@@ -345,7 +324,7 @@ const oai_settings = {
345324 openai_model: 'gpt-4-turbo',
346325 claude_model: 'claude-3-5-sonnet-20240620',
347326 google_model: 'gemini-1.5-pro',
348327 ai21_model: 'j2jamba-ultra1.5-large',
349328 mistralai_model: 'mistral-large-latest',
350329 cohere_model: 'command-r-plus',
351330 perplexity_model: 'llama-3.1-70b-instruct',
@@ -375,8 +354,6 @@ const oai_settings = {
375354 assistant_prefill: '',
376355 assistant_impersonation: '',
377356 human_sysprompt_message: default_claude_human_sysprompt_message,
378- use_ai21_tokenizer: false,
379- use_google_tokenizer: false,
380357 claude_use_sysprompt: false,
381358 use_makersuite_sysprompt: true,
382359 use_alt_scale: false,
@@ -1802,7 +1779,6 @@ async function sendOpenAIRequest(type, messages, signal) {
18021779 const isClaude = oai_settings.chat_completion_source == chat_completion_sources.CLAUDE;
18031780 const isOpenRouter = oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER;
18041781 const isScale = oai_settings.chat_completion_source == chat_completion_sources.SCALE;
1805- const isAI21 = oai_settings.chat_completion_source == chat_completion_sources.AI21;
18061782 const isGoogle = oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE;
18071783 const isOAI = oai_settings.chat_completion_source == chat_completion_sources.OPENAI;
18081784 const isMistral = oai_settings.chat_completion_source == chat_completion_sources.MISTRALAI;
@@ -1815,7 +1791,7 @@ async function sendOpenAIRequest(type, messages, signal) {
18151791 const isQuiet = type === 'quiet';
18161792 const isImpersonate = type === 'impersonate';
18171793 const isContinue = type === 'continue';
18181794 const stream = oai_settings.stream_openai && !isQuiet && !isScale && !isAI21 && !(isGoogle && oai_settings.google_model.includes('bison'));
18191795 const useLogprobs = !!power_user.request_token_probabilities;
18201796 const canMultiSwipe = oai_settings.n > 1 && !isContinue && !isImpersonate && !isQuiet && (isOAI || isCustom);
18211797
@@ -1824,15 +1800,6 @@ async function sendOpenAIRequest(type, messages, signal) {
18241800 replaceItemizedPromptText(messageId, messages);
18251801 }
18261802
1827- if (isAI21) {
1828- const joinedMsgs = messages.reduce((acc, obj) => {
1829- const prefix = getPrefixMap()[obj.role];
1830- return acc + (prefix ? (selected_group ? '\n' : prefix + ' ') : '') + obj.content + '\n';
1831- }, '');
1832- messages = substituteParams(joinedMsgs) + (isImpersonate ? `${name1}:` : `${name2}:`);
1833- replaceItemizedPromptText(messageId, messages);
1834- }
1835-
18361803 // If we're using the window.ai extension, use that instead
18371804 // Doesn't support logit bias yet
18381805 if (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI) {
@@ -1931,12 +1898,6 @@ async function sendOpenAIRequest(type, messages, signal) {
19311898 generate_data['use_makersuite_sysprompt'] = oai_settings.use_makersuite_sysprompt;
19321899 }
19331900
1934- if (isAI21) {
1935- generate_data['top_k'] = Number(oai_settings.top_k_openai);
1936- generate_data['count_pen'] = Number(oai_settings.count_pen);
1937- generate_data['stop_tokens'] = [name1 + ':', substituteParams(oai_settings.new_chat_prompt), substituteParams(oai_settings.new_group_chat_prompt)];
1938- }
1939-
19401901 if (isMistral) {
19411902 generate_data['safe_prompt'] = false; // already defaults to false, but just incase they change that in the future.
19421903 }
@@ -3008,7 +2969,6 @@ function loadOpenAISettings(data, settings) {
30082969 oai_settings.temp_openai = settings.temp_openai ?? default_settings.temp_openai;
30092970 oai_settings.freq_pen_openai = settings.freq_pen_openai ?? default_settings.freq_pen_openai;
30102971 oai_settings.pres_pen_openai = settings.pres_pen_openai ?? default_settings.pres_pen_openai;
3011- oai_settings.count_pen = settings.count_pen ?? default_settings.count_pen;
30122972 oai_settings.top_p_openai = settings.top_p_openai ?? default_settings.top_p_openai;
30132973 oai_settings.top_k_openai = settings.top_k_openai ?? default_settings.top_k_openai;
30142974 oai_settings.top_a_openai = settings.top_a_openai ?? default_settings.top_a_openai;
@@ -3080,10 +3040,12 @@ function loadOpenAISettings(data, settings) {
30803040 oai_settings.names_behavior = character_names_behavior.COMPLETION;
30813041 }
30823042
3043+ if (oai_settings.ai21_model.startsWith('j2-')) {
3044+ oai_settings.ai21_model = 'jamba-1.5-large';
3045+ }
3046+
30833047 if (settings.wrap_in_quotes !== undefined) oai_settings.wrap_in_quotes = !!settings.wrap_in_quotes;
30843048 if (settings.openai_model !== undefined) oai_settings.openai_model = settings.openai_model;
3085- if (settings.use_ai21_tokenizer !== undefined) { oai_settings.use_ai21_tokenizer = !!settings.use_ai21_tokenizer; oai_settings.use_ai21_tokenizer ? ai21_max = 8191 : ai21_max = 9200; }
3086- if (settings.use_google_tokenizer !== undefined) oai_settings.use_google_tokenizer = !!settings.use_google_tokenizer;
30873049 if (settings.claude_use_sysprompt !== undefined) oai_settings.claude_use_sysprompt = !!settings.claude_use_sysprompt;
30883050 if (settings.use_makersuite_sysprompt !== undefined) oai_settings.use_makersuite_sysprompt = !!settings.use_makersuite_sysprompt;
30893051 if (settings.use_alt_scale !== undefined) { oai_settings.use_alt_scale = !!settings.use_alt_scale; updateScaleForm(); }
@@ -3133,8 +3095,6 @@ function loadOpenAISettings(data, settings) {
31333095 $('#jailbreak_system').prop('checked', oai_settings.jailbreak_system);
31343096 $('#openai_show_external_models').prop('checked', oai_settings.show_external_models);
31353097 $('#openai_external_category').toggle(oai_settings.show_external_models);
3136- $('#use_ai21_tokenizer').prop('checked', oai_settings.use_ai21_tokenizer);
3137- $('#use_google_tokenizer').prop('checked', oai_settings.use_google_tokenizer);
31383098 $('#claude_use_sysprompt').prop('checked', oai_settings.claude_use_sysprompt);
31393099 $('#use_makersuite_sysprompt').prop('checked', oai_settings.use_makersuite_sysprompt);
31403100 $('#scale-alt').prop('checked', oai_settings.use_alt_scale);
@@ -3170,9 +3130,6 @@ function loadOpenAISettings(data, settings) {
31703130 $('#pres_pen_openai').val(oai_settings.pres_pen_openai);
31713131 $('#pres_pen_counter_openai').val(Number(oai_settings.pres_pen_openai).toFixed(2));
31723132
3173- $('#count_pen').val(oai_settings.count_pen);
3174- $('#count_pen_counter').val(Number(oai_settings.count_pen).toFixed(2));
3175-
31763133 $('#top_p_openai').val(oai_settings.top_p_openai);
31773134 $('#top_p_counter_openai').val(Number(oai_settings.top_p_openai).toFixed(2));
31783135
@@ -3392,7 +3349,6 @@ async function saveOpenAIPreset(name, settings, triggerUi = true) {
33923349 temperature: settings.temp_openai,
33933350 frequency_penalty: settings.freq_pen_openai,
33943351 presence_penalty: settings.pres_pen_openai,
3395- count_penalty: settings.count_pen,
33963352 top_p: settings.top_p_openai,
33973353 top_k: settings.top_k_openai,
33983354 top_a: settings.top_a_openai,
@@ -3427,8 +3383,6 @@ async function saveOpenAIPreset(name, settings, triggerUi = true) {
34273383 assistant_prefill: settings.assistant_prefill,
34283384 assistant_impersonation: settings.assistant_impersonation,
34293385 human_sysprompt_message: settings.human_sysprompt_message,
3430- use_ai21_tokenizer: settings.use_ai21_tokenizer,
3431- use_google_tokenizer: settings.use_google_tokenizer,
34323386 claude_use_sysprompt: settings.claude_use_sysprompt,
34333387 use_makersuite_sysprompt: settings.use_makersuite_sysprompt,
34343388 use_alt_scale: settings.use_alt_scale,
@@ -3799,7 +3753,6 @@ function onSettingsPresetChange() {
37993753 temperature: ['#temp_openai', 'temp_openai', false],
38003754 frequency_penalty: ['#freq_pen_openai', 'freq_pen_openai', false],
38013755 presence_penalty: ['#pres_pen_openai', 'pres_pen_openai', false],
3802- count_penalty: ['#count_pen', 'count_pen', false],
38033756 top_p: ['#top_p_openai', 'top_p_openai', false],
38043757 top_k: ['#top_k_openai', 'top_k_openai', false],
38053758 top_a: ['#top_a_openai', 'top_a_openai', false],
@@ -3856,8 +3809,6 @@ function onSettingsPresetChange() {
38563809 assistant_prefill: ['#claude_assistant_prefill', 'assistant_prefill', false],
38573810 assistant_impersonation: ['#claude_assistant_impersonation', 'assistant_impersonation', false],
38583811 human_sysprompt_message: ['#claude_human_sysprompt_textarea', 'human_sysprompt_message', false],
3859- use_ai21_tokenizer: ['#use_ai21_tokenizer', 'use_ai21_tokenizer', true],
3860- use_google_tokenizer: ['#use_google_tokenizer', 'use_google_tokenizer', true],
38613812 claude_use_sysprompt: ['#claude_use_sysprompt', 'claude_use_sysprompt', true],
38623813 use_makersuite_sysprompt: ['#use_makersuite_sysprompt', 'use_makersuite_sysprompt', true],
38633814 use_alt_scale: ['#use_alt_scale', 'use_alt_scale', true],
@@ -4028,6 +3979,11 @@ async function onModelChange() {
40283979 }
40293980
40303981 if ($(this).is('#model_ai21_select')) {
3982+ if (value === '' || value.startsWith('j2-')) {
3983+ value = 'jamba-1.5-large';
3984+ $('#model_ai21_select').val(value);
3985+ }
3986+
40313987 console.log('AI21 model changed to', value);
40323988 oai_settings.ai21_model = value;
40333989 }
@@ -4305,33 +4261,13 @@ async function onModelChange() {
43054261 if (oai_settings.chat_completion_source == chat_completion_sources.AI21) {
43064262 if (oai_settings.max_context_unlocked) {
43074263 $('#openai_max_context').attr('max', unlocked_max);
4308- } else {
4264+ } else if (oai_settings.ai21_model.includes('jamba-1.5') || oai_settings.ai21_model.includes('jamba-instruct')) {
43094265 $('#openai_max_context').attr('max', ai21_maxmax_256k);
43104266 }
43114267
43124268 oai_settings.openai_max_context = Math.min(oai_settings.openai_max_context, Number($('#openai_max_context').attr('max')), oai_settings.openai_max_context);
43134269 $('#openai_max_context').val(oai_settings.openai_max_context).trigger('input');
4314-
4270+ $('#temp_openai').attr('max', oai_max_temp).val(oai_settings.temp_openai).trigger('input');
4315- oai_settings.temp_openai = Math.min(claude_max_temp, oai_settings.temp_openai);
4316- $('#temp_openai').attr('max', claude_max_temp).val(oai_settings.temp_openai).trigger('input');
4317-
4318- oai_settings.freq_pen_openai = Math.min(j2_max_freq, oai_settings.freq_pen_openai < 0 ? 0 : oai_settings.freq_pen_openai);
4319- $('#freq_pen_openai').attr('min', 0).attr('max', j2_max_freq).val(oai_settings.freq_pen_openai).trigger('input');
4320-
4321- oai_settings.pres_pen_openai = Math.min(j2_max_pres, oai_settings.pres_pen_openai < 0 ? 0 : oai_settings.pres_pen_openai);
4322- $('#pres_pen_openai').attr('min', 0).attr('max', j2_max_pres).val(oai_settings.pres_pen_openai).trigger('input');
4323-
4324- oai_settings.top_k_openai = Math.min(j2_max_topk, oai_settings.top_k_openai);
4325- $('#top_k_openai').attr('max', j2_max_topk).val(oai_settings.top_k_openai).trigger('input');
4326- } else if (oai_settings.chat_completion_source != chat_completion_sources.AI21) {
4327- oai_settings.freq_pen_openai = Math.min(2.0, oai_settings.freq_pen_openai);
4328- $('#freq_pen_openai').attr('min', -2.0).attr('max', 2.0).val(oai_settings.freq_pen_openai).trigger('input');
4329-
4330- oai_settings.pres_pen_openai = Math.min(2.0, oai_settings.pres_pen_openai);
4331- $('#pres_pen_openai').attr('min', -2.0).attr('max', 2.0).val(oai_settings.pres_pen_openai).trigger('input');
4332-
4333- oai_settings.top_k_openai = Math.min(200, oai_settings.top_k_openai);
4334- $('#top_k_openai').attr('max', 200).val(oai_settings.top_k_openai).trigger('input');
43354271 }
43364272
43374273 if (oai_settings.chat_completion_source == chat_completion_sources.CUSTOM) {
@@ -4953,12 +4889,6 @@ $(document).ready(async function () {
49534889 saveSettingsDebounced();
49544890 });
49554891
4956- $('#count_pen').on('input', function () {
4957- oai_settings.count_pen = Number($(this).val());
4958- $('#count_pen_counter').val(Number($(this).val()).toFixed(2));
4959- saveSettingsDebounced();
4960- });
4961-
49624892 $('#top_p_openai').on('input', function () {
49634893 oai_settings.top_p_openai = Number($(this).val());
49644894 $('#top_p_counter_openai').val(Number($(this).val()).toFixed(2));
@@ -5017,20 +4947,6 @@ $(document).ready(async function () {
50174947 saveSettingsDebounced();
50184948 });
50194949
5020- $('#use_ai21_tokenizer').on('change', function () {
5021- oai_settings.use_ai21_tokenizer = !!$('#use_ai21_tokenizer').prop('checked');
5022- oai_settings.use_ai21_tokenizer ? ai21_max = 8191 : ai21_max = 9200;
5023- oai_settings.openai_max_context = Math.min(ai21_max, oai_settings.openai_max_context);
5024- $('#openai_max_context').attr('max', ai21_max).val(oai_settings.openai_max_context).trigger('input');
5025- $('#openai_max_context_counter').attr('max', Number($('#openai_max_context').attr('max')));
5026- saveSettingsDebounced();
5027- });
5028-
5029- $('#use_google_tokenizer').on('change', function () {
5030- oai_settings.use_google_tokenizer = !!$('#use_google_tokenizer').prop('checked');
5031- saveSettingsDebounced();
5032- });
5033-
50344950 $('#claude_use_sysprompt').on('change', function () {
50354951 oai_settings.claude_use_sysprompt = !!$('#claude_use_sysprompt').prop('checked');
50364952 $('#claude_human_sysprompt_message_block').toggle(oai_settings.claude_use_sysprompt);
public/scripts/tokenizers.js+26 -23
@@ -27,6 +27,7 @@ export const tokenizers = {
2727 CLAUDE: 11,
2828 LLAMA3: 12,
2929 GEMMA: 13,
30+ JAMBA: 14,
3031 BEST_MATCH: 99,
3132};
3233
@@ -36,6 +37,7 @@ export const SENTENCEPIECE_TOKENIZERS = [
3637 tokenizers.YI,
3738 tokenizers.LLAMA3,
3839 tokenizers.GEMMA,
40+ tokenizers.JAMBA,
3941 // uncomment when NovelAI releases Kayra and Clio weights, lol
4042 //tokenizers.NERD,
4143 //tokenizers.NERD2,
@@ -98,6 +100,11 @@ const TOKENIZER_URLS = {
98100 decode: '/api/tokenizers/gemma/decode',
99101 count: '/api/tokenizers/gemma/encode',
100102 },
103+ [tokenizers.JAMBA]: {
104+ encode: '/api/tokenizers/jamba/encode',
105+ decode: '/api/tokenizers/jamba/decode',
106+ count: '/api/tokenizers/jamba/encode',
107+ },
101108 [tokenizers.API_TEXTGENERATIONWEBUI]: {
102109 encode: '/api/tokenizers/remote/textgenerationwebui/encode',
103110 count: '/api/tokenizers/remote/textgenerationwebui/encode',
@@ -164,7 +171,7 @@ export function getAvailableTokenizers() {
164171 tokenizerId: Number(tokenizerOption.value),
165172 tokenizerKey: Object.entries(tokenizers).find(([_, value]) => value === Number(tokenizerOption.value))[0].toLocaleLowerCase(),
166173 tokenizerName: tokenizerOption.text,
167174 }));
168175}
169176
170177/**
@@ -280,6 +287,12 @@ export function getTokenizerBestMatch(forApi) {
280287 if (model.includes('gemma')) {
281288 return tokenizers.GEMMA;
282289 }
290+ if (model.includes('yi')) {
291+ return tokenizers.YI;
292+ }
293+ if (model.includes('jamba')) {
294+ return tokenizers.JAMBA;
295+ }
283296 }
284297
285298 return tokenizers.LLAMA;
@@ -497,6 +510,7 @@ export function getTokenizerModel() {
497510 const mistralTokenizer = 'mistral';
498511 const yiTokenizer = 'yi';
499512 const gemmaTokenizer = 'gemma';
513+ const jambaTokenizer = 'jamba';
500514
501515 // Assuming no one would use it for different models.. right?
502516 if (oai_settings.chat_completion_source == chat_completion_sources.SCALE) {
@@ -562,12 +576,19 @@ export function getTokenizerModel() {
562576 else if (oai_settings.openrouter_model.includes('GPT-NeoXT')) {
563577 return gpt2Tokenizer;
564578 }
579+ else if (oai_settings.openrouter_model.includes('jamba')) {
580+ return jambaTokenizer;
581+ }
565582 }
566583
567584 if (oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE) {
568585 return gemmaTokenizer;
569586 }
570587
588+ if (oai_settings.chat_completion_source == chat_completion_sources.AI21) {
589+ return jambaTokenizer;
590+ }
591+
571592 if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {
572593 return claudeTokenizer;
573594 }
@@ -626,16 +647,7 @@ export function getTokenizerModel() {
626647 * @deprecated Use countTokensOpenAIAsync instead.
627648 */
628649export function countTokensOpenAI(messages, full = false) {
629- const shouldTokenizeAI21 = oai_settings.chat_completion_source === chat_completion_sources.AI21 && oai_settings.use_ai21_tokenizer;
650+ const tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
630- const shouldTokenizeGoogle = oai_settings.chat_completion_source === chat_completion_sources.MAKERSUITE && oai_settings.use_google_tokenizer;
631- let tokenizerEndpoint = '';
632- if (shouldTokenizeAI21) {
633- tokenizerEndpoint = '/api/tokenizers/ai21/count';
634- } else if (shouldTokenizeGoogle) {
635- tokenizerEndpoint = `/api/tokenizers/google/count?model=${getTokenizerModel()}&reverse_proxy=${oai_settings.reverse_proxy}&proxy_password=${oai_settings.proxy_password}`;
636- } else {
637- tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
638- }
639651 const cacheObject = getTokenCacheObject();
640652
641653 if (!Array.isArray(messages)) {
@@ -647,7 +659,7 @@ export function countTokensOpenAI(messages, full = false) {
647659 for (const message of messages) {
648660 const model = getTokenizerModel();
649661
650662 if (model === 'claude' || shouldTokenizeAI21 || shouldTokenizeGoogle) {
651663 full = true;
652664 }
653665
@@ -687,16 +699,7 @@ export function countTokensOpenAI(messages, full = false) {
687699 * @returns {Promise<number>} Token count.
688700 */
689701export async function countTokensOpenAIAsync(messages, full = false) {
690- const shouldTokenizeAI21 = oai_settings.chat_completion_source === chat_completion_sources.AI21 && oai_settings.use_ai21_tokenizer;
702+ const tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
691- const shouldTokenizeGoogle = oai_settings.chat_completion_source === chat_completion_sources.MAKERSUITE && oai_settings.use_google_tokenizer;
692- let tokenizerEndpoint = '';
693- if (shouldTokenizeAI21) {
694- tokenizerEndpoint = '/api/tokenizers/ai21/count';
695- } else if (shouldTokenizeGoogle) {
696- tokenizerEndpoint = `/api/tokenizers/google/count?model=${getTokenizerModel()}`;
697- } else {
698- tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
699- }
700703 const cacheObject = getTokenCacheObject();
701704
702705 if (!Array.isArray(messages)) {
@@ -708,7 +711,7 @@ export async function countTokensOpenAIAsync(messages, full = false) {
708711 for (const message of messages) {
709712 const model = getTokenizerModel();
710713
711714 if (model === 'claude' || shouldTokenizeAI21 || shouldTokenizeGoogle) {
712715 full = true;
713716 }
714717
src/endpoints/backends/chat-completions.js+35 -50
@@ -5,7 +5,7 @@ const Readable = require('stream').Readable;
55const { jsonParser } = require('../../express-common');
66const { CHAT_COMPLETION_SOURCES, GEMINI_SAFETY, BISON_SAFETY, OPENROUTER_HEADERS } = require('../../constants');
77const { forwardFetchResponse, getConfigValue, tryParse, uuidv4, mergeObjectWithYaml, excludeKeysByYaml, color } = require('../../util');
88const { convertClaudeMessages, convertGooglePrompt, convertTextCompletionPrompt, convertCohereMessages, convertMistralMessages, convertCohereTools, convertAI21Messages } = require('../../prompt-converters');
99
1010const { readSecret, SECRET_KEYS } = require('../secrets');
1111const { getTokenizerModel, getSentencepiceTokenizer, getTiktokenTokenizer, sentencepieceTokenizers, TEXT_COMPLETION_MODELS } = require('../tokenizers');
@@ -19,6 +19,7 @@ const API_GROQ = 'https://api.groq.com/openai/v1';
1919const API_MAKERSUITE = 'https://generativelanguage.googleapis.com';
2020const API_01AI = 'https://api.01.ai/v1';
2121const API_BLOCKENTROPY = 'https://api.blockentropy.ai/v1';
22+const API_AI21 = 'https://api.ai21.com/studio/v1';
2223
2324/**
2425 * Applies a post-processing step to the generated messages.
@@ -413,6 +414,16 @@ async function sendAI21Request(request, response) {
413414 request.socket.on('close', function () {
414415 controller.abort();
415416 });
417+ const convertedPrompt = convertAI21Messages(request.body.messages, request.body.char_name, request.body.user_name);
418+ const body = {
419+ messages: convertedPrompt,
420+ model: request.body.model,
421+ max_tokens: request.body.max_tokens,
422+ temperature: request.body.temperature,
423+ top_p: request.body.top_p,
424+ stop: request.body.stop,
425+ stream: request.body.stream,
426+ };
416427 const options = {
417428 method: 'POST',
418429 headers: {
@@ -420,59 +431,33 @@ async function sendAI21Request(request, response) {
420431 'content-type': 'application/json',
421432 Authorization: `Bearer ${readSecret(request.user.directories, SECRET_KEYS.AI21)}`,
422433 },
423434 body: JSON.stringify({body),
424- numResults: 1,
425- maxTokens: request.body.max_tokens,
426- minTokens: 0,
427- temperature: request.body.temperature,
428- topP: request.body.top_p,
429- stopSequences: request.body.stop_tokens,
430- topKReturn: request.body.top_k,
431- frequencyPenalty: {
432- scale: request.body.frequency_penalty * 100,
433- applyToWhitespaces: false,
434- applyToPunctuations: false,
435- applyToNumbers: false,
436- applyToStopwords: false,
437- applyToEmojis: false,
438- },
439- presencePenalty: {
440- scale: request.body.presence_penalty,
441- applyToWhitespaces: false,
442- applyToPunctuations: false,
443- applyToNumbers: false,
444- applyToStopwords: false,
445- applyToEmojis: false,
446- },
447- countPenalty: {
448- scale: request.body.count_pen,
449- applyToWhitespaces: false,
450- applyToPunctuations: false,
451- applyToNumbers: false,
452- applyToStopwords: false,
453- applyToEmojis: false,
454- },
455- prompt: request.body.messages,
456- }),
457435 signal: controller.signal,
458436 };
459437
460- fetch(`https://api.ai21.com/studio/v1/${request.body.model}/complete`, options)
438+ console.log('AI21 request:', body);
461- .then(r => r.json())
462- .then(r => {
463- if (r.completions === undefined) {
464- console.log(r);
465- } else {
466- console.log(r.completions[0].data.text);
467- }
468- const reply = { choices: [{ 'message': { 'content': r.completions?.[0]?.data?.text } }] };
469- return response.send(reply);
470- })
471- .catch(err => {
472- console.error(err);
473- return response.send({ error: true });
474- });
475439
440+ try{
441+ const generateResponse = await fetch(API_AI21 + '/chat/completions', options);
442+ if (request.body.stream) {
443+ forwardFetchResponse(generateResponse, response);
444+ } else {
445+ if (!generateResponse.ok) {
446+ console.log(`AI21 API returned error: ${generateResponse.status} ${generateResponse.statusText} ${await generateResponse.text()}`);
447+ return response.status(500).send({ error: true });
448+ }
449+ const generateResponseJson = await generateResponse.json();
450+ console.log('AI21 response:', generateResponseJson);
451+ return response.send(generateResponseJson);
452+ }
453+ } catch (error) {
454+ console.log('Error communicating with MistralAI API: ', error);
455+ if (!response.headersSent) {
456+ response.send({ error: true });
457+ } else {
458+ response.end();
459+ }
460+ }
476461}
477462
478463/**
src/endpoints/tokenizers.js+27 -47
@@ -144,6 +144,7 @@ const spp_nerd_v2 = new SentencePieceTokenizer('src/tokenizers/nerdstash_v2.mode
144144const spp_mistral = new SentencePieceTokenizer('src/tokenizers/mistral.model');
145145const spp_yi = new SentencePieceTokenizer('src/tokenizers/yi.model');
146146const spp_gemma = new SentencePieceTokenizer('src/tokenizers/gemma.model');
147+const spp_jamba = new SentencePieceTokenizer('src/tokenizers/jamba.model');
147148const claude_tokenizer = new WebTokenizer('src/tokenizers/claude.json');
148149const llama3_tokenizer = new WebTokenizer('src/tokenizers/llama3.json');
149150
@@ -154,6 +155,7 @@ const sentencepieceTokenizers = [
154155 'mistral',
155156 'yi',
156157 'gemma',
158+ 'jamba',
157159];
158160
159161/**
@@ -186,6 +188,10 @@ function getSentencepiceTokenizer(model) {
186188 return spp_gemma;
187189 }
188190
191+ if (model.includes('jamba')) {
192+ return spp_jamba;
193+ }
194+
189195 return null;
190196}
191197
@@ -322,6 +328,10 @@ function getTokenizerModel(requestModel) {
322328 return 'gemma';
323329 }
324330
331+ if (requestModel.includes('jamba')) {
332+ return 'jamba';
333+ }
334+
325335 // default
326336 return 'gpt-3.5-turbo';
327337}
@@ -537,59 +547,13 @@ function createWebTokenizerDecodingHandler(tokenizer) {
537547
538548const router = express.Router();
539549
540-router.post('/ai21/count', jsonParser, async function (req, res) {
541- if (!req.body) return res.sendStatus(400);
542- const key = readSecret(req.user.directories, SECRET_KEYS.AI21);
543- const options = {
544- method: 'POST',
545- headers: {
546- accept: 'application/json',
547- 'content-type': 'application/json',
548- Authorization: `Bearer ${key}`,
549- },
550- body: JSON.stringify({ text: req.body[0].content }),
551- };
552-
553- try {
554- const response = await fetch('https://api.ai21.com/studio/v1/tokenize', options);
555- const data = await response.json();
556- return res.send({ 'token_count': data?.tokens?.length || 0 });
557- } catch (err) {
558- console.error(err);
559- return res.send({ 'token_count': 0 });
560- }
561-});
562-
563-router.post('/google/count', jsonParser, async function (req, res) {
564- if (!req.body) return res.sendStatus(400);
565- const options = {
566- method: 'POST',
567- headers: {
568- accept: 'application/json',
569- 'content-type': 'application/json',
570- },
571- body: JSON.stringify({ contents: convertGooglePrompt(req.body, String(req.query.model)).contents }),
572- };
573- try {
574- const reverseProxy = req.query.reverse_proxy?.toString() || '';
575- const proxyPassword = req.query.proxy_password?.toString() || '';
576- const apiKey = reverseProxy ? proxyPassword : readSecret(req.user.directories, SECRET_KEYS.MAKERSUITE);
577- const apiUrl = new URL(reverseProxy || API_MAKERSUITE);
578- const response = await fetch(`${apiUrl.origin}/v1beta/models/${req.query.model}:countTokens?key=${apiKey}`, options);
579- const data = await response.json();
580- return res.send({ 'token_count': data?.totalTokens || 0 });
581- } catch (err) {
582- console.error(err);
583- return res.send({ 'token_count': 0 });
584- }
585-});
586-
587550router.post('/llama/encode', jsonParser, createSentencepieceEncodingHandler(spp_llama));
588551router.post('/nerdstash/encode', jsonParser, createSentencepieceEncodingHandler(spp_nerd));
589552router.post('/nerdstash_v2/encode', jsonParser, createSentencepieceEncodingHandler(spp_nerd_v2));
590553router.post('/mistral/encode', jsonParser, createSentencepieceEncodingHandler(spp_mistral));
591554router.post('/yi/encode', jsonParser, createSentencepieceEncodingHandler(spp_yi));
592555router.post('/gemma/encode', jsonParser, createSentencepieceEncodingHandler(spp_gemma));
556+router.post('/jamba/encode', jsonParser, createSentencepieceEncodingHandler(spp_jamba));
593557router.post('/gpt2/encode', jsonParser, createTiktokenEncodingHandler('gpt2'));
594558router.post('/claude/encode', jsonParser, createWebTokenizerEncodingHandler(claude_tokenizer));
595559router.post('/llama3/encode', jsonParser, createWebTokenizerEncodingHandler(llama3_tokenizer));
@@ -599,6 +563,7 @@ router.post('/nerdstash_v2/decode', jsonParser, createSentencepieceDecodingHandl
599563router.post('/mistral/decode', jsonParser, createSentencepieceDecodingHandler(spp_mistral));
600564router.post('/yi/decode', jsonParser, createSentencepieceDecodingHandler(spp_yi));
601565router.post('/gemma/decode', jsonParser, createSentencepieceDecodingHandler(spp_gemma));
566+router.post('/jamba/decode', jsonParser, createSentencepieceDecodingHandler(spp_jamba));
602567router.post('/gpt2/decode', jsonParser, createTiktokenDecodingHandler('gpt2'));
603568router.post('/claude/decode', jsonParser, createWebTokenizerDecodingHandler(claude_tokenizer));
604569router.post('/llama3/decode', jsonParser, createWebTokenizerDecodingHandler(llama3_tokenizer));
@@ -637,6 +602,11 @@ router.post('/openai/encode', jsonParser, async function (req, res) {
637602 return handler(req, res);
638603 }
639604
605+ if (queryModel.includes('jamba')) {
606+ const handler = createSentencepieceEncodingHandler(spp_jamba);
607+ return handler(req, res);
608+ }
609+
640610 const model = getTokenizerModel(queryModel);
641611 const handler = createTiktokenEncodingHandler(model);
642612 return handler(req, res);
@@ -680,6 +650,11 @@ router.post('/openai/decode', jsonParser, async function (req, res) {
680650 return handler(req, res);
681651 }
682652
653+ if (queryModel.includes('jamba')) {
654+ const handler = createSentencepieceDecodingHandler(spp_jamba);
655+ return handler(req, res);
656+ }
657+
683658 const model = getTokenizerModel(queryModel);
684659 const handler = createTiktokenDecodingHandler(model);
685660 return handler(req, res);
@@ -731,6 +706,11 @@ router.post('/openai/count', jsonParser, async function (req, res) {
731706 return res.send({ 'token_count': num_tokens });
732707 }
733708
709+ if (model === 'jamba') {
710+ num_tokens = await countSentencepieceArrayTokens(spp_jamba, req.body);
711+ return res.send({ 'token_count': num_tokens });
712+ }
713+
734714 const tokensPerName = queryModel.includes('gpt-3.5-turbo-0301') ? -1 : 1;
735715 const tokensPerMessage = queryModel.includes('gpt-3.5-turbo-0301') ? 4 : 3;
736716 const tokensPadding = 3;
src/prompt-converters.js+73 -0
@@ -368,6 +368,78 @@ function convertGooglePrompt(messages, model, useSysPrompt = false, charName = '
368368}
369369
370370/**
371+ * Convert AI21 prompt. Classic: system message squash, user/assistant message merge.
372+ * @param {object[]} messages Array of messages
373+ * @param {string} charName Character name
374+ * @param {string} userName User name
375+ */
376+function convertAI21Messages(messages, charName = '', userName = '') {
377+ if (!Array.isArray(messages)) {
378+ return [];
379+ }
380+
381+ // Collect all the system messages up until the first instance of a non-system message, and then remove them from the messages array.
382+ let i = 0, systemPrompt = '';
383+
384+ for (i = 0; i < messages.length; i++) {
385+ if (messages[i].role !== 'system') {
386+ break;
387+ }
388+ // Append example names if not already done by the frontend (e.g. for group chats).
389+ if (userName && messages[i].name === 'example_user') {
390+ if (!messages[i].content.startsWith(`${userName}: `)) {
391+ messages[i].content = `${userName}: ${messages[i].content}`;
392+ }
393+ }
394+ if (charName && messages[i].name === 'example_assistant') {
395+ if (!messages[i].content.startsWith(`${charName}: `)) {
396+ messages[i].content = `${charName}: ${messages[i].content}`;
397+ }
398+ }
399+ systemPrompt += `${messages[i].content}\n\n`;
400+ }
401+
402+ messages.splice(0, i);
403+
404+ // Prevent erroring out if the messages array is empty.
405+ if (messages.length === 0) {
406+ messages.unshift({
407+ role: 'user',
408+ content: '[Start a new chat]',
409+ });
410+ }
411+
412+ if (systemPrompt) {
413+ messages.unshift({
414+ role: 'system',
415+ content: systemPrompt.trim(),
416+ });
417+ }
418+
419+ // Doesn't support completion names, so prepend if not already done by the frontend (e.g. for group chats).
420+ messages.forEach(msg => {
421+ if ('name' in msg) {
422+ if (msg.role !== 'system' && !msg.content.startsWith(`${msg.name}: `)) {
423+ msg.content = `${msg.name}: ${msg.content}`;
424+ }
425+ delete msg.name;
426+ }
427+ });
428+
429+ // Since the messaging endpoint only supports alternating turns, we have to merge messages with the same role if they follow each other
430+ let mergedMessages = [];
431+ messages.forEach((message) => {
432+ if (mergedMessages.length > 0 && mergedMessages[mergedMessages.length - 1].role === message.role) {
433+ mergedMessages[mergedMessages.length - 1].content += '\n\n' + message.content;
434+ } else {
435+ mergedMessages.push(message);
436+ }
437+ });
438+
439+ return mergedMessages;
440+}
441+
442+/**
371443 * Convert a prompt from the ChatML objects to the format used by MistralAI.
372444 * @param {object[]} messages Array of messages
373445 * @param {string} charName Character name
@@ -520,4 +592,5 @@ module.exports = {
520592 convertCohereMessages,
521593 convertMistralMessages,
522594 convertCohereTools,
595+ convertAI21Messages,
523596};
src/tokenizers/jamba.model+0 -0

Binary file