Update index.html

index.html

@@ -300,17 +300,37 @@
       const out = await model({ input_ids: enc.input_ids, attention_mask: enc.attention_mask });
       const dt = (performance.now() - t0) | 0;
 
-      const last = out.logits[out.logits.length - 1];
-      const m = Math.max(...last);
-      const exps = last.map(v => Math.exp(v - m));
-      const Z = exps.reduce((a,b)=>a+b,0);
-      const probs = exps.map(v => v/Z);
-
-      const idx = probs.map((p,i)=>[p,i]).sort((a,b)=>b[0]-a[0]).slice(0, topK);
+      // --- logits come as a Tensor (data: Float32Array, dims: [1, seqLen, vocab]) ---
+      const logitsT = out.logits;
+      const dims    = logitsT.dims;           // e.g., [1, seqLen, vocabSize]
+      const data    = logitsT.data;           // Float32Array
+      
+      const vocabSize = dims[dims.length - 1];
+      const seqLen    = dims[dims.length - 2];
+      
+      // Take the last time step (length = vocabSize) from the flat buffer
+      const start = (seqLen - 1) * vocabSize;
+      const last  = data.subarray(start, start + vocabSize); // typed view (no copy)
+      
+      // Softmax for probabilities
+      let m = -Infinity;
+      for (let i = 0; i < last.length; i++) if (last[i] > m) m = last[i];
+      
+      const exps = new Float32Array(last.length);
+      let Z = 0;
+      for (let i = 0; i < last.length; i++) { const e = Math.exp(last[i] - m); exps[i] = e; Z += e; }
+      
+      // Top-K
+      const K   = Math.min(parseInt(topkSel.value, 10) || 10, last.length);
+      const idx = Array.from({ length: last.length }, (_, i) => [exps[i] / Z, i])
+        .sort((a, b) => b[0] - a[0])
+        .slice(0, K);
+      
+      // Build rows
       const rows = [];
-      for (const [p, i] of idx){
-        const tok = await tokenizer.decode([i], { skip_special_tokens:false });
-        rows.push({ token: tok, p, id:i });
+      for (const [p, i] of idx) {
+        const tok = await tokenizer.decode([i], { skip_special_tokens: false });
+        rows.push({ token: tok, p, id: i });
       }
       return { rows, dt };
     }