Refactor K-Means clustering code for improved readability and performance; streamline cluster name generation and enhance progress bar updates

main.js

@@ -18,9 +18,9 @@ document.getElementById("run").onclick = async () => {
 
     // Extract cluster names from lines starting with ##
     const clusterNames = text.split(/\n/)
-    .map(x => x.trim())
-    .filter(x => x && x.startsWith('##'))
-    .map(x => x.replace(/^##\s*/, ''));
+        .map(x => x.trim())
+        .filter(x => x && x.startsWith('##'))
+        .map(x => x.replace(/^##\s*/, ''));
 
 
     const groupEmbeddings = [];
@@ -54,7 +54,7 @@ document.getElementById("run").onclick = async () => {
         sim.push(row);
     }
     // If clusterNames exist and match group count, use as axis labels
-    let xLabels = clusterNames && clusterNames.length === n ? clusterNames : Array.from({length: n}, (_, i) => `Group ${i+1}`);
+    let xLabels = clusterNames && clusterNames.length === n ? clusterNames : Array.from({ length: n }, (_, i) => `Group ${i + 1}`);
     const data = [{ z: sim, type: "heatmap", colorscale: "Viridis", zmin: 0.7, zmax: 1, x: xLabels, y: xLabels }];
     Plotly.newPlot("plot-heatmap", data, {
         xaxis: { title: "Group", scaleanchor: "y", scaleratio: 1 },
@@ -66,58 +66,98 @@ document.getElementById("run").onclick = async () => {
     });
 };
 
-// --- K-Means Clustering ---
 document.getElementById("kmeans-btn").onclick = async () => {
     const progressBar = document.getElementById("progress-bar");
     const progressBarInner = document.getElementById("progress-bar-inner");
     progressBar.style.display = "block";
-    progressBarInner.style.width = "0%";
+    progressBarInner.style.width = "0%"; // Set to 0% at the start
 
     const text = document.getElementById("input").value;
-    // Remove lines that are cluster headers (## {clusterName})
-    const lines = text.split(/\n/)
-        .map(x => x.trim())
-        .filter(x => x && !x.startsWith('##'));
+    const lines = text.split(/\n/).map(x => x.trim()).filter(x => x && !x.startsWith("##"));
     const prompts = lines.map(s => `Instruct: ${task}\nQuery:${s}`);
     const out = await embed(prompts, { pooling: "mean", normalize: true });
-    const embeddings = typeof out.tolist === 'function' ? out.tolist() : out.data;
+    const embeddings = typeof out.tolist === "function" ? out.tolist() : out.data;
+
+    const n = embeddings.length;
+    if (n < 2) return;
+
+    const requestedK = parseInt(document.getElementById("kmeans-k").value) || 3;
+    const k = Math.max(2, Math.min(requestedK, n));
+    const dim = embeddings[0].length;
 
-    // K-Means implementation
-    const k = Math.max(2, Math.min(20, parseInt(document.getElementById("kmeans-k").value) || 3));
-    const n = embeddings.length, dim = embeddings[0].length;
     let centroids = Array.from({ length: k }, () => embeddings[Math.floor(Math.random() * n)].slice());
     let labels = new Array(n).fill(0);
+
+    const reseed = () => {
+        let bestIdx = 0, bestDist = -1;
+        for (let i = 0; i < n; ++i) {
+            let minDist = 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;
+            }
+            if (minDist > bestDist) {
+                bestDist = minDist;
+                bestIdx = i;
+            }
+        }
+        return embeddings[bestIdx].slice();
+    };
+
     for (let iter = 0; iter < 100; ++iter) {
+        let changed = false;
         for (let i = 0; i < n; ++i) {
             let best = 0, 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;
-                if (dist < bestDist) { bestDist = dist; best = c; }
+                for (let d = 0; d < dim; ++d)
+                    dist += (embeddings[i][d] - centroids[c][d]) ** 2;
+                if (dist < bestDist) {
+                    bestDist = dist;
+                    best = c;
+                }
+            }
+            if (labels[i] !== best) {
+                labels[i] = best;
+                changed = true;
             }
-            labels[i] = best;
         }
+
         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) centroids[labels[i]][d] += embeddings[i][d];
+            for (let d = 0; d < dim; ++d)
+                centroids[labels[i]][d] += embeddings[i][d];
+        }
+        for (let c = 0; c < k; ++c) {
+            if (counts[c] === 0) {
+                centroids[c] = reseed();
+            } else {
+                for (let d = 0; d < dim; ++d)
+                    centroids[c][d] /= counts[c];
+            }
         }
-        for (let c = 0; c < k; ++c) if (counts[c]) for (let d = 0; d < dim; ++d) centroids[c][d] /= counts[c];
+        if (!changed) break;
     }
-    // UMAP for 2D projection
+
     const nNeighbors = Math.max(1, Math.min(lines.length - 1, 15));
     const umap = new UMAP({ nComponents: 2, nNeighbors, minDist: 0.1 });
     const proj = umap.fit(embeddings);
-    // Group lines by cluster
-    const clustered = Array.from({ length: k }, (_, c) => []);
-    for (let i = 0; i < n; ++i) clustered[labels[i]].push(lines[i]);
-    // Plot initial scatter with placeholder names before clusterNames are ready
+
+    const clustered = Array.from({ length: k }, () => []);
+    for (let i = 0; i < n; ++i)
+        clustered[labels[i]].push(lines[i]);
+
     const colors = ["red", "blue", "green", "orange", "purple", "cyan", "magenta", "yellow", "brown", "black", "lime", "navy", "teal", "olive", "maroon", "pink", "gray", "gold", "aqua", "indigo"];
     const placeholderNames = Array.from({ length: k }, (_, c) => `Cluster ${c + 1}`);
-    let traces = Array.from({ length: k }, (_, c) => ({
-        x: [], y: [], text: [], mode: "markers", type: "scatter", name: placeholderNames[c],
-        marker: { color: colors[c % colors.length], size: 12, line: { width: 1, color: '#333' } }
+    const traces = Array.from({ length: k }, (_, c) => ({
+        x: [], y: [], text: [],
+        mode: "markers", type: "scatter",
+        name: placeholderNames[c],
+        marker: { color: colors[c % colors.length], size: 12, line: { width: 1, color: "#333" } }
     }));
     for (let i = 0; i < n; ++i) {
         traces[labels[i]].x.push(proj[i][0]);
@@ -133,15 +173,14 @@ document.getElementById("kmeans-btn").onclick = async () => {
         title: `K-Means Clustering (k=${k})`,
         legend: { x: 1.05, y: 0.5, orientation: "v", xanchor: "left", yanchor: "middle" }
     });
-    // Generate cluster names using text generation pipeline (async with progress)
+
     const clusterNames = [];
     for (let c = 0; c < k; ++c) {
-        progressBarInner.style.width = `${Math.round(((c) / k) * 100)}%`;
+        progressBarInner.style.width = `${Math.round(((c + 1) / k) * 100)}%`;
+
         const joined = clustered[c].join("\n");
         const messages = [
-            {
-                role: "system", content: prompt_cluster
-            },
+            { role: "system", content: prompt_cluster },
             { role: "user", content: `Input:\n${joined}\nOutput:` }
         ];
 
@@ -150,64 +189,29 @@ document.getElementById("kmeans-btn").onclick = async () => {
             return_dict: true,
             enable_thinking: false,
         });
-        const [START_THINKING_TOKEN_ID, END_THINKING_TOKEN_ID] = tokenizer.encode("<think></think>", { add_special_tokens: false });
-        let state = "answering";
-        let startTime;
-        let numTokens = 0;
-        let tps;
-        const token_callback_function = (tokens) => {
-            startTime ??= performance.now();
-            if (numTokens++ > 0) {
-                tps = (numTokens / (performance.now() - startTime)) * 1000;
-            }
-            switch (Number(tokens[0])) {
-                case START_THINKING_TOKEN_ID:
-                    state = "thinking";
-                    break;
-                case END_THINKING_TOKEN_ID:
-                    state = "answering";
-                    break;
-            }
-            console.log(state, tokens, tokenizer.decode(tokens));
-        };
-        const callback_function = (output) => {
-            // You can update UI here if desired
-            console.log({ output, tps, numTokens, state });
-        };
-        const streamer = new TextStreamer(tokenizer, {
-            skip_prompt: true,
-            skip_special_tokens: true,
-            callback_function,
-            token_callback_function,
-        });
+
         const outputTokens = await model.generate({
             ...inputs,
             max_new_tokens: 1024,
             do_sample: true,
-            temperature: 0.6,
-            // streamer,
+            temperature: 0.6
         });
-        let rawName = tokenizer
-            .decode(outputTokens[0], { skip_special_tokens: false })
-            .trim();
 
-        const THINK_TAG = '</think>';
-        const END_TAG = '<|im_end|>';
+        let rawName = tokenizer.decode(outputTokens[0], { skip_special_tokens: false }).trim();
+
+        const THINK_TAG = "</think>";
+        const END_TAG = "<|im_end|>";
 
         if (rawName.includes(THINK_TAG)) {
-            // take everything after the last </think>
             rawName = rawName.substring(rawName.lastIndexOf(THINK_TAG) + THINK_TAG.length).trim();
         }
         if (rawName.includes(END_TAG)) {
-            // take everything before the first <|im_end|>
             rawName = rawName.substring(0, rawName.indexOf(END_TAG)).trim();
         }
-        // use a default if name is empty
+
         clusterNames.push(rawName || `Cluster ${c + 1}`);
-        // After all names are generated, update the trace names and render once
-        for (let c = 0; c < k; ++c) {
-            traces[c].name = clusterNames[c];
-        }
+        traces[c].name = clusterNames[c];
+
         Plotly.react("plot-scatter", traces, {
             xaxis: { title: "UMAP-1", scaleanchor: "y", scaleratio: 1 },
             yaxis: { title: "UMAP-2", scaleanchor: "x", scaleratio: 1 },
@@ -217,9 +221,13 @@ document.getElementById("kmeans-btn").onclick = async () => {
             title: `K-Means Clustering (k=${k})`,
             legend: { x: 1.05, y: 0.5, orientation: "v", xanchor: "left", yanchor: "middle" }
         });
-        // Update textarea: group by cluster, separated by triple newlines, with cluster name as header
-        document.getElementById("input").value = clustered.map((g, i) => `## ${clusterNames[i]}\n` + g.join("\n")).join("\n\n\n");
-        // Re-run heatmap after updating textarea
+
+        document.getElementById("input").value = clustered.map((g, i) =>
+            `## ${clusterNames[i]}\n${g.join("\n")}`
+        ).join("\n\n\n");
+
         document.getElementById("run").onclick();
     }
+
+    progressBarInner.style.width = "100%"; // Set to 100% after all clusters are named
 };