Add tokenizer for Gemma/Gemini

e707def7dddb00ed9d72b04df9d7fb07136901d0

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

Parents: ce8b0aae960ee00f3b50a85e3b2628e05b2f8148
5 files changed, +52 -3Showing whitespace changes
public/index.html+1 -0
@@ -3449,6 +3449,7 @@
34493449 <!-- Option #2 was a legacy GPT-2/3 tokenizer -->
34503450 <option value="3">Llama 1/2</option>
34513451 <option value="12">Llama 3</option>
3452+ <option value="13">Gemma / Gemini</option>
34523453 <option value="4">NerdStash (NovelAI Clio)</option>
34533454 <option value="5">NerdStash v2 (NovelAI Kayra)</option>
34543455 <option value="7">Mistral</option>
public/scripts/textgen-models.js+4 -0
@@ -599,6 +599,10 @@ export function getCurrentOpenRouterModelTokenizer() {
599599 return tokenizers.YI;
600600 case 'Mistral':
601601 return tokenizers.MISTRAL;
602+ case 'Gemini':
603+ return tokenizers.GEMMA;
604+ case 'Claude':
605+ return tokenizers.CLAUDE;
602606 default:
603607 return tokenizers.OPENAI;
604608 }
public/scripts/tokenizers.js+22 -1
@@ -26,6 +26,7 @@ export const tokenizers = {
2626 API_KOBOLD: 10,
2727 CLAUDE: 11,
2828 LLAMA3: 12,
29+ GEMMA: 13,
2930 BEST_MATCH: 99,
3031};
3132
@@ -34,6 +35,7 @@ export const SENTENCEPIECE_TOKENIZERS = [
3435 tokenizers.MISTRAL,
3536 tokenizers.YI,
3637 tokenizers.LLAMA3,
38+ tokenizers.GEMMA,
3739 // uncomment when NovelAI releases Kayra and Clio weights, lol
3840 //tokenizers.NERD,
3941 //tokenizers.NERD2,
@@ -91,6 +93,11 @@ const TOKENIZER_URLS = {
9193 decode: '/api/tokenizers/llama3/decode',
9294 count: '/api/tokenizers/llama3/encode',
9395 },
96+ [tokenizers.GEMMA]: {
97+ encode: '/api/tokenizers/gemma/encode',
98+ decode: '/api/tokenizers/gemma/decode',
99+ count: '/api/tokenizers/gemma/encode',
100+ },
94101 [tokenizers.API_TEXTGENERATIONWEBUI]: {
95102 encode: '/api/tokenizers/remote/textgenerationwebui/encode',
96103 count: '/api/tokenizers/remote/textgenerationwebui/encode',
@@ -232,6 +239,9 @@ export function getTokenizerBestMatch(forApi) {
232239 if (model.includes('mistral') || model.includes('mixtral')) {
233240 return tokenizers.MISTRAL;
234241 }
242+ if (model.includes('gemma')) {
243+ return tokenizers.GEMMA;
244+ }
235245 }
236246
237247 return tokenizers.LLAMA;
@@ -441,12 +451,14 @@ export function getTokenizerModel() {
441451 const turbo0301Tokenizer = 'gpt-3.5-turbo-0301';
442452 const turboTokenizer = 'gpt-3.5-turbo';
443453 const gpt4Tokenizer = 'gpt-4';
454+ const gpt4oTokenizer = 'gpt-4o';
444455 const gpt2Tokenizer = 'gpt2';
445456 const claudeTokenizer = 'claude';
446457 const llamaTokenizer = 'llama';
447458 const llama3Tokenizer = 'llama3';
448459 const mistralTokenizer = 'mistral';
449460 const yiTokenizer = 'yi';
461+ const gemmaTokenizer = 'gemma';
450462
451463 // Assuming no one would use it for different models.. right?
452464 if (oai_settings.chat_completion_source == chat_completion_sources.SCALE) {
@@ -491,6 +503,12 @@ export function getTokenizerModel() {
491503 else if (model?.architecture?.tokenizer === 'Yi') {
492504 return yiTokenizer;
493505 }
506+ else if (model?.architecture?.tokenizer === 'Gemini') {
507+ return gemmaTokenizer;
508+ }
509+ else if (oai_settings.openrouter_model.includes('gpt-4o')) {
510+ return gpt4oTokenizer;
511+ }
494512 else if (oai_settings.openrouter_model.includes('gpt-4')) {
495513 return gpt4Tokenizer;
496514 }
@@ -509,7 +527,7 @@ export function getTokenizerModel() {
509527 }
510528
511529 if (oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE) {
512530 return oai_settings.google_modelgemmaTokenizer;
513531 }
514532
515533 if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {
@@ -543,6 +561,9 @@ export function getTokenizerModel() {
543561 if (oai_settings.groq_model.includes('mistral') || oai_settings.groq_model.includes('mixtral')) {
544562 return mistralTokenizer;
545563 }
564+ if (oai_settings.groq_model.includes('gemma')) {
565+ return gemmaTokenizer;
566+ }
546567 }
547568
548569 if (oai_settings.chat_completion_source === chat_completion_sources.ZEROONEAI) {
src/endpoints/tokenizers.js+25 -2
@@ -143,6 +143,7 @@ const spp_nerd = new SentencePieceTokenizer('src/tokenizers/nerdstash.model');
143143const spp_nerd_v2 = new SentencePieceTokenizer('src/tokenizers/nerdstash_v2.model');
144144const spp_mistral = new SentencePieceTokenizer('src/tokenizers/mistral.model');
145145const spp_yi = new SentencePieceTokenizer('src/tokenizers/yi.model');
146+const spp_gemma = new SentencePieceTokenizer('src/tokenizers/gemma.model');
146147const claude_tokenizer = new WebTokenizer('src/tokenizers/claude.json');
147148const llama3_tokenizer = new WebTokenizer('src/tokenizers/llama3.json');
148149
@@ -152,6 +153,7 @@ const sentencepieceTokenizers = [
152153 'nerdstash_v2',
153154 'mistral',
154155 'yi',
156+ 'gemma',
155157];
156158
157159/**
@@ -180,6 +182,10 @@ function getSentencepiceTokenizer(model) {
180182 return spp_yi;
181183 }
182184
185+ if (model.includes('gemma')) {
186+ return spp_gemma;
187+ }
188+
183189 return null;
184190}
185191
@@ -312,8 +318,8 @@ function getTokenizerModel(requestModel) {
312318 return 'yi';
313319 }
314320
315321 if (requestModel.includes('gemma') || requestModel.includes('gemini')) {
316322 return 'gpt-4ogemma';
317323 }
318324
319325 // default
@@ -583,6 +589,7 @@ router.post('/nerdstash/encode', jsonParser, createSentencepieceEncodingHandler(
583589router.post('/nerdstash_v2/encode', jsonParser, createSentencepieceEncodingHandler(spp_nerd_v2));
584590router.post('/mistral/encode', jsonParser, createSentencepieceEncodingHandler(spp_mistral));
585591router.post('/yi/encode', jsonParser, createSentencepieceEncodingHandler(spp_yi));
592+router.post('/gemma/encode', jsonParser, createSentencepieceEncodingHandler(spp_gemma));
586593router.post('/gpt2/encode', jsonParser, createTiktokenEncodingHandler('gpt2'));
587594router.post('/claude/encode', jsonParser, createWebTokenizerEncodingHandler(claude_tokenizer));
588595router.post('/llama3/encode', jsonParser, createWebTokenizerEncodingHandler(llama3_tokenizer));
@@ -591,6 +598,7 @@ router.post('/nerdstash/decode', jsonParser, createSentencepieceDecodingHandler(
591598router.post('/nerdstash_v2/decode', jsonParser, createSentencepieceDecodingHandler(spp_nerd_v2));
592599router.post('/mistral/decode', jsonParser, createSentencepieceDecodingHandler(spp_mistral));
593600router.post('/yi/decode', jsonParser, createSentencepieceDecodingHandler(spp_yi));
601+router.post('/gemma/decode', jsonParser, createSentencepieceDecodingHandler(spp_gemma));
594602router.post('/gpt2/decode', jsonParser, createTiktokenDecodingHandler('gpt2'));
595603router.post('/claude/decode', jsonParser, createWebTokenizerDecodingHandler(claude_tokenizer));
596604router.post('/llama3/decode', jsonParser, createWebTokenizerDecodingHandler(llama3_tokenizer));
@@ -624,6 +632,11 @@ router.post('/openai/encode', jsonParser, async function (req, res) {
624632 return handler(req, res);
625633 }
626634
635+ if (queryModel.includes('gemma') || queryModel.includes('gemini')) {
636+ const handler = createSentencepieceEncodingHandler(spp_gemma);
637+ return handler(req, res);
638+ }
639+
627640 const model = getTokenizerModel(queryModel);
628641 const handler = createTiktokenEncodingHandler(model);
629642 return handler(req, res);
@@ -662,6 +675,11 @@ router.post('/openai/decode', jsonParser, async function (req, res) {
662675 return handler(req, res);
663676 }
664677
678+ if (queryModel.includes('gemma') || queryModel.includes('gemini')) {
679+ const handler = createSentencepieceDecodingHandler(spp_gemma);
680+ return handler(req, res);
681+ }
682+
665683 const model = getTokenizerModel(queryModel);
666684 const handler = createTiktokenDecodingHandler(model);
667685 return handler(req, res);
@@ -708,6 +726,11 @@ router.post('/openai/count', jsonParser, async function (req, res) {
708726 return res.send({ 'token_count': num_tokens });
709727 }
710728
729+ if (model === 'gemma' || model === 'gemini') {
730+ num_tokens = await countSentencepieceArrayTokens(spp_gemma, req.body);
731+ return res.send({ 'token_count': num_tokens });
732+ }
733+
711734 const tokensPerName = queryModel.includes('gpt-3.5-turbo-0301') ? -1 : 1;
712735 const tokensPerMessage = queryModel.includes('gpt-3.5-turbo-0301') ? 4 : 3;
713736 const tokensPadding = 3;
src/tokenizers/gemma.model+0 -0

Binary file