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

3700396b145a4576836e8842b4dc57d673deb420

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

Signed

Parents: ff834efde3bd173036ebef30790d4fc1c4ba0183 662e29911c6a07352051ae69f40556c8edeb6046

default/content/presets/openai/Default.json+1 -4 Viewed

     "openrouter_force_instruct": false,
     "openrouter_group_models": false,
     "openrouter_sort_models": "alphabetically",
-    "ai21_model": "j2-ultra",
+    "ai21_model": "jamba-1.5-large",
     "mistralai_model": "mistral-large-latest",
     "custom_model": "",
     "custom_url": "",

     "temperature": 1,
     "frequency_penalty": 0,
     "presence_penalty": 0,
-    "count_penalty": 0,
     "top_p": 1,
     "top_k": 0,
     "top_a": 0,

     "assistant_prefill": "",
     "assistant_impersonation": "",
     "human_sysprompt_message": "Let's get started. Please generate your response based on the information and instructions provided above.",
-    "use_ai21_tokenizer": false,
-    "use_google_tokenizer": false,
     "claude_use_sysprompt": false,
     "use_alt_scale": false,
     "squash_system_messages": false,

default/content/settings.json+2 -4 Viewed

         "temp_openai": 1.0,
         "freq_pen_openai": 0,
         "pres_pen_openai": 0,
-        "count_pen": 0,
         "top_p_openai": 1,
         "top_k_openai": 0,
         "stream_openai": true,

         "wi_format": "{0}",
         "openai_model": "gpt-4-turbo",
         "claude_model": "claude-3-5-sonnet-20240620",
-        "ai21_model": "j2-ultra",
+        "ai21_model": "jamba-1.5-large",
         "windowai_model": "",
         "openrouter_model": "OR_Website",
         "jailbreak_system": true,

@@ -625,7 +624,6 @@
625	"show_external_models": false,	624	"show_external_models": false,
626	"proxy_password": "",	625	"proxy_password": "",
627	"assistant_prefill": "",	626	"assistant_prefill": "",
628	"assistant_impersonation": "",	627	"assistant_impersonation": ""
629	"use_ai21_tokenizer": false
630	}	628	}
631	}	629	}

public/index.html+10 -36 Viewed

                                         </div>
                                     </div>
                                 </div>
-                                <div data-newbie-hidden class="range-block" data-source="openai,openrouter,ai21,custom,cohere,perplexity,groq">
+                                <div data-newbie-hidden class="range-block" data-source="openai,openrouter,custom,cohere,perplexity,groq">
                                     <div class="range-block-title" data-i18n="Frequency Penalty">
                                         Frequency Penalty
                                     </div>

                                         </div>
                                     </div>
                                 </div>
-                                <div data-newbie-hidden class="range-block" data-source="openai,openrouter,ai21,custom,cohere,perplexity,groq">
+                                <div data-newbie-hidden class="range-block" data-source="openai,openrouter,custom,cohere,perplexity,groq">
                                     <div class="range-block-title" data-i18n="Presence Penalty">
                                         Presence Penalty
                                     </div>

                                         </div>
                                     </div>
                                 </div>
-                                <div data-newbie-hidden class="range-block" data-source="ai21">
+                                <div data-newbie-hidden class="range-block" data-source="claude,openrouter,makersuite,cohere,perplexity">
-                                    <div class="range-block-title" data-i18n="Count Penalty">
-                                        Count Penalty
-                                    </div>
-                                    <div class="range-block-range-and-counter">
-                                        <div class="range-block-range">
-                                            <input type="range" id="count_pen" name="volume" min="0" max="1" step="0.01">
-                                        </div>
-                                        <div class="range-block-counter">
-                                            <input type="number" min="0" max="1" step="0.01" data-for="count_pen" id="count_pen_counter">
-                                        </div>
-                                    </div>
-                                </div>
-                                <div data-newbie-hidden class="range-block" data-source="claude,openrouter,ai21,makersuite,cohere,perplexity">
                                     <div class="range-block-title" data-i18n="Top K">
                                         Top K
                                     </div>

                                             </select>
                                         </div>
                                     </div>
-                                    <div class="range-block" data-source="ai21">
-                                        <label for="use_ai21_tokenizer" title="Use AI21 Tokenizer" data-i18n="[title]Use AI21 Tokenizer" class="checkbox_label widthFreeExpand">
-                                            <input id="use_ai21_tokenizer" type="checkbox" /><span data-i18n="Use AI21 Tokenizer">Use AI21 Tokenizer</span>
-                                        </label>
-                                        <div class="toggle-description justifyLeft">
-                                            <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>
-                                        </div>
-                                    </div>
-                                    <div class="range-block" data-source="makersuite">
-                                        <label for="use_google_tokenizer" title="Use Google Tokenizer" data-i18n="[title]Use Google Tokenizer" class="checkbox_label widthFreeExpand">
-                                            <input id="use_google_tokenizer" type="checkbox" /><span data-i18n="Use Google Tokenizer">Use Google Tokenizer</span>
-                                        </label>
-                                        <div class="toggle-description justifyLeft">
-                                            <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>
-                                        </div>
-                                    </div>
                                     <div class="range-block" data-source="makersuite">
                                         <label for="use_makersuite_sysprompt" class="checkbox_label widthFreeExpand">
                                             <input id="use_makersuite_sysprompt" type="checkbox" />

                             <div>
                                 <h4 data-i18n="AI21 Model">AI21 Model</h4>
                                 <select id="model_ai21_select">
-                                    <optgroup label="Latest">
+                                    <optgroup label="Jamba 1.5">
-                                        <option value="j2-ultra">j2-ultra</option>
+                                        <option value="jamba-1.5-mini">jamba-1.5-mini</option>
-                                        <option value="j2-mid">j2-mid</option>
+                                        <option value="jamba-1.5-large">jamba-1.5-large</option>
-                                        <option value="j2-light">j2-light</option>
+                                    </optgroup>
+                                    <optgroup label="Jamba-Instruct (Deprecated)">
+                                        <option value="jamba-instruct-preview">jamba-instruct-preview</option>
                                     </optgroup>
                                 </select>
                             </div>

                                 <option value="3">Llama 1/2</option>
                                 <option value="12">Llama 3</option>
                                 <option value="13">Gemma / Gemini</option>
+                                <option value="14">Jamba</option>
                                 <option value="4">NerdStash (NovelAI Clio)</option>
                                 <option value="5">NerdStash v2 (NovelAI Kayra)</option>
                                 <option value="7">Mistral</option>

public/script.js+1 -1 Viewed

+}
 export function isStreamingEnabled() {
-    const noStreamSources = [chat_completion_sources.SCALE, chat_completion_sources.AI21];
+    const noStreamSources = [chat_completion_sources.SCALE];
     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')))
         || (main_api == 'kobold' && kai_settings.streaming_kobold && kai_flags.can_use_streaming)
         || (main_api == 'novel' && nai_settings.streaming_novel)

public/scripts/openai.js+17 -101 Viewed

 const scale_max = 8191;
 const claude_max = 9000; // We have a proper tokenizer, so theoretically could be larger (up to 9k)
 const claude_100k_max = 99000;
-let ai21_max = 9200; //can easily fit 9k gpt tokens because j2's tokenizer is efficient af
 const unlocked_max = max_200k;
 const oai_max_temp = 2.0;
-const claude_max_temp = 1.0; //same as j2
+const claude_max_temp = 1.0;
-const j2_max_topk = 10.0;
-const j2_max_freq = 5.0;
-const j2_max_pres = 5.0;
 const openrouter_website_model = 'OR_Website';
 const openai_max_stop_strings = 4;

     'custom_include_headers',
 ];
-function getPrefixMap() {
-    return selected_group ? {
-        assistant: '',
-        user: '',
-        system: 'OOC: ',
-    }
-        : {
-            assistant: '{{char}}:',
-            user: '{{user}}:',
-            system: '',
-        };
-}
 const default_settings = {
     preset_settings_openai: 'Default',
     temp_openai: 1.0,
     freq_pen_openai: 0,
     pres_pen_openai: 0,
-    count_pen: 0.0,
     top_p_openai: 1.0,
     top_k_openai: 0,
     min_p_openai: 0,

     openai_model: 'gpt-4-turbo',
     claude_model: 'claude-3-5-sonnet-20240620',
     google_model: 'gemini-1.5-pro',
-    ai21_model: 'j2-ultra',
+    ai21_model: 'jamba-1.5-large',
     mistralai_model: 'mistral-large-latest',
     cohere_model: 'command-r-plus',
     perplexity_model: 'llama-3.1-70b-instruct',

     assistant_prefill: '',
     assistant_impersonation: '',
     human_sysprompt_message: default_claude_human_sysprompt_message,
-    use_ai21_tokenizer: false,
-    use_google_tokenizer: false,
     claude_use_sysprompt: false,
     use_makersuite_sysprompt: true,
     use_alt_scale: false,

     temp_openai: 1.0,
     freq_pen_openai: 0,
     pres_pen_openai: 0,
-    count_pen: 0.0,
     top_p_openai: 1.0,
     top_k_openai: 0,
     min_p_openai: 0,

     openai_model: 'gpt-4-turbo',
     claude_model: 'claude-3-5-sonnet-20240620',
     google_model: 'gemini-1.5-pro',
-    ai21_model: 'j2-ultra',
+    ai21_model: 'jamba-1.5-large',
     mistralai_model: 'mistral-large-latest',
     cohere_model: 'command-r-plus',
     perplexity_model: 'llama-3.1-70b-instruct',

     assistant_prefill: '',
     assistant_impersonation: '',
     human_sysprompt_message: default_claude_human_sysprompt_message,
-    use_ai21_tokenizer: false,
-    use_google_tokenizer: false,
     claude_use_sysprompt: false,
     use_makersuite_sysprompt: true,
     use_alt_scale: false,

     const isClaude = oai_settings.chat_completion_source == chat_completion_sources.CLAUDE;
     const isOpenRouter = oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER;
     const isScale = oai_settings.chat_completion_source == chat_completion_sources.SCALE;
-    const isAI21 = oai_settings.chat_completion_source == chat_completion_sources.AI21;
     const isGoogle = oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE;
     const isOAI = oai_settings.chat_completion_source == chat_completion_sources.OPENAI;
     const isMistral = oai_settings.chat_completion_source == chat_completion_sources.MISTRALAI;

     const isQuiet = type === 'quiet';
     const isImpersonate = type === 'impersonate';
     const isContinue = type === 'continue';
-    const stream = oai_settings.stream_openai && !isQuiet && !isScale && !isAI21 && !(isGoogle && oai_settings.google_model.includes('bison'));
+    const stream = oai_settings.stream_openai && !isQuiet && !isScale && !(isGoogle && oai_settings.google_model.includes('bison'));
     const useLogprobs = !!power_user.request_token_probabilities;
     const canMultiSwipe = oai_settings.n > 1 && !isContinue && !isImpersonate && !isQuiet && (isOAI || isCustom);

         replaceItemizedPromptText(messageId, messages);
+    }
-    if (isAI21) {
-        const joinedMsgs = messages.reduce((acc, obj) => {
-            const prefix = getPrefixMap()[obj.role];
-            return acc + (prefix ? (selected_group ? '\n' : prefix + ' ') : '') + obj.content + '\n';
-        }, '');
-        messages = substituteParams(joinedMsgs) + (isImpersonate ? `${name1}:` : `${name2}:`);
-        replaceItemizedPromptText(messageId, messages);
-    }
     // If we're using the window.ai extension, use that instead
     // Doesn't support logit bias yet
     if (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI) {

         generate_data['use_makersuite_sysprompt'] = oai_settings.use_makersuite_sysprompt;
+    }
-    if (isAI21) {
-        generate_data['top_k'] = Number(oai_settings.top_k_openai);
-        generate_data['count_pen'] = Number(oai_settings.count_pen);
-        generate_data['stop_tokens'] = [name1 + ':', substituteParams(oai_settings.new_chat_prompt), substituteParams(oai_settings.new_group_chat_prompt)];
-    }
     if (isMistral) {
         generate_data['safe_prompt'] = false; // already defaults to false, but just incase they change that in the future.
+    }

     oai_settings.temp_openai = settings.temp_openai ?? default_settings.temp_openai;
     oai_settings.freq_pen_openai = settings.freq_pen_openai ?? default_settings.freq_pen_openai;
     oai_settings.pres_pen_openai = settings.pres_pen_openai ?? default_settings.pres_pen_openai;
-    oai_settings.count_pen = settings.count_pen ?? default_settings.count_pen;
     oai_settings.top_p_openai = settings.top_p_openai ?? default_settings.top_p_openai;
     oai_settings.top_k_openai = settings.top_k_openai ?? default_settings.top_k_openai;
     oai_settings.top_a_openai = settings.top_a_openai ?? default_settings.top_a_openai;

         oai_settings.names_behavior = character_names_behavior.COMPLETION;
+    }
+    if (oai_settings.ai21_model.startsWith('j2-')) {
+        oai_settings.ai21_model = 'jamba-1.5-large';
+    }
     if (settings.wrap_in_quotes !== undefined) oai_settings.wrap_in_quotes = !!settings.wrap_in_quotes;
     if (settings.openai_model !== undefined) oai_settings.openai_model = settings.openai_model;
-    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; }
-    if (settings.use_google_tokenizer !== undefined) oai_settings.use_google_tokenizer = !!settings.use_google_tokenizer;
     if (settings.claude_use_sysprompt !== undefined) oai_settings.claude_use_sysprompt = !!settings.claude_use_sysprompt;
     if (settings.use_makersuite_sysprompt !== undefined) oai_settings.use_makersuite_sysprompt = !!settings.use_makersuite_sysprompt;
     if (settings.use_alt_scale !== undefined) { oai_settings.use_alt_scale = !!settings.use_alt_scale; updateScaleForm(); }

     $('#jailbreak_system').prop('checked', oai_settings.jailbreak_system);
     $('#openai_show_external_models').prop('checked', oai_settings.show_external_models);
     $('#openai_external_category').toggle(oai_settings.show_external_models);
-    $('#use_ai21_tokenizer').prop('checked', oai_settings.use_ai21_tokenizer);
-    $('#use_google_tokenizer').prop('checked', oai_settings.use_google_tokenizer);
     $('#claude_use_sysprompt').prop('checked', oai_settings.claude_use_sysprompt);
     $('#use_makersuite_sysprompt').prop('checked', oai_settings.use_makersuite_sysprompt);
     $('#scale-alt').prop('checked', oai_settings.use_alt_scale);

     $('#pres_pen_openai').val(oai_settings.pres_pen_openai);
     $('#pres_pen_counter_openai').val(Number(oai_settings.pres_pen_openai).toFixed(2));
-    $('#count_pen').val(oai_settings.count_pen);
-    $('#count_pen_counter').val(Number(oai_settings.count_pen).toFixed(2));
     $('#top_p_openai').val(oai_settings.top_p_openai);
     $('#top_p_counter_openai').val(Number(oai_settings.top_p_openai).toFixed(2));

         temperature: settings.temp_openai,
         frequency_penalty: settings.freq_pen_openai,
         presence_penalty: settings.pres_pen_openai,
-        count_penalty: settings.count_pen,
         top_p: settings.top_p_openai,
         top_k: settings.top_k_openai,
         top_a: settings.top_a_openai,

         assistant_prefill: settings.assistant_prefill,
         assistant_impersonation: settings.assistant_impersonation,
         human_sysprompt_message: settings.human_sysprompt_message,
-        use_ai21_tokenizer: settings.use_ai21_tokenizer,
-        use_google_tokenizer: settings.use_google_tokenizer,
         claude_use_sysprompt: settings.claude_use_sysprompt,
         use_makersuite_sysprompt: settings.use_makersuite_sysprompt,
         use_alt_scale: settings.use_alt_scale,

         temperature: ['#temp_openai', 'temp_openai', false],
         frequency_penalty: ['#freq_pen_openai', 'freq_pen_openai', false],
         presence_penalty: ['#pres_pen_openai', 'pres_pen_openai', false],
-        count_penalty: ['#count_pen', 'count_pen', false],
         top_p: ['#top_p_openai', 'top_p_openai', false],
         top_k: ['#top_k_openai', 'top_k_openai', false],
         top_a: ['#top_a_openai', 'top_a_openai', false],

         assistant_prefill: ['#claude_assistant_prefill', 'assistant_prefill', false],
         assistant_impersonation: ['#claude_assistant_impersonation', 'assistant_impersonation', false],
         human_sysprompt_message: ['#claude_human_sysprompt_textarea', 'human_sysprompt_message', false],
-        use_ai21_tokenizer: ['#use_ai21_tokenizer', 'use_ai21_tokenizer', true],
-        use_google_tokenizer: ['#use_google_tokenizer', 'use_google_tokenizer', true],
         claude_use_sysprompt: ['#claude_use_sysprompt', 'claude_use_sysprompt', true],
         use_makersuite_sysprompt: ['#use_makersuite_sysprompt', 'use_makersuite_sysprompt', true],
         use_alt_scale: ['#use_alt_scale', 'use_alt_scale', true],

@@ -4028,6 +3979,11 @@ async function onModelChange() {
4028	}	3979	}
4029		3980
4030	if ($(this).is('#model_ai21_select')) {	3981	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
4031	console.log('AI21 model changed to', value);	3987	console.log('AI21 model changed to', value);
4032	oai_settings.ai21_model = value;	3988	oai_settings.ai21_model = value;
4033	}	3989	}

     if (oai_settings.chat_completion_source == chat_completion_sources.AI21) {
         if (oai_settings.max_context_unlocked) {
             $('#openai_max_context').attr('max', unlocked_max);
-        } else {
+        } else if (oai_settings.ai21_model.includes('jamba-1.5') || oai_settings.ai21_model.includes('jamba-instruct')) {
-            $('#openai_max_context').attr('max', ai21_max);
+            $('#openai_max_context').attr('max', max_256k);
+        }
-        oai_settings.openai_max_context = Math.min(oai_settings.openai_max_context, Number($('#openai_max_context').attr('max')));
+        oai_settings.openai_max_context = Math.min(Number($('#openai_max_context').attr('max')), oai_settings.openai_max_context);
         $('#openai_max_context').val(oai_settings.openai_max_context).trigger('input');
+        $('#temp_openai').attr('max', oai_max_temp).val(oai_settings.temp_openai).trigger('input');
-        oai_settings.temp_openai = Math.min(claude_max_temp, oai_settings.temp_openai);
-        $('#temp_openai').attr('max', claude_max_temp).val(oai_settings.temp_openai).trigger('input');
-        oai_settings.freq_pen_openai = Math.min(j2_max_freq, oai_settings.freq_pen_openai < 0 ? 0 : oai_settings.freq_pen_openai);
-        $('#freq_pen_openai').attr('min', 0).attr('max', j2_max_freq).val(oai_settings.freq_pen_openai).trigger('input');
-        oai_settings.pres_pen_openai = Math.min(j2_max_pres, oai_settings.pres_pen_openai < 0 ? 0 : oai_settings.pres_pen_openai);
-        $('#pres_pen_openai').attr('min', 0).attr('max', j2_max_pres).val(oai_settings.pres_pen_openai).trigger('input');
-        oai_settings.top_k_openai = Math.min(j2_max_topk, oai_settings.top_k_openai);
-        $('#top_k_openai').attr('max', j2_max_topk).val(oai_settings.top_k_openai).trigger('input');
-    } else if (oai_settings.chat_completion_source != chat_completion_sources.AI21) {
-        oai_settings.freq_pen_openai = Math.min(2.0, oai_settings.freq_pen_openai);
-        $('#freq_pen_openai').attr('min', -2.0).attr('max', 2.0).val(oai_settings.freq_pen_openai).trigger('input');
-        oai_settings.pres_pen_openai = Math.min(2.0, oai_settings.pres_pen_openai);
-        $('#pres_pen_openai').attr('min', -2.0).attr('max', 2.0).val(oai_settings.pres_pen_openai).trigger('input');
-        oai_settings.top_k_openai = Math.min(200, oai_settings.top_k_openai);
-        $('#top_k_openai').attr('max', 200).val(oai_settings.top_k_openai).trigger('input');
+    }
     if (oai_settings.chat_completion_source == chat_completion_sources.CUSTOM) {

         saveSettingsDebounced();
     });
-    $('#count_pen').on('input', function () {
-        oai_settings.count_pen = Number($(this).val());
-        $('#count_pen_counter').val(Number($(this).val()).toFixed(2));
-        saveSettingsDebounced();
-    });
     $('#top_p_openai').on('input', function () {
         oai_settings.top_p_openai = Number($(this).val());
         $('#top_p_counter_openai').val(Number($(this).val()).toFixed(2));

         saveSettingsDebounced();
     });
-    $('#use_ai21_tokenizer').on('change', function () {
-        oai_settings.use_ai21_tokenizer = !!$('#use_ai21_tokenizer').prop('checked');
-        oai_settings.use_ai21_tokenizer ? ai21_max = 8191 : ai21_max = 9200;
-        oai_settings.openai_max_context = Math.min(ai21_max, oai_settings.openai_max_context);
-        $('#openai_max_context').attr('max', ai21_max).val(oai_settings.openai_max_context).trigger('input');
-        $('#openai_max_context_counter').attr('max', Number($('#openai_max_context').attr('max')));
-        saveSettingsDebounced();
-    });
-    $('#use_google_tokenizer').on('change', function () {
-        oai_settings.use_google_tokenizer = !!$('#use_google_tokenizer').prop('checked');
-        saveSettingsDebounced();
-    });
     $('#claude_use_sysprompt').on('change', function () {
         oai_settings.claude_use_sysprompt = !!$('#claude_use_sysprompt').prop('checked');
         $('#claude_human_sysprompt_message_block').toggle(oai_settings.claude_use_sysprompt);

public/scripts/tokenizers.js+26 -23 Viewed

     CLAUDE: 11,
     LLAMA3: 12,
     GEMMA: 13,
+    JAMBA: 14,
     BEST_MATCH: 99,
 };

     tokenizers.YI,
     tokenizers.LLAMA3,
     tokenizers.GEMMA,
+    tokenizers.JAMBA,
     // uncomment when NovelAI releases Kayra and Clio weights, lol
     //tokenizers.NERD,
     //tokenizers.NERD2,

         decode: '/api/tokenizers/gemma/decode',
         count: '/api/tokenizers/gemma/encode',
     },
+    [tokenizers.JAMBA]: {
+        encode: '/api/tokenizers/jamba/encode',
+        decode: '/api/tokenizers/jamba/decode',
+        count: '/api/tokenizers/jamba/encode',
+    },
     [tokenizers.API_TEXTGENERATIONWEBUI]: {
         encode: '/api/tokenizers/remote/textgenerationwebui/encode',
         count: '/api/tokenizers/remote/textgenerationwebui/encode',

         tokenizerId: Number(tokenizerOption.value),
         tokenizerKey: Object.entries(tokenizers).find(([_, value]) => value === Number(tokenizerOption.value))[0].toLocaleLowerCase(),
         tokenizerName: tokenizerOption.text,
-    }))
+    }));
+}
 /**

             if (model.includes('gemma')) {
                 return tokenizers.GEMMA;
+            }
+            if (model.includes('yi')) {
+                return tokenizers.YI;
+            }
+            if (model.includes('jamba')) {
+                return tokenizers.JAMBA;
+            }
+        }
         return tokenizers.LLAMA;

     const mistralTokenizer = 'mistral';
     const yiTokenizer = 'yi';
     const gemmaTokenizer = 'gemma';
+    const jambaTokenizer = 'jamba';
     // Assuming no one would use it for different models.. right?
     if (oai_settings.chat_completion_source == chat_completion_sources.SCALE) {

@@ -562,12 +576,19 @@ export function getTokenizerModel() {
562	else if (oai_settings.openrouter_model.includes('GPT-NeoXT')) {	576	else if (oai_settings.openrouter_model.includes('GPT-NeoXT')) {
563	return gpt2Tokenizer;	577	return gpt2Tokenizer;
564	}	578	}
		579	else if (oai_settings.openrouter_model.includes('jamba')) {
		580	return jambaTokenizer;
		581	}
565	}	582	}
566		583
567	if (oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE) {	584	if (oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE) {
568	return gemmaTokenizer;	585	return gemmaTokenizer;
569	}	586	}
570		587
		588	if (oai_settings.chat_completion_source == chat_completion_sources.AI21) {
		589	return jambaTokenizer;
		590	}
		591
571	if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {	592	if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {
572	return claudeTokenizer;	593	return claudeTokenizer;
573	}	594	}

  * @deprecated Use countTokensOpenAIAsync instead.
  */
 export function countTokensOpenAI(messages, full = false) {
-    const shouldTokenizeAI21 = oai_settings.chat_completion_source === chat_completion_sources.AI21 && oai_settings.use_ai21_tokenizer;
+    const tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
-    const shouldTokenizeGoogle = oai_settings.chat_completion_source === chat_completion_sources.MAKERSUITE && oai_settings.use_google_tokenizer;
-    let tokenizerEndpoint = '';
-    if (shouldTokenizeAI21) {
-        tokenizerEndpoint = '/api/tokenizers/ai21/count';
-    } else if (shouldTokenizeGoogle) {
-        tokenizerEndpoint = `/api/tokenizers/google/count?model=${getTokenizerModel()}&reverse_proxy=${oai_settings.reverse_proxy}&proxy_password=${oai_settings.proxy_password}`;
-    } else {
-        tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
-    }
     const cacheObject = getTokenCacheObject();
     if (!Array.isArray(messages)) {

     for (const message of messages) {
         const model = getTokenizerModel();
-        if (model === 'claude' || shouldTokenizeAI21 || shouldTokenizeGoogle) {
+        if (model === 'claude') {
             full = true;
+        }

  * @returns {Promise<number>} Token count.
  */
 export async function countTokensOpenAIAsync(messages, full = false) {
-    const shouldTokenizeAI21 = oai_settings.chat_completion_source === chat_completion_sources.AI21 && oai_settings.use_ai21_tokenizer;
+    const tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
-    const shouldTokenizeGoogle = oai_settings.chat_completion_source === chat_completion_sources.MAKERSUITE && oai_settings.use_google_tokenizer;
-    let tokenizerEndpoint = '';
-    if (shouldTokenizeAI21) {
-        tokenizerEndpoint = '/api/tokenizers/ai21/count';
-    } else if (shouldTokenizeGoogle) {
-        tokenizerEndpoint = `/api/tokenizers/google/count?model=${getTokenizerModel()}`;
-    } else {
-        tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
-    }
     const cacheObject = getTokenCacheObject();
     if (!Array.isArray(messages)) {

     for (const message of messages) {
         const model = getTokenizerModel();
-        if (model === 'claude' || shouldTokenizeAI21 || shouldTokenizeGoogle) {
+        if (model === 'claude') {
             full = true;
+        }

src/endpoints/backends/chat-completions.js+35 -50 Viewed

 const { jsonParser } = require('../../express-common');
 const { CHAT_COMPLETION_SOURCES, GEMINI_SAFETY, BISON_SAFETY, OPENROUTER_HEADERS } = require('../../constants');
 const { forwardFetchResponse, getConfigValue, tryParse, uuidv4, mergeObjectWithYaml, excludeKeysByYaml, color } = require('../../util');
-const { convertClaudeMessages, convertGooglePrompt, convertTextCompletionPrompt, convertCohereMessages, convertMistralMessages, convertCohereTools } = require('../../prompt-converters');
+const { convertClaudeMessages, convertGooglePrompt, convertTextCompletionPrompt, convertCohereMessages, convertMistralMessages, convertCohereTools, convertAI21Messages } = require('../../prompt-converters');
 const { readSecret, SECRET_KEYS } = require('../secrets');
 const { getTokenizerModel, getSentencepiceTokenizer, getTiktokenTokenizer, sentencepieceTokenizers, TEXT_COMPLETION_MODELS } = require('../tokenizers');

 const API_MAKERSUITE = 'https://generativelanguage.googleapis.com';
 const API_01AI = 'https://api.01.ai/v1';
 const API_BLOCKENTROPY = 'https://api.blockentropy.ai/v1';
+const API_AI21 = 'https://api.ai21.com/studio/v1';
 /**
  * Applies a post-processing step to the generated messages.

     request.socket.on('close', function () {
         controller.abort();
     });
+    const convertedPrompt = convertAI21Messages(request.body.messages, request.body.char_name, request.body.user_name);
+    const body = {
+        messages: convertedPrompt,
+        model: request.body.model,
+        max_tokens: request.body.max_tokens,
+        temperature: request.body.temperature,
+        top_p: request.body.top_p,
+        stop: request.body.stop,
+        stream: request.body.stream,
+    };
     const options = {
         method: 'POST',
         headers: {

             'content-type': 'application/json',
             Authorization: `Bearer ${readSecret(request.user.directories, SECRET_KEYS.AI21)}`,
         },
-        body: JSON.stringify({
+        body: JSON.stringify(body),
-            numResults: 1,
-            maxTokens: request.body.max_tokens,
-            minTokens: 0,
-            temperature: request.body.temperature,
-            topP: request.body.top_p,
-            stopSequences: request.body.stop_tokens,
-            topKReturn: request.body.top_k,
-            frequencyPenalty: {
-                scale: request.body.frequency_penalty * 100,
-                applyToWhitespaces: false,
-                applyToPunctuations: false,
-                applyToNumbers: false,
-                applyToStopwords: false,
-                applyToEmojis: false,
-            },
-            presencePenalty: {
-                scale: request.body.presence_penalty,
-                applyToWhitespaces: false,
-                applyToPunctuations: false,
-                applyToNumbers: false,
-                applyToStopwords: false,
-                applyToEmojis: false,
-            },
-            countPenalty: {
-                scale: request.body.count_pen,
-                applyToWhitespaces: false,
-                applyToPunctuations: false,
-                applyToNumbers: false,
-                applyToStopwords: false,
-                applyToEmojis: false,
-            },
-            prompt: request.body.messages,
-        }),
         signal: controller.signal,
     };
-    fetch(`https://api.ai21.com/studio/v1/${request.body.model}/complete`, options)
+    console.log('AI21 request:', body);
-        .then(r => r.json())
-        .then(r => {
-            if (r.completions === undefined) {
-                console.log(r);
-            } else {
-                console.log(r.completions[0].data.text);
-            }
-            const reply = { choices: [{ 'message': { 'content': r.completions?.[0]?.data?.text } }] };
-            return response.send(reply);
-        })
-        .catch(err => {
-            console.error(err);
-            return response.send({ error: true });
-        });
+    try{
+        const generateResponse = await fetch(API_AI21 + '/chat/completions', options);
+        if (request.body.stream) {
+            forwardFetchResponse(generateResponse, response);
+        } else {
+            if (!generateResponse.ok) {
+                console.log(`AI21 API returned error: ${generateResponse.status} ${generateResponse.statusText} ${await generateResponse.text()}`);
+                return response.status(500).send({ error: true });
+            }
+            const generateResponseJson = await generateResponse.json();
+            console.log('AI21 response:', generateResponseJson);
+            return response.send(generateResponseJson);
+        }
+    } catch (error) {
+        console.log('Error communicating with MistralAI API: ', error);
+        if (!response.headersSent) {
+            response.send({ error: true });
+        } else {
+            response.end();
+        }
+    }
+}
 /**

src/endpoints/tokenizers.js+27 -47 Viewed

 const spp_mistral = new SentencePieceTokenizer('src/tokenizers/mistral.model');
 const spp_yi = new SentencePieceTokenizer('src/tokenizers/yi.model');
 const spp_gemma = new SentencePieceTokenizer('src/tokenizers/gemma.model');
+const spp_jamba = new SentencePieceTokenizer('src/tokenizers/jamba.model');
 const claude_tokenizer = new WebTokenizer('src/tokenizers/claude.json');
 const llama3_tokenizer = new WebTokenizer('src/tokenizers/llama3.json');

     'mistral',
     'yi',
     'gemma',
+    'jamba',
 ];
 /**

@@ -186,6 +188,10 @@ function getSentencepiceTokenizer(model) {
186	return spp_gemma;	188	return spp_gemma;
187	}	189	}
188		190
		191	if (model.includes('jamba')) {
		192	return spp_jamba;
		193	}
		194
189	return null;	195	return null;
190	}	196	}
191		197

@@ -322,6 +328,10 @@ function getTokenizerModel(requestModel) {
322	return 'gemma';	328	return 'gemma';
323	}	329	}
324		330
		331	if (requestModel.includes('jamba')) {
		332	return 'jamba';
		333	}
		334
325	// default	335	// default
326	return 'gpt-3.5-turbo';	336	return 'gpt-3.5-turbo';
327	}	337	}

 const router = express.Router();
-router.post('/ai21/count', jsonParser, async function (req, res) {
-    if (!req.body) return res.sendStatus(400);
-    const key = readSecret(req.user.directories, SECRET_KEYS.AI21);
-    const options = {
-        method: 'POST',
-        headers: {
-            accept: 'application/json',
-            'content-type': 'application/json',
-            Authorization: `Bearer ${key}`,
-        },
-        body: JSON.stringify({ text: req.body[0].content }),
-    };
-    try {
-        const response = await fetch('https://api.ai21.com/studio/v1/tokenize', options);
-        const data = await response.json();
-        return res.send({ 'token_count': data?.tokens?.length || 0 });
-    } catch (err) {
-        console.error(err);
-        return res.send({ 'token_count': 0 });
-    }
-});
-router.post('/google/count', jsonParser, async function (req, res) {
-    if (!req.body) return res.sendStatus(400);
-    const options = {
-        method: 'POST',
-        headers: {
-            accept: 'application/json',
-            'content-type': 'application/json',
-        },
-        body: JSON.stringify({ contents: convertGooglePrompt(req.body, String(req.query.model)).contents }),
-    };
-    try {
-        const reverseProxy = req.query.reverse_proxy?.toString() || '';
-        const proxyPassword = req.query.proxy_password?.toString() || '';
-        const apiKey = reverseProxy ? proxyPassword : readSecret(req.user.directories, SECRET_KEYS.MAKERSUITE);
-        const apiUrl = new URL(reverseProxy || API_MAKERSUITE);
-        const response = await fetch(`${apiUrl.origin}/v1beta/models/${req.query.model}:countTokens?key=${apiKey}`, options);
-        const data = await response.json();
-        return res.send({ 'token_count': data?.totalTokens || 0 });
-    } catch (err) {
-        console.error(err);
-        return res.send({ 'token_count': 0 });
-    }
-});
 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('/gemma/encode', jsonParser, createSentencepieceEncodingHandler(spp_gemma));
+router.post('/jamba/encode', jsonParser, createSentencepieceEncodingHandler(spp_jamba));
 router.post('/gpt2/encode', jsonParser, createTiktokenEncodingHandler('gpt2'));
 router.post('/claude/encode', jsonParser, createWebTokenizerEncodingHandler(claude_tokenizer));
 router.post('/llama3/encode', jsonParser, createWebTokenizerEncodingHandler(llama3_tokenizer));

 router.post('/mistral/decode', jsonParser, createSentencepieceDecodingHandler(spp_mistral));
 router.post('/yi/decode', jsonParser, createSentencepieceDecodingHandler(spp_yi));
 router.post('/gemma/decode', jsonParser, createSentencepieceDecodingHandler(spp_gemma));
+router.post('/jamba/decode', jsonParser, createSentencepieceDecodingHandler(spp_jamba));
 router.post('/gpt2/decode', jsonParser, createTiktokenDecodingHandler('gpt2'));
 router.post('/claude/decode', jsonParser, createWebTokenizerDecodingHandler(claude_tokenizer));
 router.post('/llama3/decode', jsonParser, createWebTokenizerDecodingHandler(llama3_tokenizer));

             return handler(req, res);
+        }
+        if (queryModel.includes('jamba')) {
+            const handler = createSentencepieceEncodingHandler(spp_jamba);
+            return handler(req, res);
+        }
         const model = getTokenizerModel(queryModel);
         const handler = createTiktokenEncodingHandler(model);
         return handler(req, res);

             return handler(req, res);
+        }
+        if (queryModel.includes('jamba')) {
+            const handler = createSentencepieceDecodingHandler(spp_jamba);
+            return handler(req, res);
+        }
         const model = getTokenizerModel(queryModel);
         const handler = createTiktokenDecodingHandler(model);
         return handler(req, res);

             return res.send({ 'token_count': num_tokens });
+        }
+        if (model === 'jamba') {
+            num_tokens = await countSentencepieceArrayTokens(spp_jamba, req.body);
+            return res.send({ 'token_count': num_tokens });
+        }
         const tokensPerName = queryModel.includes('gpt-3.5-turbo-0301') ? -1 : 1;
         const tokensPerMessage = queryModel.includes('gpt-3.5-turbo-0301') ? 4 : 3;
         const tokensPadding = 3;

src/prompt-converters.js+73 -0 Viewed

+}
 /**
+ * Convert AI21 prompt. Classic: system message squash, user/assistant message merge.
+ * @param {object[]} messages Array of messages
+ * @param {string} charName Character name
+ * @param {string} userName User name
+ */
+function convertAI21Messages(messages, charName = '', userName = '') {
+    if (!Array.isArray(messages)) {
+        return [];
+    }
+    // Collect all the system messages up until the first instance of a non-system message, and then remove them from the messages array.
+    let i = 0, systemPrompt = '';
+    for (i = 0; i < messages.length; i++) {
+        if (messages[i].role !== 'system') {
+            break;
+        }
+        // Append example names if not already done by the frontend (e.g. for group chats).
+        if (userName && messages[i].name === 'example_user') {
+            if (!messages[i].content.startsWith(`${userName}: `)) {
+                messages[i].content = `${userName}: ${messages[i].content}`;
+            }
+        }
+        if (charName && messages[i].name === 'example_assistant') {
+            if (!messages[i].content.startsWith(`${charName}: `)) {
+                messages[i].content = `${charName}: ${messages[i].content}`;
+            }
+        }
+        systemPrompt += `${messages[i].content}\n\n`;
+    }
+    messages.splice(0, i);
+    // Prevent erroring out if the messages array is empty.
+    if (messages.length === 0) {
+        messages.unshift({
+            role: 'user',
+            content: '[Start a new chat]',
+        });
+    }
+    if (systemPrompt) {
+        messages.unshift({
+            role: 'system',
+            content: systemPrompt.trim(),
+        });
+    }
+    // Doesn't support completion names, so prepend if not already done by the frontend (e.g. for group chats).
+    messages.forEach(msg => {
+        if ('name' in msg) {
+            if (msg.role !== 'system' && !msg.content.startsWith(`${msg.name}: `)) {
+                msg.content = `${msg.name}: ${msg.content}`;
+            }
+            delete msg.name;
+        }
+    });
+    // Since the messaging endpoint only supports alternating turns, we have to merge messages with the same role if they follow each other
+    let mergedMessages = [];
+    messages.forEach((message) => {
+        if (mergedMessages.length > 0 && mergedMessages[mergedMessages.length - 1].role === message.role) {
+            mergedMessages[mergedMessages.length - 1].content += '\n\n' + message.content;
+        } else {
+            mergedMessages.push(message);
+        }
+    });
+    return mergedMessages;
+}
+/**
  * Convert a prompt from the ChatML objects to the format used by MistralAI.
  * @param {object[]} messages Array of messages
  * @param {string} charName Character name

     convertCohereMessages,
     convertMistralMessages,
     convertCohereTools,
+    convertAI21Messages,
 };

src/tokenizers/jamba.model+0 -0 Viewed

Binary file