move to main

main.js

@@ -0,0 +1,176 @@
+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";
+
+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" });
+
+const task = "Given a textual input sentence, retrieve relevant categories that best describe it.";
+
+document.getElementById("run").onclick = async () => {
+    const text = document.getElementById("input").value;
+    const groups = text.split(/\n{3,}/);
+    const groupEmbeddings = [];
+    for (const g of groups) {
+        const lines = g.split(/\n/).filter(x => x.trim() != "");
+        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);
+    }
+    const n = groupEmbeddings.length;
+    const sim = [];
+    for (let i = 0; i < n; i++) {
+        const row = [];
+        for (let j = 0; j < n; j++) {
+            let dot = 0, na = 0, nb = 0;
+            for (let k = 0; k < groupEmbeddings[i].length; k++) {
+                dot += groupEmbeddings[i][k] * groupEmbeddings[j][k];
+                na += groupEmbeddings[i][k] ** 2;
+                nb += groupEmbeddings[j][k] ** 2;
+            }
+            row.push(dot / Math.sqrt(na * nb));
+        }
+        sim.push(row);
+    }
+    const data = [{ z: sim, type: "heatmap", colorscale: "Viridis", zmin: 0, zmax: 1 }];
+    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 progressBar = document.getElementById("progress-bar");
+    const progressBarInner = document.getElementById("progress-bar-inner");
+    progressBar.style.display = "block";
+    progressBarInner.style.width = "0%";
+
+    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;
+    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) {
+        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;
+        }
+        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);
+    // Group lines by cluster
+    const clustered = Array.from({ length: k }, (_, c) => []);
+    for (let i = 0; i < n; ++i) clustered[labels[i]].push(lines[i]);
+    // 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)}%`;
+        const joined = clustered[c].join("\n");
+        const messages = [
+            { role: "system", content: "You are a helpful assistant." },
+            { role: "user", content: `Given the following texts, provide a short, descriptive name for this group:\n\n${joined}` }
+        ];
+        const reasonEnabled = false;
+        const inputs = tokenizer.apply_chat_template(messages, {
+            add_generation_prompt: true,
+            return_dict: true,
+            enable_thinking: reasonEnabled,
+        });
+        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: 32,
+            do_sample: false,
+            streamer,
+        });
+        let name = tokenizer.decode(outputTokens[0], { skip_special_tokens: false }).trim();
+        clusterNames.push(name.length > 0 ? name : `Cluster ${c + 1}`);
+    }
+    progressBarInner.style.width = "100%";
+    setTimeout(() => { progressBar.style.display = "none"; }, 400);
+    // 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: clusterNames[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]);
+        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})`
+    });
+    // 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();
+};