Add cluster name extraction for heatmap axis labels in K-Means clustering

main.js

@@ -13,7 +13,16 @@ const task = "Given a textual input sentence, retrieve relevant categories that
 
 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) {
         const lines = g.split(/\n/).filter(x => x.trim() != "");
@@ -41,7 +50,9 @@ document.getElementById("run").onclick = async () => {
         }
         sim.push(row);
     }
-    const data = [{ z: sim, type: "heatmap", colorscale: "Viridis", zmin: 0, zmax: 1 }];
+    // 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}`);
+    const data = [{ z: sim, type: "heatmap", colorscale: "Viridis", zmin: 0, 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 },
@@ -208,4 +219,4 @@ document.getElementById("kmeans-btn").onclick = async () => {
         // Re-run heatmap after updating textarea
         document.getElementById("run").onclick();
     }
-};
+};