Support gzip compression for tokenizer downloads (#4348) * Support gzip compression for tokenizer downloads * Use Uint8Array for buffer conversion * Add gunzip support for manually downloaded tokenizer files * Remove underlying buffer access * Remove additional buffer allocations

4db44af11062724fe3181ba3c9ba5721da24a48a

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

Signed

src/endpoints/tokenizers.js+33 -8 Viewed

 import fs from 'node:fs';
 import path from 'node:path';
 import { Buffer } from 'node:buffer';
 +import zlib from 'node:zlib';
 +import { promisify } from 'node:util';
 import express from 'express';
 import fetch from 'node-fetch';

 const CHARS_PER_TOKEN = 3.35;
 const IS_DOWNLOAD_ALLOWED = getConfigValue('enableDownloadableTokenizers', true, 'boolean');
 +const gunzip = promisify(zlib.gunzip);
 /**
  * Gets a path to the tokenizer model. Downloads the model if it's a URL.
  * @param {string} model Model URL or path
  * @param {string|undefined} fallbackModel Fallback model path\
  * @returns {Promise<string>} Path to the tokenizer model
  */
 async function getPathToTokenizer(model, fallbackModel) {

             fs.mkdirSync(CACHE_PATH, { recursive: true });
+        }
 +        // If an uncompressed version exists, return it
 +        const isCompressed = path.extname(fileName) === '.gz';
 +        const uncompressedName = path.basename(fileName, '.gz');
 +        const uncompressedPath = path.join(CACHE_PATH, uncompressedName);
 +        if (isCompressed && fs.existsSync(uncompressedPath)) {
 +            return uncompressedPath;
 +        }
++
         const cachedFile = path.join(CACHE_PATH, fileName);
         if (fs.existsSync(cachedFile)) {
 +            // If the file was downloaded manually
 +            if (isCompressed) {
 +                const compressedBuffer = await fs.promises.readFile(cachedFile);
 +                const decompressedBuffer = await gunzip(compressedBuffer);
 +                writeFileAtomicSync(uncompressedPath, decompressedBuffer);
 +                await fs.promises.unlink(cachedFile);
 +                return uncompressedPath;
 +            }
             return cachedFile;
+        }

+        }
         const arrayBuffer = await response.arrayBuffer();
 +        if (isCompressed) {
 +            const decompressedBuffer = await gunzip(arrayBuffer);
 +            writeFileAtomicSync(uncompressedPath, decompressedBuffer);
 +            return uncompressedPath;
 +        }
++
         writeFileAtomicSync(cachedFile, Buffer.from(arrayBuffer));
         return cachedFile;
     } catch (error) {

         try {
             const pathToModel = await getPathToTokenizer(this.#model, this.#fallbackModel);
             const arrayBufferfileBuffer = await fs.readFileSyncpromises.readFile(pathToModel).buffer;
             this.#instance = await Tokenizer.fromJSON(arrayBufferfileBuffer);
             console.info('Instantiated the tokenizer for', path.parse(pathToModel).name);
             return this.#instance;
         } catch (error) {

 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');
 const commandRTokenizer = new WebTokenizer('https://github.com/SillyTavern/SillyTavern-Tokenizers/raw/main/command-r.json.gz', 'src/tokenizers/llama3.json');
 const commandATokenizer = new WebTokenizer('https://github.com/SillyTavern/SillyTavern-Tokenizers/raw/main/command-a.json.gz', 'src/tokenizers/llama3.json');
 const qwen2Tokenizer = new WebTokenizer('https://github.com/SillyTavern/SillyTavern-Tokenizers/raw/main/qwen2.json.gz', 'src/tokenizers/llama3.json');
 const nemoTokenizer = new WebTokenizer('https://github.com/SillyTavern/SillyTavern-Tokenizers/raw/main/nemo.json.gz', 'src/tokenizers/llama3.json');
 const deepseekTokenizer = new WebTokenizer('https://github.com/SillyTavern/SillyTavern-Tokenizers/raw/main/deepseek.json.gz', 'src/tokenizers/llama3.json');
 export const sentencepieceTokenizers = [
     'llama',