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

Parents: f81e14a2e8ce335254e83bf52ec82f4436bbe8e1

src/endpoints/tokenizers.js+34 -9 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\
+ * @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 arrayBuffer = fs.readFileSync(pathToModel).buffer;
+            const fileBuffer = await fs.promises.readFile(pathToModel);
-            this.#instance = await Tokenizer.fromJSON(arrayBuffer);
+            this.#instance = await Tokenizer.fromJSON(fileBuffer);
             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', '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', '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', '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', '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', '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',

         /** @type {any} */
         const data = await result.json();
-        const count =  (data?.length ?? data?.count ?? data?.value ?? data?.tokens?.length);
+        const count = (data?.length ?? data?.count ?? data?.value ?? data?.tokens?.length);
         const ids = (data?.tokens ?? data?.ids ?? []);
         return response.send({ count, ids });