Fix attention accumulator un-initialized logic, add comment on why on…

…nx webgpu won't work for now

README.md

@@ -38,6 +38,10 @@ make serve
 
 Processing a test sample with 2638 cells took 67 seconds in the browser vs. 34 seconds in python on the same machine.
 
+# Leveraging a GPU
+
+ONNX Web Runtime does have support for GPUs, but unfortunately they don't support all operators yet. Specifically TopK is not [supported](https://github.com/microsoft/onnxruntime/blob/main/js/web/docs/webgpu-operators.md)
+
 # References
 
 [Open Neural Network Exchange (ONNX)](https://onnx.ai/)

worker.js

@@ -1,5 +1,4 @@
 self.importScripts(
-  "https://cdnjs.cloudflare.com/ajax/libs/onnxruntime-web/1.20.1/ort.min.js",
   "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js",
   "https://cdn.jsdelivr.net/npm/[email protected]/dist/iife/h5wasm.min.js",
   "https://cdn.jsdelivr.net/npm/[email protected]/lib/umap-js.min.js"
@@ -22,7 +21,7 @@ self.addEventListener("message", async function (event) {
       genes: self.model.genes,
     });
   } else if (type === "resetAttentionAccumulator") {
-    attentionAccumulator = new Float32Array(self.model.genes.length);
+    attentionAccumulator = null;
     self.postMessage({ type: "attentionAccumulatorReset" });
   }
 });
@@ -251,6 +250,10 @@ async function predict(event) {
 
       encodings.push(output.encoding.cpuData);
 
+      if (!attentionAccumulator) {
+        attentionAccumulator = new Float32Array(genes.length);
+      }
+
       for (let i = 0; i < attentionAccumulator.length; i++) {
         attentionAccumulator[i] += output.attention.cpuData[i];
       }
@@ -300,4 +303,4 @@ async function predict(event) {
     FS.unmount("/work");
     self.postMessage({ type: "error", error: error.message });
   }
-};
+}