radames/Candle-T5-Generation-Wasm

1//load Candle Bert Module wasm module2let init, ModelConditionalGeneration;3 4async function fetchArrayBuffer(url) {5  const cacheName = "t5-candle-cache";6  const cache = await caches.open(cacheName);7  const cachedResponse = await cache.match(url);8  if (cachedResponse) {9    const data = await cachedResponse.arrayBuffer();10    return new Uint8Array(data);11  }12  const res = await fetch(url, { cache: "force-cache" });13  cache.put(url, res.clone());14  return new Uint8Array(await res.arrayBuffer());15}16class ConditionalGeneration {17  static instance = {};18 19  static async getInstance(weightsURL, tokenizerURL, configURL, modelID) {20    if (modelID.includes("quantized")) {21      ({ default: init, ModelConditionalGeneration } = await import(22        "./build/m-quantized.js"23      ));24    } else {25      ({ default: init, ModelConditionalGeneration } = await import(26        "./build/m.js"27      ));28    }29    if (!this.instance[modelID]) {30      await init();31 32      self.postMessage({ status: "loading", message: "Loading Model" });33      const [weightsArrayU8, tokenizerArrayU8, configArrayU8] =34        await Promise.all([35          fetchArrayBuffer(weightsURL),36          fetchArrayBuffer(tokenizerURL),37          fetchArrayBuffer(configURL),38        ]);39 40      this.instance[modelID] = new ModelConditionalGeneration(41        weightsArrayU8,42        tokenizerArrayU8,43        configArrayU844      );45    } else {46      self.postMessage({ status: "ready", message: "Model Already Loaded" });47    }48    return this.instance[modelID];49  }50}51 52self.addEventListener("message", async (event) => {53  const { weightsURL, tokenizerURL, configURL, modelID, prompt, params } =54    event.data;55  let {56    temperature = 0.0,57    seed = 299792458,58    repeat_penalty = 1.1,59    repeat_last_n = 64,60    top_p = 1,61  } = { ...params };62  try {63    self.postMessage({64      status: "ready",65      message: "Starting T5 Conditional Generation",66    });67    const model = await ConditionalGeneration.getInstance(68      weightsURL,69      tokenizerURL,70      configURL,71      modelID72    );73    self.postMessage({74      status: "decoding",75      message: "Decoding Prompt",76    });77    const output = model.decode({78      prompt,79      temperature,80      seed,81      top_p,82      repeat_penalty,83      repeat_last_n,84    });85    self.postMessage({86      status: "complete",87      message: "complete",88      output: output,89    });90  } catch (e) {91    self.postMessage({ error: e });92  }93});94