Fix README short description capitalization and enhance index.html layout for K-Means clustering visualization

README.md

@@ -6,7 +6,7 @@ colorTo: yellow
 sdk: static
 pinned: false
 license: apache-2.0
-short_description: 'exploring text embeddings and group similarity '
+short_description: 'Exploring text embeddings and group similarity '
 ---
 
 # Embedding WebGPU Playground
@@ -23,6 +23,10 @@ This is a browser-based playground for exploring text embeddings and group simil
 - Cosine similarity is calculated between all group embeddings, resulting in a group-by-group similarity matrix.
 - The similarity matrix is visualized as a heatmap using Plotly (color range locked to 0–1).
 
+
+K-Means Clustering
+re group text by using K-Means and number of group
+
 ## Tech stack
 - [@huggingface/transformers](https://www.npmjs.com/package/@huggingface/transformers) (ESM, WebGPU)
 - [ONNX Qwen3-Embedding-0.6B-ONNX](https://huggingface.co/onnx-community/Qwen3-Embedding-0.6B-ONNX)

index.html

@@ -21,10 +21,16 @@
             margin-top: 10px
         }
 
-        #plot {
+        #plot-heatmap,
+        #plot-scatter {
             width: 100%;
             height: 600px
         }
+
+        .plot-container {
+            display: flex;
+            gap: 20px;
+        }
     </style>
 </head>
 
@@ -32,10 +38,16 @@
     <h1>Embedding Similarity Heatmap</h1>
     <textarea id="input"></textarea>
     <button id="run">Run</button>
-    <div id="plot"></div>
+    <input id="kmeans-k" type="number" min="2" max="20" value="3" style="width:60px; margin-left:10px;"> <button
+        id="kmeans-btn">K-Means Clustering</button>
+    <div class="plot-container">
+        <div id="plot-heatmap" style="width:500px; height:500px;"></div>
+        <div id="plot-scatter" style="width:500px; height:500px;"></div>
+    </div>
     <script src="https://cdn.plot.ly/plotly-2.32.0.min.js"></script>
     <script type="module">
         import { pipeline } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.5.2";
+        import { UMAP } from "https://cdn.jsdelivr.net/npm/umap-js@1.4.0/+esm";
 
         const embed = await pipeline(
             "feature-extraction",
@@ -75,7 +87,72 @@
                 sim.push(row);
             }
             const data = [{ z: sim, type: "heatmap", colorscale: "Viridis", zmin: 0, zmax: 1 }];
-            Plotly.newPlot("plot", data, { xaxis: { title: "Group" }, yaxis: { title: "Group" } });
+            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: 40, r: 10, b: 40 },
+                title: "Group Similarity Heatmap"
+            });
+        };
+
+        // --- K-Means Clustering ---
+        document.getElementById("kmeans-btn").onclick = async () => {
+            const text = document.getElementById("input").value;
+            const lines = text.split(/\n/).map(x => x.trim()).filter(x => x);
+            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;
+
+            // 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;
+            // Randomly initialize centroids
+            let centroids = Array.from({length: k}, () => embeddings[Math.floor(Math.random()*n)].slice());
+            let labels = new Array(n).fill(0);
+            for (let iter = 0; iter < 20; ++iter) {
+                // Assign
+                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; }
+                    }
+                    labels[i] = best;
+                }
+                // Update
+                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]) for (let d = 0; d < dim; ++d) centroids[c][d] /= counts[c];
+            }
+            // UMAP for 2D projection
+            const umap = new UMAP({ nComponents: 2 });
+            const proj = umap.fit(embeddings);
+            // Plot
+            const colors = ["red","blue","green","orange","purple","cyan","magenta","yellow","brown","black","lime","navy","teal","olive","maroon","pink","gray","gold","aqua","indigo"];
+            const traces = Array.from({length: k}, (_, c) => ({
+                x: [], y: [], text: [], mode: "markers", type: "scatter", name: `Cluster ${c+1}`,
+                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]);
+                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: 500,
+                height: 500,
+                margin: { t: 40, l: 40, r: 10, b: 40 },
+                title: `K-Means Clustering (k=${k})`
+            });
         };
     </script>
 </body>