Refactor K-Means clustering implementation; modularize embedding and clustering logic, enhance heatmap and scatter plot functions, and improve cluster naming process.

cluster_naming.js

@@ -0,0 +1,34 @@
+import { AutoTokenizer, AutoModelForCausalLM } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.0';
+import { prompt_cluster } from "./prompt_cluster.js";
+
+const tokenizer = await AutoTokenizer.from_pretrained("onnx-community/Qwen3-0.6B-ONNX");
+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");
+    const messages = [
+        { role: "system", content: prompt_cluster },
+        { role: "user", content: `Input:\n${joined}\nOutput:` }
+    ];
+    const inputs = tokenizer.apply_chat_template(messages, {
+        add_generation_prompt: true,
+        return_dict: true,
+        enable_thinking: false,
+    });
+    const outputTokens = await model.generate({
+        ...inputs,
+        max_new_tokens: 1024,
+        do_sample: true,
+        temperature: 0.6
+    });
+    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)) {
+        rawName = rawName.substring(rawName.lastIndexOf(THINK_TAG) + THINK_TAG.length).trim();
+    }
+    if (rawName.includes(END_TAG)) {
+        rawName = rawName.substring(0, rawName.indexOf(END_TAG)).trim();
+    }
+    return rawName;
+}

clustering.js

@@ -0,0 +1,68 @@
+import { UMAP } from "https://cdn.jsdelivr.net/npm/umap-js@1.4.0/+esm";
+
+export function kmeans(embeddings, k, maxIter = 100) {
+    const n = embeddings.length;
+    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);
+
+    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 < maxIter; ++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;
+                }
+            }
+            if (labels[i] !== best) {
+                labels[i] = best;
+                changed = true;
+            }
+        }
+        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 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];
+            }
+        }
+        if (!changed) break;
+    }
+    return { labels, centroids };
+}
+
+export function runUMAP(embeddings, nNeighbors = 15) {
+    const umap = new UMAP({ nComponents: 2, nNeighbors, minDist: 0.1 });
+    return umap.fit(embeddings);
+}

embedding.js

@@ -0,0 +1,32 @@
+import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.0';
+
+const embed = await pipeline(
+    "feature-extraction",
+    "onnx-community/Qwen3-Embedding-0.6B-ONNX",
+    { device: "webgpu", dtype: "q4f16" },
+);
+
+export async function getGroupEmbeddings(groups, task) {
+    const groupEmbeddings = [];
+    for (const g of groups) {
+        // Remove lines starting with ##
+        const lines = g.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 dim = embeddings[0].length;
+        const avg = new Float32Array(dim);
+        for (const e of embeddings) { for (let i = 0; i < dim; i++) avg[i] += e[i]; }
+        for (let i = 0; i < dim; i++) avg[i] /= embeddings.length;
+        groupEmbeddings.push(avg);
+    }
+    return groupEmbeddings;
+}
+
+export async function getLineEmbeddings(lines, task) {
+    const prompts = lines.map(s => `Instruct: ${task}\nQuery:${s}`);
+    const out = await embed(prompts, { pooling: "mean", normalize: true });
+    return typeof out.tolist === 'function' ? out.tolist() : out.data;
+}

main.js

@@ -1,44 +1,27 @@
-import { pipeline, TextStreamer, AutoTokenizer, AutoModelForCausalLM } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.0';
-import { UMAP } from "https://cdn.jsdelivr.net/npm/umap-js@1.4.0/+esm";
-import { prompt_cluster } from "./prompt_cluster.js"
-const embed = await pipeline(
-    "feature-extraction",
-    "onnx-community/Qwen3-Embedding-0.6B-ONNX",
-    { device: "webgpu", dtype: "q4f16" },
-);
-const tokenizer = await AutoTokenizer.from_pretrained("onnx-community/Qwen3-0.6B-ONNX");
-const model = await AutoModelForCausalLM.from_pretrained("onnx-community/Qwen3-0.6B-ONNX", { device: "webgpu", dtype: "q4f16" });
+import { getGroupEmbeddings, getLineEmbeddings } from './embedding.js';
+import { kmeans } from './clustering.js';
+import { plotHeatmap, plotScatter, updateScatter } from './plotting.js';
+import { nameCluster } from './cluster_naming.js';
+import { prompt_cluster } from './prompt_cluster.js';
 
 const task = "Given a textual input sentence, retrieve relevant categories that best describe it.";
 
+// Heatmap event
+// Handles group similarity heatmap
+// Uses group-level embeddings
+
 document.getElementById("run").onclick = async () => {
     const text = document.getElementById("input").value;
-
     const groups = text.split(/\n{3,}/);
-
     // 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*/, ''));
-
-
-    const groupEmbeddings = [];
-    for (const g of groups) {
-        // Remove lines starting with ##
-        const lines = g.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 dim = embeddings[0].length;
-        const avg = new Float32Array(dim);
-        for (const e of embeddings) { for (let i = 0; i < dim; i++) avg[i] += e[i]; }
-        for (let i = 0; i < dim; i++) avg[i] /= embeddings.length;
-        groupEmbeddings.push(avg);
-    }
+    // Get group embeddings (removes ## lines internally)
+    const groupEmbeddings = await getGroupEmbeddings(groups, task);
     const n = groupEmbeddings.length;
+    // Cosine similarity matrix
     const sim = [];
     for (let i = 0; i < n; i++) {
         const row = [];
@@ -53,104 +36,39 @@ document.getElementById("run").onclick = async () => {
         }
         sim.push(row);
     }
-    // If clusterNames exist and match group count, use as axis labels
+    // Use cluster names as axis labels if available
     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 },
-        yaxis: { title: "Group", scaleanchor: "x", scaleratio: 1 },
-        width: 500,
-        height: 500,
-        margin: { t: 40, l: 200, r: 10, b: 200 },
-        title: "Group Similarity Heatmap"
-    });
+    plotHeatmap(sim, xLabels, xLabels);
 };
 
+// K-Means + UMAP + Cluster Naming event
+
 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%"; // Set to 0% at the start
+    progressBarInner.style.width = "0%";
 
     const text = document.getElementById("input").value;
+    // Remove ## lines for embedding
     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 = await getLineEmbeddings(lines, task);
     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;
-
-    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;
-                }
-            }
-            if (labels[i] !== best) {
-                labels[i] = best;
-                changed = true;
-            }
-        }
-
-        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 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];
-            }
-        }
-        if (!changed) break;
-    }
-
+    // K-Means clustering
+    const { labels } = kmeans(embeddings, k);
+    // UMAP projection
+    const { UMAP } = await import('https://cdn.jsdelivr.net/npm/umap-js@1.4.0/+esm');
     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 }, () => []);
     for (let i = 0; i < n; ++i)
         clustered[labels[i]].push(lines[i]);
-
+    // Prepare scatter plot traces
     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}`);
     const traces = Array.from({ length: k }, (_, c) => ({
@@ -164,70 +82,21 @@ document.getElementById("kmeans-btn").onclick = async () => {
         traces[labels[i]].y.push(proj[i][1]);
         traces[labels[i]].text.push(lines[i]);
     }
-    Plotly.newPlot("plot-scatter", traces, {
-        xaxis: { title: "UMAP-1", scaleanchor: "y", scaleratio: 1 },
-        yaxis: { title: "UMAP-2", scaleanchor: "x", scaleratio: 1 },
-        width: 1000,
-        height: 500,
-        margin: { t: 40, l: 40, r: 10, b: 40 },
-        title: `K-Means Clustering (k=${k})`,
-        legend: { x: 1.05, y: 0.5, orientation: "v", xanchor: "left", yanchor: "middle" }
-    });
-
+    plotScatter(traces, k);
+    // Cluster naming
     const clusterNames = [];
     for (let c = 0; c < k; ++c) {
         progressBarInner.style.width = `${Math.round(((c + 1) / k) * 100)}%`;
-
-        const joined = clustered[c].join("\n");
-        const messages = [
-            { role: "system", content: prompt_cluster },
-            { role: "user", content: `Input:\n${joined}\nOutput:` }
-        ];
-
-        const inputs = tokenizer.apply_chat_template(messages, {
-            add_generation_prompt: true,
-            return_dict: true,
-            enable_thinking: false,
-        });
-
-        const outputTokens = await model.generate({
-            ...inputs,
-            max_new_tokens: 1024,
-            do_sample: true,
-            temperature: 0.6
-        });
-
-        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)) {
-            rawName = rawName.substring(rawName.lastIndexOf(THINK_TAG) + THINK_TAG.length).trim();
-        }
-        if (rawName.includes(END_TAG)) {
-            rawName = rawName.substring(0, rawName.indexOf(END_TAG)).trim();
-        }
-
-        clusterNames.push(rawName || `Cluster ${c + 1}`);
+        const name = await nameCluster(clustered[c]);
+        clusterNames.push(name || `Cluster ${c + 1}`);
         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 },
-            width: 1000,
-            height: 500,
-            margin: { t: 40, l: 40, r: 10, b: 40 },
-            title: `K-Means Clustering (k=${k})`,
-            legend: { x: 1.05, y: 0.5, orientation: "v", xanchor: "left", yanchor: "middle" }
-        });
-
+        updateScatter(traces, k);
+        // Update textarea with cluster names as markdown headers
         document.getElementById("input").value = clustered.map((g, i) =>
             `## ${clusterNames[i]}\n${g.join("\n")}`
         ).join("\n\n\n");
-
+        // Update heatmap with new cluster names
         document.getElementById("run").onclick();
     }
-
-    progressBarInner.style.width = "100%"; // Set to 100% after all clusters are named
+    progressBarInner.style.width = "100%";
 };

plotting.js

@@ -0,0 +1,35 @@
+export function plotHeatmap(sim, xLabels, yLabels) {
+    const data = [{ z: sim, type: "heatmap", colorscale: "Viridis", zmin: 0.7, zmax: 1, x: xLabels, y: yLabels }];
+    Plotly.newPlot("plot-heatmap", data, {
+        xaxis: { title: "Group", scaleanchor: "y", scaleratio: 1 },
+        yaxis: { title: "Group", scaleanchor: "x", scaleratio: 1 },
+        width: 500,
+        height: 500,
+        margin: { t: 40, l: 200, r: 10, b: 200 },
+        title: "Group Similarity Heatmap"
+    });
+}
+
+export function plotScatter(traces, k) {
+    Plotly.newPlot("plot-scatter", traces, {
+        xaxis: { title: "UMAP-1", scaleanchor: "y", scaleratio: 1 },
+        yaxis: { title: "UMAP-2", scaleanchor: "x", scaleratio: 1 },
+        width: 1000,
+        height: 500,
+        margin: { t: 40, l: 40, r: 10, b: 40 },
+        title: `K-Means Clustering (k=${k})`,
+        legend: { x: 1.05, y: 0.5, orientation: "v", xanchor: "left", yanchor: "middle" }
+    });
+}
+
+export function updateScatter(traces, k) {
+    Plotly.react("plot-scatter", traces, {
+        xaxis: { title: "UMAP-1", scaleanchor: "y", scaleratio: 1 },
+        yaxis: { title: "UMAP-2", scaleanchor: "x", scaleratio: 1 },
+        width: 1000,
+        height: 500,
+        margin: { t: 40, l: 40, r: 10, b: 40 },
+        title: `K-Means Clustering (k=${k})`,
+        legend: { x: 1.05, y: 0.5, orientation: "v", xanchor: "left", yanchor: "middle" }
+    });
+}