add quantized models support (#4)

- add quantized models support (e5f217ff864b7c048b430487ac42ad6e0ca23095)


Co-authored-by: Radamés Ajna <[email protected]>

build/m.d.ts

@@ -8,8 +8,14 @@ export class Decoder {
 * @param {Uint8Array} weights
 * @param {Uint8Array} tokenizer
 * @param {Uint8Array} mel_filters
+* @param {Uint8Array} config
+* @param {boolean} quantized
+* @param {boolean} is_multilingual
+* @param {boolean} timestamps
+* @param {string | undefined} task
+* @param {string | undefined} language
 */
-  constructor(weights: Uint8Array, tokenizer: Uint8Array, mel_filters: Uint8Array);
+  constructor(weights: Uint8Array, tokenizer: Uint8Array, mel_filters: Uint8Array, config: Uint8Array, quantized: boolean, is_multilingual: boolean, timestamps: boolean, task?: string, language?: string);
 /**
 * @param {Uint8Array} wav_input
 * @returns {string}
@@ -22,11 +28,12 @@ export type InitInput = RequestInfo | URL | Response | BufferSource | WebAssembl
 export interface InitOutput {
   readonly memory: WebAssembly.Memory;
   readonly __wbg_decoder_free: (a: number) => void;
-  readonly decoder_new: (a: number, b: number, c: number, d: number, e: number, f: number, g: number) => void;
+  readonly decoder_new: (a: number, b: number, c: number, d: number, e: number, f: number, g: number, h: number, i: number, j: number, k: number, l: number, m: number, n: number, o: number, p: number) => void;
   readonly decoder_decode: (a: number, b: number, c: number, d: number) => void;
   readonly main: (a: number, b: number) => number;
   readonly __wbindgen_add_to_stack_pointer: (a: number) => number;
   readonly __wbindgen_malloc: (a: number, b: number) => number;
+  readonly __wbindgen_realloc: (a: number, b: number, c: number, d: number) => number;
   readonly __wbindgen_free: (a: number, b: number, c: number) => void;
   readonly __wbindgen_start: () => void;
 }

build/m.js

@@ -42,6 +42,63 @@ function passArray8ToWasm0(arg, malloc) {
     return ptr;
 }
 
+const cachedTextEncoder = (typeof TextEncoder !== 'undefined' ? new TextEncoder('utf-8') : { encode: () => { throw Error('TextEncoder not available') } } );
+
+const encodeString = (typeof cachedTextEncoder.encodeInto === 'function'
+    ? function (arg, view) {
+    return cachedTextEncoder.encodeInto(arg, view);
+}
+    : function (arg, view) {
+    const buf = cachedTextEncoder.encode(arg);
+    view.set(buf);
+    return {
+        read: arg.length,
+        written: buf.length
+    };
+});
+
+function passStringToWasm0(arg, malloc, realloc) {
+
+    if (realloc === undefined) {
+        const buf = cachedTextEncoder.encode(arg);
+        const ptr = malloc(buf.length, 1) >>> 0;
+        getUint8Memory0().subarray(ptr, ptr + buf.length).set(buf);
+        WASM_VECTOR_LEN = buf.length;
+        return ptr;
+    }
+
+    let len = arg.length;
+    let ptr = malloc(len, 1) >>> 0;
+
+    const mem = getUint8Memory0();
+
+    let offset = 0;
+
+    for (; offset < len; offset++) {
+        const code = arg.charCodeAt(offset);
+        if (code > 0x7F) break;
+        mem[ptr + offset] = code;
+    }
+
+    if (offset !== len) {
+        if (offset !== 0) {
+            arg = arg.slice(offset);
+        }
+        ptr = realloc(ptr, len, len = offset + arg.length * 3, 1) >>> 0;
+        const view = getUint8Memory0().subarray(ptr + offset, ptr + len);
+        const ret = encodeString(arg, view);
+
+        offset += ret.written;
+    }
+
+    WASM_VECTOR_LEN = offset;
+    return ptr;
+}
+
+function isLikeNone(x) {
+    return x === undefined || x === null;
+}
+
 let cachedInt32Memory0 = null;
 
 function getInt32Memory0() {
@@ -91,8 +148,14 @@ export class Decoder {
     * @param {Uint8Array} weights
     * @param {Uint8Array} tokenizer
     * @param {Uint8Array} mel_filters
+    * @param {Uint8Array} config
+    * @param {boolean} quantized
+    * @param {boolean} is_multilingual
+    * @param {boolean} timestamps
+    * @param {string | undefined} task
+    * @param {string | undefined} language
     */
-    constructor(weights, tokenizer, mel_filters) {
+    constructor(weights, tokenizer, mel_filters, config, quantized, is_multilingual, timestamps, task, language) {
         try {
             const retptr = wasm.__wbindgen_add_to_stack_pointer(-16);
             const ptr0 = passArray8ToWasm0(weights, wasm.__wbindgen_malloc);
@@ -101,7 +164,13 @@ export class Decoder {
             const len1 = WASM_VECTOR_LEN;
             const ptr2 = passArray8ToWasm0(mel_filters, wasm.__wbindgen_malloc);
             const len2 = WASM_VECTOR_LEN;
-            wasm.decoder_new(retptr, ptr0, len0, ptr1, len1, ptr2, len2);
+            const ptr3 = passArray8ToWasm0(config, wasm.__wbindgen_malloc);
+            const len3 = WASM_VECTOR_LEN;
+            var ptr4 = isLikeNone(task) ? 0 : passStringToWasm0(task, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
+            var len4 = WASM_VECTOR_LEN;
+            var ptr5 = isLikeNone(language) ? 0 : passStringToWasm0(language, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
+            var len5 = WASM_VECTOR_LEN;
+            wasm.decoder_new(retptr, ptr0, len0, ptr1, len1, ptr2, len2, ptr3, len3, quantized, is_multilingual, timestamps, ptr4, len4, ptr5, len5);
             var r0 = getInt32Memory0()[retptr / 4 + 0];
             var r1 = getInt32Memory0()[retptr / 4 + 1];
             var r2 = getInt32Memory0()[retptr / 4 + 2];
@@ -183,15 +252,9 @@ function __wbg_get_imports() {
         const ret = new Error(getStringFromWasm0(arg0, arg1));
         return addHeapObject(ret);
     };
-    imports.wbg.__wbg_log_f448472545eafac4 = function(arg0, arg1) {
+    imports.wbg.__wbg_log_0d9af0379e7a06b8 = function(arg0, arg1) {
         console.log(getStringFromWasm0(arg0, arg1));
     };
-    imports.wbg.__wbg_time_fa135a7c2786e907 = function(arg0, arg1) {
-        console.time(getStringFromWasm0(arg0, arg1));
-    };
-    imports.wbg.__wbg_timeEnd_594d82f147c9776f = function(arg0, arg1) {
-        console.timeEnd(getStringFromWasm0(arg0, arg1));
-    };
     imports.wbg.__wbindgen_throw = function(arg0, arg1) {
         throw new Error(getStringFromWasm0(arg0, arg1));
     };

build/m_bg.wasm

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fa61030475868a6533b28628cd20d4d49c7a00f0e2a044c1f141a3d80f3d8a72
-size 3659953
+oid sha256:393c1add1a180c1f0403cf5bb26db587ec59d19bec0c756b613f89b5e12fa512
+size 4070269

build/m_bg.wasm.d.ts

@@ -2,10 +2,11 @@
 /* eslint-disable */
 export const memory: WebAssembly.Memory;
 export function __wbg_decoder_free(a: number): void;
-export function decoder_new(a: number, b: number, c: number, d: number, e: number, f: number, g: number): void;
+export function decoder_new(a: number, b: number, c: number, d: number, e: number, f: number, g: number, h: number, i: number, j: number, k: number, l: number, m: number, n: number, o: number, p: number): void;
 export function decoder_decode(a: number, b: number, c: number, d: number): void;
 export function main(a: number, b: number): number;
 export function __wbindgen_add_to_stack_pointer(a: number): number;
 export function __wbindgen_malloc(a: number, b: number): number;
+export function __wbindgen_realloc(a: number, b: number, c: number, d: number): number;
 export function __wbindgen_free(a: number, b: number, c: number): void;
 export function __wbindgen_start(): void;

index.html

@@ -6,7 +6,7 @@
   <body></body>
 </html>
 
-<!doctype html>
+<!DOCTYPE html>
 <html>
   <head>
     <meta charset="UTF-8" />
@@ -26,9 +26,30 @@
 
       // models base url
       const MODELS = {
+        tiny_multilingual: {
+          base_url: "https://huggingface.co/openai/whisper-tiny/resolve/main/",
+          model: "model.safetensors",
+          tokenizer: "tokenizer.json",
+          config: "config.json",
+        },
         tiny_en: {
           base_url:
-            "https://huggingface.co/openai/whisper-tiny.en/resolve/refs%2Fpr%2F17/",
+            "https://huggingface.co/openai/whisper-tiny.en/resolve/main/",
+          model: "model.safetensors",
+          tokenizer: "tokenizer.json",
+          config: "config.json",
+        },
+        tiny_quantized_multilingual_q80: {
+          base_url: "https://huggingface.co/lmz/candle-whisper/resolve/main/",
+          model: "model-tiny-q80.gguf",
+          tokenizer: "tokenizer-tiny.json",
+          config: "config-tiny.json",
+        },
+        tiny_en_quantized_q80: {
+          base_url: "https://huggingface.co/lmz/candle-whisper/resolve/main/",
+          model: "model-tiny-q80.gguf",
+          tokenizer: "tokenizer-tiny-en.json",
+          config: "config-tiny-en.json",
         },
       };
       const whisperWorker = new Worker("./whisperWorker.js", {
@@ -39,6 +60,7 @@
         weightsURL, // URL to the weights file
         modelID, // model ID
         tokenizerURL, // URL to the tokenizer file
+        configURL, // model config URL
         mel_filtersURL, // URL to the mel filters file
         audioURL, // URL to the audio file
         updateStatus // function to update the status
@@ -48,21 +70,25 @@
             weightsURL,
             modelID,
             tokenizerURL,
+            configURL,
             mel_filtersURL,
             audioURL,
           });
-          whisperWorker.addEventListener("message", (event) => {
+          function messageHandler(event) {
             console.log(event.data);
             if ("status" in event.data) {
               updateStatus(event.data);
             }
             if ("error" in event.data) {
+              whisperWorker.removeEventListener("message", messageHandler);
               reject(new Error(event.data.error));
             }
             if (event.data.status === "complete") {
+              whisperWorker.removeEventListener("message", messageHandler);
               resolve(event.data);
             }
-          });
+          }
+          whisperWorker.addEventListener("message", messageHandler);
         });
       }
 
@@ -125,13 +151,16 @@
           return;
         }
         const modelID = document.querySelector("#model").value;
-        const modelURL = MODELS[modelID].base_url + "model.safetensors";
-        const tokenizerURL = MODELS[modelID].base_url + "tokenizer.json";
+        const model = MODELS[modelID];
+        const modelURL = model.base_url + model.model;
+        const tokenizerURL = model.base_url + model.tokenizer;
+        const configURL = model.base_url + model.config;
 
         classifyAudio(
           modelURL,
           modelID,
           tokenizerURL,
+          configURL,
           "mel_filters.safetensors",
           audioURL,
           updateStatus
@@ -175,8 +204,7 @@
           <a
             href="https://huggingface.co/openai/"
             target="_blank"
-            class="underline hover:text-blue-500 hover:no-underline"
-          >
+            class="underline hover:text-blue-500 hover:no-underline">
             OpenAI Whisper models
           </a>
           and WASM runtime built with
@@ -193,37 +221,38 @@
         <label for="model" class="font-medium">Models Options: </label>
         <select
           id="model"
-          class="border-2 border-gray-500 rounded-md font-light"
-        >
+          class="border-2 border-gray-500 rounded-md font-light">
+          <option value="tiny_multilingual" selected>tiny.en (151 MB)</option>
           <option value="tiny_en" selected>tiny.en (151 MB)</option>
+          <option value="tiny_quantized_multilingual_q80">
+            tiny quantized q80 (41.5 MB)
+          </option>
+          <option value="tiny_en_quantized_q80">
+            tiny.en quantized q80 (41.8 MB)
+          </option>
         </select>
       </div>
       <!-- drag and drop area -->
       <div class="relative">
         <div
           id="drop-area"
-          class="flex flex-col items-center justify-center border-2 border-gray-300 border-dashed rounded-xl relative h-48 w-full overflow-hidden"
-        >
+          class="flex flex-col items-center justify-center border-2 border-gray-300 border-dashed rounded-xl relative h-48 w-full overflow-hidden">
           <div
-            class="flex flex-col items-center justify-center space-y-1 text-center"
-          >
+            class="flex flex-col items-center justify-center space-y-1 text-center">
             <svg
               width="25"
               height="25"
               viewBox="0 0 25 25"
               fill="none"
-              xmlns="http://www.w3.org/2000/svg"
-            >
+              xmlns="http://www.w3.org/2000/svg">
               <path
                 d="M3.5 24.3a3 3 0 0 1-1.9-.8c-.5-.5-.8-1.2-.8-1.9V2.9c0-.7.3-1.3.8-1.9.6-.5 1.2-.7 2-.7h18.6c.7 0 1.3.2 1.9.7.5.6.7 1.2.7 2v18.6c0 .7-.2 1.4-.7 1.9a3 3 0 0 1-2 .8H3.6Zm0-2.7h18.7V2.9H3.5v18.7Zm2.7-2.7h13.3c.3 0 .5 0 .6-.3v-.7l-3.7-5a.6.6 0 0 0-.6-.2c-.2 0-.4 0-.5.3l-3.5 4.6-2.4-3.3a.6.6 0 0 0-.6-.3c-.2 0-.4.1-.5.3l-2.7 3.6c-.1.2-.2.4 0 .7.1.2.3.3.6.3Z"
-                fill="#000"
-              />
+                fill="#000" />
             </svg>
             <div class="flex text-sm text-gray-600">
               <label
                 for="file-upload"
-                class="relative cursor-pointer bg-white rounded-md font-medium text-blue-950 hover:text-blue-700"
-              >
+                class="relative cursor-pointer bg-white rounded-md font-medium text-blue-950 hover:text-blue-700">
                 <span>Drag and drop your audio here</span>
                 <span class="block text-xs">or</span>
                 <span class="block text-xs">Click to upload</span>
@@ -234,15 +263,13 @@
               name="file-upload"
               type="file"
               accept="audio/*"
-              class="sr-only"
-            />
+              class="sr-only" />
           </div>
           <audio
             id="audio"
             hidden
             controls
-            class="w-full p-2 select-none"
-          ></audio>
+            class="w-full p-2 select-none"></audio>
         </div>
       </div>
       <div>
@@ -250,43 +277,37 @@
           <h3 class="font-medium">Examples:</h3>
           <button
             data-value="samples_jfk.wav"
-            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline"
-          >
+            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline">
             <span>jfk.wav</span>
             <span class="text-xs block"> (352 kB)</span>
           </button>
           <button
             data-value="samples_a13.wav"
-            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline"
-          >
+            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline">
             <span>a13.wav</span>
             <span class="text-xs block"> (960 kB)</span>
           </button>
           <button
             data-value="samples_mm0.wav"
-            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline"
-          >
+            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline">
             <span>mm0.wav</span>
             <span class="text-xs block new"> (957 kB)</span>
           </button>
           <button
             data-value="samples_gb0.wav"
-            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline"
-          >
+            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline">
             <span>gb0.wav </span>
             <span class="text-xs block">(4.08 MB)</span>
           </button>
           <button
             data-value="samples_gb1.wav"
-            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline"
-          >
+            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline">
             <span>gb1.wav </span>
             <span class="text-xs block">(6.36 MB)</span>
           </button>
           <button
             data-value="samples_hp0.wav"
-            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline"
-          >
+            class="text-gray-500 border border-gray-500 rounded-md p-2 underline hover:no-underline">
             <span>hp0.wav </span>
             <span class="text-xs block">(8.75 MB)</span>
           </button>
@@ -297,16 +318,14 @@
         <button
           id="detect"
           disabled
-          class="bg-gray-700 hover:bg-gray-800 text-white font-normal py-2 px-4 rounded disabled:bg-gray-300 disabled:cursor-not-allowed"
-        >
+          class="bg-gray-700 hover:bg-gray-800 text-white font-normal py-2 px-4 rounded disabled:bg-gray-300 disabled:cursor-not-allowed">
           Transcribe Audio
         </button>
       </div>
       <div>
         <h3 class="font-medium">Transcription:</h3>
         <div
-          class="min-h-[250px] bg-slate-100 text-gray-500 p-4 rounded-md flex flex-col gap-2"
-        >
+          class="min-h-[250px] bg-slate-100 text-gray-500 p-4 rounded-md flex flex-col gap-2">
           <p hidden id="output-generation" class="grid-rows-2"></p>
           <span id="output-status" class="m-auto font-light"
             >No transcription results yet</span

whisperWorker.js

@@ -17,23 +17,46 @@ class Whisper {
   static instance = {};
   // Retrieve the Whisper model. When called for the first time,
   // this will load the model and save it for future use.
-  static async getInstance(weightsURL, modelID, tokenizerURL, mel_filtersURL) {
+  static async getInstance(params) {
+    const {
+      weightsURL,
+      modelID,
+      tokenizerURL,
+      mel_filtersURL,
+      configURL,
+      quantized,
+      is_multilingual,
+      timestamps,
+      task,
+      language,
+    } = params;
     // load individual modelID only once
     if (!this.instance[modelID]) {
       await init();
 
       self.postMessage({ status: "loading", message: "Loading Model" });
-      const [weightsArrayU8, tokenizerArrayU8, mel_filtersArrayU8] =
-        await Promise.all([
-          fetchArrayBuffer(weightsURL),
-          fetchArrayBuffer(tokenizerURL),
-          fetchArrayBuffer(mel_filtersURL),
-        ]);
+      const [
+        weightsArrayU8,
+        tokenizerArrayU8,
+        mel_filtersArrayU8,
+        configArrayU8,
+      ] = await Promise.all([
+        fetchArrayBuffer(weightsURL),
+        fetchArrayBuffer(tokenizerURL),
+        fetchArrayBuffer(mel_filtersURL),
+        fetchArrayBuffer(configURL),
+      ]);
 
       this.instance[modelID] = new Decoder(
         weightsArrayU8,
         tokenizerArrayU8,
-        mel_filtersArrayU8
+        mel_filtersArrayU8,
+        configArrayU8,
+        quantized,
+        is_multilingual,
+        timestamps,
+        task,
+        language
       );
     } else {
       self.postMessage({ status: "loading", message: "Model Already Loaded" });
@@ -43,17 +66,37 @@ class Whisper {
 }
 
 self.addEventListener("message", async (event) => {
-  const { weightsURL, modelID, tokenizerURL, mel_filtersURL, audioURL } =
-    event.data;
+  const {
+    weightsURL,
+    modelID,
+    tokenizerURL,
+    configURL,
+    mel_filtersURL,
+    audioURL,
+  } = event.data;
   try {
     self.postMessage({ status: "decoding", message: "Starting Decoder" });
-
-    const decoder = await Whisper.getInstance(
+    let quantized = false;
+    if (modelID.includes("quantized")) {
+      quantized = true;
+    }
+    let is_multilingual = false;
+    if (modelID.includes("multilingual")) {
+      is_multilingual = true;
+    }
+    let timestamps = true;
+    const decoder = await Whisper.getInstance({
       weightsURL,
       modelID,
       tokenizerURL,
-      mel_filtersURL
-    );
+      mel_filtersURL,
+      configURL,
+      quantized,
+      is_multilingual,
+      timestamps,
+      task: null,
+      language: null,
+    });
 
     self.postMessage({ status: "decoding", message: "Loading Audio" });
     const audioArrayU8 = await fetchArrayBuffer(audioURL);