Update K-Means input range and enhance cluster name extraction logic for improved accuracy

index.html

@@ -57,7 +57,7 @@
       document.getElementById("input").value = sentences.join("\n");
     </script>
     <label for="kmeans-k" style="margin-left:10px;">Clusters:</label>
-    <input id="kmeans-k" type="number" min="2" max="20" value="3" style="width:60px;">
+    <input id="kmeans-k" type="number" min="2" max="100" value="7" style="width:60px;">
     <button id="kmeans-btn">K-Means Clustering</button>
     <button id="run">Similarity Heatmap</button>
     <div id="progress-bar">

main.js

@@ -124,7 +124,7 @@ document.getElementById("kmeans-btn").onclick = async () => {
             { role: "system", content: "Given the following texts, provide a short 1-3 word summary in plaintext" },
             { role: "user", content: `${joined}` }
         ];
-        
+
         const inputs = tokenizer.apply_chat_template(messages, {
             add_generation_prompt: true,
             return_dict: true,
@@ -165,30 +165,41 @@ document.getElementById("kmeans-btn").onclick = async () => {
             max_new_tokens: 1024,
             do_sample: true,
             temperature: 0.6,
-            streamer,
+            // streamer,
         });
-        let rawName = tokenizer.decode(outputTokens[0], { skip_special_tokens: false }).trim();
-        // Extract cluster name after last '</think>' if present, and remove '<|im_end|>' if present
-        let idx = rawName.lastIndexOf('</think>');
-        let name = idx !== -1 ? rawName.slice(idx + 8).trim() : rawName;
-        if (name.endsWith('<|im_end|>')) name = name.slice(0, -11).trim();
-        clusterNames.push(name.length > 0 ? name : `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];
+        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];
+        }
+        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" }
+        });
+        // Update textarea: group by cluster, separated by triple newlines
+        document.getElementById("input").value = clustered.map(g => g.join("\n")).join("\n\n\n");
+        // Re-run heatmap after updating textarea
+        document.getElementById("run").onclick();
     }
-    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" }
-    });
-    // Update textarea: group by cluster, separated by triple newlines
-    document.getElementById("input").value = clustered.map(g => g.join("\n")).join("\n\n\n");
-    // Re-run heatmap after updating textarea
-    document.getElementById("run").onclick();
 };