refactor: [breaking] use Intl.Segmenter to generate cluster labels

package.json

@@ -1,4 +1,5 @@
 {
+  "name": "@embedding-atlas/workspace",
   "private": true,
   "workspaces": [
     "packages/component",

packages/component/package.json

@@ -47,6 +47,7 @@
     "publint": "^0.3.12",
     "quickselect": "^3.0.0",
     "simplify-js": "^1.2.4",
+    "stemmer": "^2.0.1",
     "svelte": "^5.37.3",
     "typescript": "^5.9.2",
     "vite": "^7.0.6",

packages/component/src/demo/EmbeddingViewDemo.svelte

@@ -37,8 +37,8 @@
     return { x: data.x[minIndex], y: data.y[minIndex], text: dataset[minIndex].text, fields: {} };
   }
 
-  async function queryClusterLabels(rects: Rectangle[]): Promise<string | null> {
-    return "label";
+  async function queryClusterLabels(clusters: Rectangle[][]): Promise<(string | null)[]> {
+    return clusters.map(() => "label");
   }
 </script>
 

packages/component/src/lib/embedding_view/EmbeddingViewImpl.svelte

@@ -18,7 +18,7 @@
     totalCount: number | null;
     maxDensity: number | null;
     automaticLabels: AutomaticLabelsConfig | boolean;
-    queryClusterLabels: ((rects: Rectangle[]) => Promise<string | null>) | null;
+    queryClusterLabels: ((clusters: Rectangle[][]) => Promise<(string | null)[]>) | null;
     tooltip: Selection | null;
     selection: Selection[] | null;
     querySelection: ((x: number, y: number, unitDistance: number) => Promise<Selection | null>) | null;
@@ -595,12 +595,11 @@
     let newClusters = await generateClusters(renderer, 10, viewport);
     newClusters = newClusters.concat(await generateClusters(renderer, 5, viewport));
 
-    statusMessage = "Generating labels (initializing)...";
+    statusMessage = "Generating labels...";
     if (queryClusterLabels) {
+      let labels = await queryClusterLabels(newClusters.map((x) => x.rects));
       for (let i = 0; i < newClusters.length; i++) {
-        let label = await queryClusterLabels(newClusters[i].rects);
-        newClusters[i].label = label;
-        statusMessage = `Generating labels (${(((i + 1) / newClusters.length) * 100).toFixed(0)}%)...`;
+        newClusters[i].label = labels[i];
       }
     }
 

packages/component/src/lib/embedding_view/EmbeddingViewMosaic.svelte

@@ -6,7 +6,6 @@
 
   import EmbeddingViewImpl from "./EmbeddingViewImpl.svelte";
 
-  import { TextSummarizer } from "../text_summarizer/text_summarizer.js";
   import { deepEquals, type Point, type Rectangle, type ViewportState } from "../utils.js";
   import type { EmbeddingViewMosaicProps } from "./embedding_view_mosaic_api.js";
   import {
@@ -17,6 +16,12 @@
   } from "./mosaic_client.js";
   import { makeClient } from "./mosaic_helper.js";
   import type { DataPoint, DataPointID } from "./types.js";
+  import {
+    textSummarizerAdd,
+    textSummarizerCreate,
+    textSummarizerDestroy,
+    textSummarizerSummarize,
+  } from "./worker/index.js";
 
   let {
     coordinator = defaultCoordinator(),
@@ -328,20 +333,52 @@
   }
 
   // Cluster Labels
-  let textSummarizer = $derived(
-    text != null ? new TextSummarizer({ coordinator: coordinator, table: table, x: x, y: y, text: text }) : null,
-  );
-
-  async function queryClusterLabels(rects: Rectangle[]): Promise<string | null> {
-    if (textSummarizer == null) {
-      return null;
+  async function queryClusterLabels(clusters: Rectangle[][]): Promise<(string | null)[]> {
+    if (text == null) {
+      return clusters.map(() => null);
     }
-    let list = await textSummarizer.summarize(rects, 4);
-    if (list.length > 0) {
-      return list.slice(0, 2).join("-") + "-\n" + list.slice(2).join("-");
-    } else {
-      return null;
+    // Create text summarizer (in the worker)
+    let summarizer = await textSummarizerCreate({ regions: clusters });
+    // Add text data to the summarizer
+    let start = 0;
+    let chunkSize = 10000;
+    let lastAdd: Promise<unknown> | null = null;
+    while (true) {
+      let r: any = await coordinator.query(
+        SQL.Query.from(table)
+          .select({ x: SQL.column(x), y: SQL.column(y), text: SQL.column(text) })
+          .offset(start)
+          .limit(chunkSize),
+      );
+      let data = {
+        x: r.getChild("x").toArray(),
+        y: r.getChild("y").toArray(),
+        text: r.getChild("text").toArray(),
+      };
+      if (lastAdd != null) {
+        await lastAdd;
+      }
+      lastAdd = textSummarizerAdd(summarizer, data);
+      if (r.getChild("text").length < chunkSize) {
+        break;
+      }
+      start += chunkSize;
+    }
+    if (lastAdd != null) {
+      await lastAdd;
     }
+    let summarizeResult = await textSummarizerSummarize(summarizer);
+    await textSummarizerDestroy(summarizer);
+
+    return summarizeResult.map((words) => {
+      if (words.length == 0) {
+        return null;
+      } else if (words.length > 2) {
+        return words.slice(0, 2).join("-") + "-\n" + words.slice(2).join("-");
+      } else {
+        return words.join("-");
+      }
+    });
   }
 </script>
 

packages/component/src/lib/embedding_view/embedding_view_api.ts

@@ -66,8 +66,8 @@ export interface EmbeddingViewProps {
   /** A function to query selected point given (x, y) location, and a unit distance (distance of 1pt in data units). */
   querySelection?: ((x: number, y: number, unitDistance: number) => Promise<DataPoint | null>) | null;
 
-  /** A function that returns a summary label for points covered by the union of the given rectangles. */
-  queryClusterLabels?: ((rects: Rectangle[]) => Promise<string | null>) | null;
+  /** A function that returns summary labels for clusters. Each cluster is given by a list of rectangles that approximate its shape. */
+  queryClusterLabels?: ((clusters: Rectangle[][]) => Promise<(string | null)[]>) | null;
 
   /** A callback for when viewportState changes. */
   onViewportState?: ((value: ViewportState) => void) | null;

packages/component/src/lib/embedding_view/theme.ts

@@ -31,7 +31,7 @@ const defaultThemeConfig: { light: ThemeConfig; dark: ThemeConfig } = {
   },
   dark: {
     fontFamily: "system-ui,sans-serif",
-    clusterLabelColor: "#fff",
+    clusterLabelColor: "#ccc",
     clusterLabelOutlineColor: "rgba(0,0,0,0.8)",
     clusterLabelOpacity: 0.8,
     statusBar: true,

packages/component/src/lib/embedding_view/worker/clustering.worker.js

@@ -1,16 +1,32 @@
 // Copyright (c) 2025 Apple Inc. Licensed under MIT License.
 
-import { dynamicLabelPlacement, findClusters } from "./worker_functions.js";
+import {
+  dynamicLabelPlacement,
+  findClusters,
+  textSummarizerAdd,
+  textSummarizerCreate,
+  textSummarizerDestroy,
+  textSummarizerSummarize,
+} from "./worker_functions.js";
+
+/** @type Record<string, (...args: any[]) => any> */
+let functions = {
+  dynamicLabelPlacement,
+  findClusters,
+  textSummarizerCreate,
+  textSummarizerAdd,
+  textSummarizerDestroy,
+  textSummarizerSummarize,
+};
 
 onmessage = async (msg) => {
-  if (msg.data.name == "findClusters") {
-    let args = msg.data.payload;
-    let clusters = await findClusters(args.density_map, args.width, args.height, args.options);
-    postMessage({ id: msg.data.id, payload: clusters });
-  }
-  if (msg.data.name == "dynamicLabelPlacement") {
+  if (functions[msg.data.name]) {
+    let func = functions[msg.data.name];
     let args = msg.data.payload;
-    let result = dynamicLabelPlacement(args.labels, args.options);
+    let result = func(...args);
+    if (result instanceof Promise) {
+      result = await result;
+    }
     postMessage({ id: msg.data.id, payload: result });
   }
 };

packages/component/src/lib/embedding_view/worker/index.ts

@@ -40,10 +40,26 @@ function invokeWorker(name: string, payload: any, transfer: Transferable[] = [])
 
 type PromiseReturn<F extends (...args: any) => any> = (...args: Parameters<F>) => Promise<ReturnType<F>>;
 
-export let findClusters: PromiseReturn<typeof Functions.findClusters> = (density_map, width, height, options) => {
-  return invokeWorker("findClusters", { density_map, width, height, options: options }, [density_map.buffer]);
+export let findClusters: PromiseReturn<typeof Functions.findClusters> = (densityMap, width, height, options) => {
+  return invokeWorker("findClusters", [densityMap, width, height, options], [densityMap.buffer]);
 };
 
-export let dynamicLabelPlacement: PromiseReturn<typeof Functions.dynamicLabelPlacement> = (labels, options) => {
-  return invokeWorker("dynamicLabelPlacement", { labels, options });
+export let dynamicLabelPlacement: PromiseReturn<typeof Functions.dynamicLabelPlacement> = (...args) => {
+  return invokeWorker("dynamicLabelPlacement", args);
+};
+
+export let textSummarizerCreate: PromiseReturn<typeof Functions.textSummarizerCreate> = (...args) => {
+  return invokeWorker("textSummarizerCreate", args);
+};
+
+export let textSummarizerDestroy: PromiseReturn<typeof Functions.textSummarizerDestroy> = (...args) => {
+  return invokeWorker("textSummarizerDestroy", args);
+};
+
+export let textSummarizerAdd: PromiseReturn<typeof Functions.textSummarizerAdd> = (...args) => {
+  return invokeWorker("textSummarizerAdd", args);
+};
+
+export let textSummarizerSummarize: PromiseReturn<typeof Functions.textSummarizerSummarize> = (...args) => {
+  return invokeWorker("textSummarizerSummarize", args);
 };

packages/component/src/lib/embedding_view/worker/worker_functions.ts

@@ -2,5 +2,30 @@
 
 import { findClusters } from "@embedding-atlas/density-clustering";
 import { dynamicLabelPlacement } from "../../dynamic_label_placement/dynamic_label_placement.js";
+import { TextSummarizer } from "../../text_summarizer/text_summarizer.js";
+import type { Rectangle } from "../../utils.js";
 
 export { dynamicLabelPlacement, findClusters };
+
+let textSummarizers = new Map<string, TextSummarizer>();
+
+export function textSummarizerCreate(options: { regions: Rectangle[][]; stopWords?: string[] }) {
+  let key = new Date().getTime() + "-" + Math.random();
+  textSummarizers.set(key, new TextSummarizer(options));
+  return key;
+}
+
+export function textSummarizerDestroy(key: string) {
+  return textSummarizers.delete(key);
+}
+
+export function textSummarizerAdd(
+  key: string,
+  data: { x: ArrayLike<number>; y: ArrayLike<number>; text: ArrayLike<string> }
+) {
+  textSummarizers.get(key)?.add(data);
+}
+
+export function textSummarizerSummarize(key: string) {
+  return textSummarizers.get(key)?.summarize() ?? [];
+}