developerjeremylive/colbert-tool-selection-etheroi

1"""Precompute the 2D UMAP layout of the tool embeddings for the background constellation.2 3Run once locally (needs the model + HF auth + umap-learn):4 5    python build_layout.py6 7Embeds every tool's routing_text with LFM2.5-Embedding-350M (the same text and prompt8the live search uses), projects to 2D with UMAP (cosine metric), normalises to [0,1]9with a small inset, and writes static/layout.json — a list of10{key: "domain|name", name, domain, x, y}. The frontend renders one dot per tool and11lights up the ones a query retrieves. The Space itself does not need umap-learn.12"""13 14from __future__ import annotations15 16import json17import pathlib18 19import numpy as np20 21import search as S22import toolset as T23 24OUT = pathlib.Path(__file__).resolve().parent / "static" / "layout.json"25 26 27def _normalize(xy: np.ndarray) -> np.ndarray:28    lo, hi = xy.min(axis=0), xy.max(axis=0)29    span = np.where(hi - lo == 0, 1.0, hi - lo)30    return (xy - lo) / span                         # -> [0, 1]; the frontend adds border padding31 32 33def main() -> None:34    import umap  # heavy (numba); only needed for this offline build35 36    tools = T.load_catalog()37    print(f"embedding {len(tools)} tools with {S.MODEL_ID} ...", flush=True)38    emb = S._encode([t.routing_text for t in tools], prompt_name="document")39 40    n_neighbors = min(15, max(2, len(tools) - 1))41    reducer = umap.UMAP(42        n_components=2, metric="cosine", n_neighbors=n_neighbors,43        min_dist=0.12, random_state=42,44    )45    xy = _normalize(np.asarray(reducer.fit_transform(emb), dtype=np.float64))46 47    pts = [48        {"key": f"{t.domain}|{t.name}", "name": t.name, "domain": t.domain,49         "x": round(float(xy[i, 0]), 4), "y": round(float(xy[i, 1]), 4)}50        for i, t in enumerate(tools)51    ]52    OUT.write_text(json.dumps(pts), encoding="utf-8")53    print(f"wrote {len(pts)} points to {OUT}", flush=True)54 55 56if __name__ == "__main__":57    main()58