test: add script to generate demo data

Author: donghaoren

.github/workflows/ci.yml

@@ -31,5 +31,15 @@ jobs:
       - name: Install dependencies
         run: npm install
 
+      - uses: actions/cache@v4
+        id: cache-demo-data
+        with:
+          path: packages/docs/demo-data
+          key: ${{ runner.os }}-${{ hashFiles('packages/docs/generate_demo_data.py') }}
+
+      - name: Generate demo data
+        if: steps.cache-demo-data.outputs.cache-hit != 'true'
+        run: cd packages/docs && uv run generate_demo_data.py
+
       - name: Run custom build script
         run: ./scripts/build.sh

packages/docs/.gitignore

@@ -2,3 +2,4 @@
 .vitepress/dist
 public/upload
 public/demo
+demo-data

packages/docs/generate_assets.sh

@@ -2,14 +2,14 @@
 
 set -euxo pipefail
 
+# Create the upload page
 rm -rf public/upload
 cp -r ../viewer/dist public/upload
 python -c "fn='public/upload/index.html';c=open(fn).read().replace('viewer','upload');open(fn,'w').write(c);"
 
-DEMO_DATA_FOLDER=../../../embedding-atlas-demo/data
-
-rm -rf public/demo
-if [ -d "$DEMO_DATA_FOLDER" ]; then
-cp -r ../viewer/dist public/demo
-cp -r "$DEMO_DATA_FOLDER" public/demo/data
+# Create the demo page
+if [ -d "demo-data" ]; then
+    rm -rf public/demo
+    cp -r ../viewer/dist public/demo
+    cp -r demo-data public/demo/data
 fi

packages/docs/generate_demo_data.py

@@ -0,0 +1,80 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = ["click", "datasets", "pandas", "sentence-transformers", "umap-learn"]
+# ///
+
+import json
+import os
+import shutil
+
+import click
+import pandas as pd
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+from umap import UMAP
+from umap.umap_ import nearest_neighbors
+
+
+def add_embedding_projection(df: pd.DataFrame, text: str):
+    texts = list(df[text])
+
+    transformer = SentenceTransformer("all-MiniLM-L6-v2")
+    hidden_vectors = transformer.encode(texts)
+
+    knn = nearest_neighbors(
+        hidden_vectors,
+        n_neighbors=15,
+        metric="cosine",
+        metric_kwds=None,
+        angular=False,
+        random_state=None,
+    )
+
+    proj = UMAP(metric="cosine", precomputed_knn=knn).fit_transform(hidden_vectors)
+
+    df["projection_x"] = proj[:, 0]  # type: ignore
+    df["projection_y"] = proj[:, 1]  # type: ignore
+    df["__neighbors"] = [{"distances": b, "ids": a} for a, b in zip(knn[0], knn[1])]
+
+
+@click.command()
+@click.option("--output", default="demo-data")
+def main(output: str):
+    shutil.rmtree(output, ignore_errors=True)
+    os.makedirs(output, exist_ok=True)
+
+    name = "spawn99/wine-reviews"
+    columns = [
+        "country",
+        "province",
+        "description",
+        "points",
+        "price",
+        "variety",
+        "designation",
+    ]
+
+    ds = load_dataset(name, split="train")
+    df = ds.to_pandas().sample(100)[columns]  # type: ignore
+
+    add_embedding_projection(df, text="description")
+
+    df.to_parquet(os.path.join(output, "dataset.parquet"), index=False)
+
+    metadata = {
+        "columns": {
+            "id": "_row_index",
+            "text": "description",
+            "embedding": {"x": "projection_x", "y": "projection_y"},
+            "neighbors": "__neighbors",
+        },
+        "is_static": True,
+        "database": {"type": "wasm", "load": True},
+    }
+
+    with open(os.path.join(output, "metadata.json"), "wb") as f:
+        f.write(json.dumps(metadata).encode("utf-8"))
+
+
+if __name__ == "__main__":
+    main()