Enhance cluster naming and K-Means functionality; implement random selection for cluster names, improve centroid initialization, and adjust UMAP parameters for better projection accuracy.

cluster_naming.js

@@ -5,7 +5,20 @@ const tokenizer = await AutoTokenizer.from_pretrained("onnx-community/Qwen3-0.6B
 const model = await AutoModelForCausalLM.from_pretrained("onnx-community/Qwen3-0.6B-ONNX", { device: "webgpu", dtype: "q4f16" });
 
 export async function nameCluster(lines) {
-    const joined = lines.join("\n");
+    // If more than 5 lines, randomly pick 5
+    let selected = lines;
+    if (lines.length > 5) {
+        selected = [];
+        const used = new Set();
+        while (selected.length < 5) {
+            const idx = Math.floor(Math.random() * lines.length);
+            if (!used.has(idx)) {
+                used.add(idx);
+                selected.push(lines[idx]);
+            }
+        }
+    }
+    const joined = selected.join("\n");
     const messages = [
         { role: "system", content: prompt_cluster },
         { role: "user", content: `Input:\n${joined}\nOutput:` }

clustering.js

@@ -1,37 +1,69 @@
 import { UMAP } from "https://cdn.jsdelivr.net/npm/[email protected]/+esm";
 
+function kmeansPlusPlusInit(embeddings, k) {
+    const n = embeddings.length;
+    const dim = embeddings[0].length;
+    const centroids = [embeddings[Math.floor(Math.random() * n)].slice()];
+    const d2 = new Float64Array(n);
+    for (let c = 1; c < k; ++c) {
+        let total = 0;
+        for (let i = 0; i < n; ++i) {
+            let dist = 0;
+            for (let d = 0; d < dim; ++d) {
+                const diff = embeddings[i][d] - centroids[c - 1][d];
+                dist += diff * diff;
+            }
+            if (c === 1 || dist < d2[i]) d2[i] = dist;
+            total += d2[i];
+        }
+        let r = Math.random() * total;
+        let idx = 0;
+        while (r > d2[idx]) r -= d2[idx++];
+        centroids.push(embeddings[idx].slice());
+    }
+    return centroids;
+}
+
 export function kmeans(embeddings, k, maxIter = 100) {
     const n = embeddings.length;
+    if (n === 0) return { labels: [], centroids: [] };
+    k = Math.max(2, Math.min(k, n));
     const dim = embeddings[0].length;
-    let centroids = Array.from({ length: k }, () => embeddings[Math.floor(Math.random() * n)].slice());
-    let labels = new Array(n).fill(0);
+    let centroids = kmeansPlusPlusInit(embeddings, k);
+    const labels = new Uint32Array(n);
 
     const reseed = () => {
-        let bestIdx = 0, bestDist = -1;
+        let farIdx = 0;
+        let farDist = -1;
         for (let i = 0; i < n; ++i) {
-            let minDist = Infinity;
+            let min = Infinity;
             for (let c = 0; c < k; ++c) {
                 let dist = 0;
-                for (let d = 0; d < dim; ++d)
-                    dist += (embeddings[i][d] - centroids[c][d]) ** 2;
-                if (dist < minDist) minDist = dist;
+                for (let d = 0; d < dim; ++d) {
+                    const diff = embeddings[i][d] - centroids[c][d];
+                    dist += diff * diff;
+                }
+                if (dist < min) min = dist;
             }
-            if (minDist > bestDist) {
-                bestDist = minDist;
-                bestIdx = i;
+            if (min > farDist) {
+                farDist = min;
+                farIdx = i;
             }
         }
-        return embeddings[bestIdx].slice();
+        return embeddings[farIdx].slice();
     };
 
     for (let iter = 0; iter < maxIter; ++iter) {
-        let changed = false;
+        let moved = false;
         for (let i = 0; i < n; ++i) {
-            let best = 0, bestDist = Infinity;
+            let best = 0;
+            let bestDist = Infinity;
             for (let c = 0; c < k; ++c) {
                 let dist = 0;
-                for (let d = 0; d < dim; ++d)
-                    dist += (embeddings[i][d] - centroids[c][d]) ** 2;
+                for (let d = 0; d < dim; ++d) {
+                    const diff = embeddings[i][d] - centroids[c][d];
+                    dist += diff * diff;
+                }
                 if (dist < bestDist) {
                     bestDist = dist;
                     best = c;
@@ -39,11 +71,11 @@ export function kmeans(embeddings, k, maxIter = 100) {
             }
             if (labels[i] !== best) {
                 labels[i] = best;
-                changed = true;
+                moved = true;
             }
         }
+        const counts = new Uint32Array(k);
         centroids = Array.from({ length: k }, () => new Array(dim).fill(0));
-        const counts = new Array(k).fill(0);
         for (let i = 0; i < n; ++i) {
             counts[labels[i]]++;
             for (let d = 0; d < dim; ++d)
@@ -53,16 +85,20 @@ export function kmeans(embeddings, k, maxIter = 100) {
             if (counts[c] === 0) {
                 centroids[c] = reseed();
             } else {
-                for (let d = 0; d < dim; ++d)
-                    centroids[c][d] /= counts[c];
+                const inv = 1 / counts[c];
+                for (let d = 0; d < dim; ++d) centroids[c][d] *= inv;
             }
         }
-        if (!changed) break;
+        if (!moved) break;
     }
-    return { labels, centroids };
+    return { labels: Array.from(labels), centroids };
 }
 
 export function runUMAP(embeddings, nNeighbors = 15) {
-    const umap = new UMAP({ nComponents: 2, nNeighbors, minDist: 0.1 });
+    const umap = new UMAP({
+        nComponents: 2,
+        nNeighbors: Math.max(1, Math.min(nNeighbors, embeddings.length - 1)),
+        minDist: 0.1
+    });
     return umap.fit(embeddings);
 }

main.js

@@ -62,7 +62,7 @@ document.getElementById("kmeans-btn").onclick = async () => {
     // UMAP projection
     const { UMAP } = await import('https://cdn.jsdelivr.net/npm/[email protected]/+esm');
     const nNeighbors = Math.max(1, Math.min(lines.length - 1, 15));
-    const umap = new UMAP({ nComponents: 2, nNeighbors, minDist: 0.5, metric: "cosine" });
+    const umap = new UMAP({ nComponents: 2, nNeighbors, minDist: 0.2, metric: "cosine" });
     const proj = umap.fit(embeddings);
     // Group lines by cluster
     const clustered = Array.from({ length: k }, () => []);