scripts/compute-centroids.mjs

#!/usr/bin/env node
/**
 * Offline k-means centroid computation for the k-SemStamp demo.
 *
 * Embeds a corpus with the SAME encoder, the SAME task prefix and the SAME
 * call shape the browser uses (EmbeddingGemma via AutoModel ->
 * `sentence_embedding`). Any difference here puts the centroids in a
 * different space from the runtime assignments, and every cluster decision
 * becomes noise.
 *
 * Corpus: wikitext-2 sentences from the HF datasets server, plus narrative and
 * Korean sentences so the clusters cover what the demo actually generates.
 *
 * Usage: node scripts/compute-centroids.mjs
 */

import { AutoModel, AutoTokenizer } from '@huggingface/transformers';
import { mkdirSync, writeFileSync } from 'node:fs';
import { dirname, join } from 'node:path';
import { fileURLToPath } from 'node:url';

const ROOT = join(dirname(fileURLToPath(import.meta.url)), '..');
const OUT = join(ROOT, 'public', 'centroids', 'embeddinggemma-k8.json');
const ENCODER = 'onnx-community/embeddinggemma-300m-ONNX';
const CLUSTERING_PREFIX = 'task: clustering | query: ';
/**
 * Quantization changes the geometry: q4 and fp32 embeddings of the same
 * sentence sit at cosine ~0.98, enough to flip cluster assignments near a
 * boundary. The browser must use this same dtype, so it travels with the
 * centroids and is checked at load.
 */
const DTYPE = 'q4';
const K = 8;
const TARGET_SENTENCES = 400;
const SEED = 12345;

// Deterministic PRNG (mulberry32) so centroids are reproducible.
function mulberry32(seed) {
  let a = seed >>> 0;
  return () => {
    a |= 0;
    a = (a + 0x6d2b79f5) | 0;
    let t = Math.imul(a ^ (a >>> 15), 1 | a);
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}

async function fetchWikitextSentences() {
  const sentences = [];
  for (let offset = 0; offset < 1200 && sentences.length < TARGET_SENTENCES; offset += 100) {
    const url = `https://datasets-server.huggingface.co/rows?dataset=Salesforce%2Fwikitext&config=wikitext-2-raw-v1&split=train&offset=${offset}&length=100`;
    const res = await fetch(url);
    if (!res.ok) throw new Error(`datasets-server HTTP ${res.status}`);
    const json = await res.json();
    for (const row of json.rows ?? []) {
      const text = (row.row?.text ?? '').trim();
      if (!text || text.startsWith('=')) continue;
      for (const s of text.split(/(?<=[.!?])\s+/)) {
        const t = s.trim();
        if (t.length >= 40 && t.length <= 300 && /[.!?]$/.test(t)) sentences.push(t);
        if (sentences.length >= TARGET_SENTENCES) break;
      }
      if (sentences.length >= TARGET_SENTENCES) break;
    }
  }
  return sentences;
}

const FALLBACK = [
  'The committee approved the budget after a lengthy debate about infrastructure spending.',
  'Photosynthesis converts sunlight, water, and carbon dioxide into glucose and oxygen.',
  'The orchestra performed the symphony with remarkable precision and emotional depth.',
  'Quarterly earnings exceeded analyst expectations despite rising material costs.',
  'The hikers reached the summit just before the storm clouds rolled in.',
  'Researchers discovered a new species of frog in the remote rainforest valley.',
  'The museum unveiled a restored painting that had been hidden for decades.',
  'Global temperatures have risen steadily over the past century, reshaping ecosystems.',
  'The chef combined traditional techniques with unexpected local ingredients.',
  'Engineers tested the bridge design under simulated earthquake conditions.',
  'The novel follows three generations of a family through war and migration.',
  'Voter turnout reached a record high in several coastal districts this year.',
  'The spacecraft transmitted its first images of the icy moon at dawn.',
  'Local farmers adopted drip irrigation to cope with the prolonged drought.',
  'The startup raised new funding to expand its battery recycling operations.',
  'Archaeologists uncovered pottery fragments dating back three thousand years.',
  'The goalkeeper made a spectacular save in the final minute of the match.',
  'New regulations require companies to disclose their supply chain emissions.',
  'The pianist practiced the difficult passage until it sounded effortless.',
  'Coral reefs support a quarter of all marine species despite covering little area.',
];

function kmeans(vectors, k, rng) {
  const n = vectors.length;
  const dim = vectors[0].length;
  // k-means++ init
  const centroids = [vectors[Math.floor(rng() * n)].slice()];
  while (centroids.length < k) {
    const d2 = vectors.map((v) => {
      let best = Infinity;
      for (const c of centroids) {
        let dot = 0;
        for (let i = 0; i < dim; i++) dot += v[i] * c[i];
        const d = 1 - dot;
        if (d < best) best = d;
      }
      return best * best;
    });
    const sum = d2.reduce((a, b) => a + b, 0);
    let r = rng() * sum;
    let idx = 0;
    for (; idx < n - 1; idx++) {
      r -= d2[idx];
      if (r <= 0) break;
    }
    centroids.push(vectors[idx].slice());
  }

  const assign = new Array(n).fill(0);
  for (let iter = 0; iter < 100; iter++) {
    let changed = 0;
    for (let i = 0; i < n; i++) {
      let best = 0;
      let bestD = Infinity;
      for (let c = 0; c < k; c++) {
        let dot = 0;
        for (let j = 0; j < dim; j++) dot += vectors[i][j] * centroids[c][j];
        const d = 1 - dot;
        if (d < bestD) {
          bestD = d;
          best = c;
        }
      }
      if (assign[i] !== best) {
        assign[i] = best;
        changed++;
      }
    }
    // recompute (spherical: mean then renormalize)
    for (let c = 0; c < k; c++) {
      const sum = new Array(dim).fill(0);
      let count = 0;
      for (let i = 0; i < n; i++) {
        if (assign[i] !== c) continue;
        count++;
        for (let j = 0; j < dim; j++) sum[j] += vectors[i][j];
      }
      if (count === 0) continue;
      let norm = 0;
      for (let j = 0; j < dim; j++) norm += sum[j] * sum[j];
      norm = Math.sqrt(norm) || 1;
      for (let j = 0; j < dim; j++) centroids[c][j] = sum[j] / norm;
    }
    if (changed === 0) break;
  }
  const sizes = new Array(k).fill(0);
  for (const a of assign) sizes[a]++;
  return { centroids, sizes };
}

// Narrative-register sentences: the LLM demo mostly generates stories, and
// clusters trained on encyclopedic text alone would collapse story sentences
// into one or two clusters, starving the rejection sampler.
const NARRATIVE = [
  'The old keeper climbed the spiral staircase as the wind howled outside.',
  'She had never seen a light quite like the one glowing beneath the waves.',
  'Every evening he wrote a single line in his weather-beaten journal.',
  'The village children whispered stories about the tower on the cliff.',
  'A strange ship appeared on the horizon just before midnight.',
  'His lantern flickered, and for a moment the darkness felt alive.',
  'The storm had passed, leaving the shore littered with pale driftwood.',
  'She wrapped the mysterious object in cloth and hid it beneath the floorboards.',
  'Nobody in town believed him, but the light kept returning every night.',
  'The fog rolled in thick, swallowing the beam before it reached the water.',
  'He remembered his grandfather saying the sea keeps what it takes.',
  'At dawn the tide revealed a door in the rocks that had never been there.',
  'Her heart pounded as the signal repeated: three flashes, then silence.',
  'The keeper brewed his tea and watched the horizon with tired eyes.',
  'Something metallic glinted in the sand where the gulls had gathered.',
  'The letters stopped coming in October, and the winter felt endless.',
  'He rowed out at first light, following the trail of green phosphorescence.',
  'The melody drifted up from the cellar, soft and impossibly old.',
  'They said the previous keeper had vanished without taking his coat.',
  'By morning the writing on the glass had faded, but she remembered every word.',
  'The captain lowered his voice and told them what he had seen in the trench.',
  'A key turned in a lock that no one had opened for forty years.',
  'The child drew the same spiral again and again in her notebook.',
  'When the generator failed, the stars felt suddenly very close.',
  'He counted the ships as they passed, and one of them counted him back.',
  'The tide pool held a color that did not belong to this coast.',
  'Her boots left the only prints on the long gray beach.',
  'The radio crackled with a voice reciting coordinates that did not exist.',
  'In the attic they found charts of a coastline no map had ever shown.',
  'The bell buoy rang twice, though the sea was perfectly still.',
];

// The encoder is multilingual, so the corpus should be too: clusters trained
// only on English would still "work" on Korean, but they would be describing
// a region of the space the Korean sentences barely enter.
const KOREAN = [
  '레몬 치킨 파스타 레시피를 알려줘.',
  '올리브유에 마늘을 볶다가 닭가슴살을 넣고 노릇하게 익힌다.',
  '면은 소금을 넉넉히 넣은 물에 알 덴테로 삶는다.',
  '마지막에 레몬즙과 파슬리를 뿌리면 향이 살아난다.',
  '등대지기는 매일 밤 수평선을 바라보며 기록을 남겼다.',
  '폭풍이 지나간 아침, 해변에는 낯선 나무 조각들이 밀려와 있었다.',
  '마을 사람들은 절벽 위 탑에 대한 이야기를 소곤거렸다.',
  '그는 등불을 들고 나선 계단을 천천히 올라갔다.',
  '연구진은 새로운 개구리 종을 열대 우림에서 발견했다.',
  '지난 100년 동안 지구의 평균 기온은 꾸준히 상승했다.',
  '이번 분기 실적은 원자재 가격 상승에도 예상치를 웃돌았다.',
  '위원회는 오랜 논의 끝에 예산안을 승인했다.',
  '박물관은 수십 년간 보관돼 있던 그림을 복원해 공개했다.',
  '광합성은 햇빛과 물, 이산화탄소를 포도당과 산소로 바꾼다.',
  '기술이 발전하면서 번역의 품질도 크게 좋아졌다.',
  '그 소설은 전쟁과 이주를 겪은 세 세대의 이야기를 따라간다.',
  '아이는 공책에 같은 나선을 몇 번이고 다시 그렸다.',
  '새벽이 되자 밀물이 빠지며 바위 사이의 문이 드러났다.',
  '엔지니어들은 지진을 가정한 조건에서 다리 설계를 시험했다.',
  '주말 농장에서는 물을 아끼기 위해 점적 관수를 도입했다.',
];

async function main() {
  let sentences;
  let corpus;
  try {
    const wiki = await fetchWikitextSentences();
    sentences = [...wiki, ...NARRATIVE, ...KOREAN];
    corpus = `wikitext-2-raw-v1 (${wiki.length}) + narrative (${NARRATIVE.length}) + Korean (${KOREAN.length})`;
    console.log(`Fetched ${wiki.length} wikitext + ${NARRATIVE.length} narrative + ${KOREAN.length} Korean`);
  } catch (e) {
    console.warn('wikitext fetch failed, using built-in fallback corpus:', e.message);
    sentences = [...FALLBACK, ...NARRATIVE, ...KOREAN];
    corpus = `built-in fallback (${sentences.length} sentences)`;
  }

  console.log('Loading EmbeddingGemma...');
  const tokenizer = await AutoTokenizer.from_pretrained(ENCODER);
  const model = await AutoModel.from_pretrained(ENCODER, { dtype: DTYPE });

  console.log(`Embedding ${sentences.length} sentences...`);
  const vectors = [];
  const BATCH = 16;
  for (let i = 0; i < sentences.length; i += BATCH) {
    const batch = sentences.slice(i, i + BATCH).map((s) => CLUSTERING_PREFIX + s);
    const inputs = await tokenizer(batch, { padding: true });
    const { sentence_embedding } = await model(inputs);
    const [n, d] = sentence_embedding.dims;
    for (let r = 0; r < n; r++) {
      vectors.push(Array.from(sentence_embedding.data.slice(r * d, (r + 1) * d)));
    }
    console.log(`  ${Math.min(i + BATCH, sentences.length)}/${sentences.length}`);
  }

  // Mean-center before clustering. Sentence-transformer embeddings share a
  // large common component; without removing it every English sentence sits
  // in one dominant cone and k-means degenerates into a single attractor
  // (observed: all candidates landing in C0). The same mean is subtracted at
  // assignment time in the browser.
  const dim = vectors[0].length;
  const meanVec = new Array(dim).fill(0);
  for (const v of vectors) for (let j = 0; j < dim; j++) meanVec[j] += v[j] / vectors.length;
  const centered = vectors.map((v) => {
    const c = v.map((x, j) => x - meanVec[j]);
    const norm = Math.hypot(...c) || 1;
    return c.map((x) => x / norm);
  });

  console.log('Running spherical k-means on mean-centered embeddings...');
  const rng = mulberry32(SEED);
  const { centroids, sizes } = kmeans(centered, K, rng);
  console.log('Cluster sizes:', sizes);

  const payload = {
    vectors: centroids.map((c) => c.map((x) => Number(x.toFixed(6)))),
    mean: meanVec.map((x) => Number(x.toFixed(6))),
    dim,
    prefix: CLUSTERING_PREFIX,
    dtype: DTYPE,
    encoder: ENCODER,
    corpus,
    k: K,
    seed: SEED,
    generatedAt: new Date().toISOString(),
  };
  mkdirSync(dirname(OUT), { recursive: true });
  writeFileSync(OUT, JSON.stringify(payload));
  console.log(`Wrote ${OUT}`);
}

main().catch((e) => {
  console.error(e);
  process.exit(1);
});