From 681e632b58a359c45d2dffc92e16254a0dee3dc4 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Apr 2026 16:04:08 +0000 Subject: [PATCH] Add unit-circle embedding benchmark project MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Full Replit-ready benchmark for testing whether a unit-circle / number-base coordinate remapping improves retrieval, clustering, and neighborhood preservation on top of proven embedding models. - src/transforms.py: full A0–A5 ablation ladder (identity, PCA, random rotation, UC direct, UC rank, UC broken-phase control) - src/benchmark.py: retrieval (Recall@k, MRR, nDCG), clustering (silhouette, NMI, ARI), geometry (neighbor overlap, Spearman, trustworthiness/continuity) - src/embed.py: SentenceTransformer wrapper with task-prefix support - src/config.py: model presets (nomic, bge-m3) and ablation configs - src/main.py: full ablation runner with JSON output and success checks - src/metrics.py: table formatting, delta-vs-baseline, success criteria - data/: 20-doc toy corpus (4 topics) + 10 queries + qrels - notebooks/sanity_checks.ipynb: bar charts, neighbor overlap, UC geometry - requirements.txt, pyproject.toml, .replit https://claude.ai/code/session_01EWFjWpMFhdAcdpST3syGPE --- unit-circle-benchmark/.replit | 17 ++ unit-circle-benchmark/README.md | 112 +++++++++++ unit-circle-benchmark/data/corpus.jsonl | 20 ++ unit-circle-benchmark/data/qrels.jsonl | 12 ++ unit-circle-benchmark/data/queries.jsonl | 10 + .../notebooks/sanity_checks.ipynb | 185 ++++++++++++++++++ unit-circle-benchmark/pyproject.toml | 28 +++ unit-circle-benchmark/requirements.txt | 12 ++ unit-circle-benchmark/src/__init__.py | 0 unit-circle-benchmark/src/benchmark.py | 153 +++++++++++++++ unit-circle-benchmark/src/config.py | 57 ++++++ unit-circle-benchmark/src/data.py | 51 +++++ unit-circle-benchmark/src/embed.py | 47 +++++ unit-circle-benchmark/src/main.py | 138 +++++++++++++ unit-circle-benchmark/src/metrics.py | 98 ++++++++++ unit-circle-benchmark/src/transforms.py | 151 ++++++++++++++ unit-circle-benchmark/src/utils.py | 25 +++ 17 files changed, 1116 insertions(+) create mode 100644 unit-circle-benchmark/.replit create mode 100644 unit-circle-benchmark/README.md create mode 100644 unit-circle-benchmark/data/corpus.jsonl create mode 100644 unit-circle-benchmark/data/qrels.jsonl create mode 100644 unit-circle-benchmark/data/queries.jsonl create mode 100644 unit-circle-benchmark/notebooks/sanity_checks.ipynb create mode 100644 unit-circle-benchmark/pyproject.toml create mode 100644 unit-circle-benchmark/requirements.txt create mode 100644 unit-circle-benchmark/src/__init__.py create mode 100644 unit-circle-benchmark/src/benchmark.py create mode 100644 unit-circle-benchmark/src/config.py create mode 100644 unit-circle-benchmark/src/data.py create mode 100644 unit-circle-benchmark/src/embed.py create mode 100644 unit-circle-benchmark/src/main.py create mode 100644 unit-circle-benchmark/src/metrics.py create mode 100644 unit-circle-benchmark/src/transforms.py create mode 100644 unit-circle-benchmark/src/utils.py diff --git a/unit-circle-benchmark/.replit b/unit-circle-benchmark/.replit new file mode 100644 index 0000000..564ba51 --- /dev/null +++ b/unit-circle-benchmark/.replit @@ -0,0 +1,17 @@ +run = "python -m src.main" +entrypoint = "src/main.py" + +[nix] +channel = "stable-24_05" + +[deployment] +run = ["sh", "-c", "python -m src.main"] +deploymentTarget = "cloudrun" + +[env] +# Switch to bge-m3 for Phase 2: MODEL_PRESET=bge-m3 +MODEL_PRESET = "nomic" + +[[ports]] +localPort = 8080 +externalPort = 80 diff --git a/unit-circle-benchmark/README.md b/unit-circle-benchmark/README.md new file mode 100644 index 0000000..fd83bff --- /dev/null +++ b/unit-circle-benchmark/README.md @@ -0,0 +1,112 @@ +# Unit Circle Embedding Benchmark + +This repository tests whether a unit-circle / number-base coordinate remapping can improve the usefulness of existing embedding spaces. + +## Principle + +We do not train a new encoder first. +We begin with a proven embedding model, then test whether a derived coordinate system improves retrieval, clustering, and neighborhood structure. + +## Initial substrates + +- `nomic-ai/nomic-embed-text-v1.5` — Phase 1 (Matryoshka-style dim reduction, long context) +- `BAAI/bge-m3` — Phase 2 (dense + sparse + multi-vector, multilingual) +- `text-embedding-3-large` — optional Phase 3 hosted ceiling-check + +## Core question + +Does the transform preserve or improve semantic structure better than: +1. native vectors +2. matched-dimension PCA +3. negative-control transforms (random rotation, broken phase) + +--- + +## Ablation ladder + +| Step | Name | Description | +|------|-----------------|----------------------------------------------------------| +| A0 | `native` | Plain L2-normalised vectors | +| A1 | `pca_256/64` | PCA to same target dim — is gain just dimensional cleanup? | +| A2 | `rot_768` | Random orthogonal rotation — is gain just reorientation? | +| A3 | `uc_256/64` | Unit-circle direct angle map — main hypothesis | +| A3b | `uc_256_rank` | Unit-circle rank angle map | +| A4 | *(in notebook)* | UC → inverse PCA back to Euclidean | +| A5 | `uc_256_broken` | Shuffled phase — negative control / theory destroyer | + +--- + +## Unit-circle transform variants + +**T1 — direct angle map** (`mode="direct"`) +``` +theta_i = pi * x_i +output_i = [cos(theta_i), sin(theta_i)] +``` + +**T2 — rank angle map** (`mode="rank"`) +``` +theta_i = 2*pi * rank(x_i) / (D-1) +output_i = [|x_i|*cos(theta_i), |x_i|*sin(theta_i)] +``` + +--- + +## Metrics + +| Category | Metrics | +|-------------|-----------------------------------------------| +| Retrieval | Recall@1/5/10, MRR, nDCG@10 | +| Clustering | Silhouette, NMI, ARI | +| Geometry | Neighbor overlap@k, Spearman rank correlation, trustworthiness/continuity penalties | + +--- + +## Quick start + +```bash +pip install -r requirements.txt +python -m src.main # Phase 1: nomic +MODEL_PRESET=bge-m3 python -m src.main # Phase 2: BGE-M3 +``` + +Results are written to `outputs/metrics/`. + +--- + +## Project layout + +``` +unit-circle-benchmark/ +├── src/ +│ ├── main.py — runner (ablation loop) +│ ├── config.py — model presets, transform configs +│ ├── data.py — JSONL loaders + validators +│ ├── embed.py — SentenceTransformer wrapper +│ ├── transforms.py — all A0–A5 transforms +│ ├── benchmark.py — retrieval, clustering, geometry suites +│ ├── metrics.py — table formatting, delta, success checks +│ └── utils.py — JSON I/O, dir helpers +├── data/ +│ ├── corpus.jsonl — {id, text, label} +│ ├── queries.jsonl — {id, text} +│ └── qrels.jsonl — {query_id, doc_id, relevance} +├── outputs/ +│ ├── metrics/ — JSON result files +│ ├── runs/ — per-run artefacts +│ └── plots/ — visualisations from notebook +└── notebooks/ + └── sanity_checks.ipynb +``` + +--- + +## Success criteria (spec) + +Call it promising only if: +- gain appears on **at least two corpora or two models** +- gain **survives against PCA control** +- gain is **not confined to a single metric** +- **negative control breaks** as expected (`uc_256_broken` < `native`) + +A single bump on one dataset is noise until proven otherwise. diff --git a/unit-circle-benchmark/data/corpus.jsonl b/unit-circle-benchmark/data/corpus.jsonl new file mode 100644 index 0000000..84a0e5b --- /dev/null +++ b/unit-circle-benchmark/data/corpus.jsonl @@ -0,0 +1,20 @@ +{"id":"d1","text":"The mitochondria are the powerhouse of the cell, responsible for generating most of the cell's supply of ATP through oxidative phosphorylation.","label":"biology"} +{"id":"d2","text":"Photosynthesis is the process by which green plants and algae convert sunlight, water, and carbon dioxide into glucose and oxygen.","label":"biology"} +{"id":"d3","text":"DNA replication occurs during the S phase of the cell cycle and involves unwinding the double helix and synthesising a complementary strand.","label":"biology"} +{"id":"d4","text":"The Krebs cycle, also known as the citric acid cycle, is a series of chemical reactions used to generate energy through the oxidation of acetyl-CoA.","label":"biology"} +{"id":"d5","text":"Neurons transmit signals through a combination of electrical impulses and chemical neurotransmitters released across synaptic gaps.","label":"biology"} +{"id":"d6","text":"Quantum entanglement is a phenomenon where two particles remain connected such that the state of one instantly influences the other regardless of distance.","label":"physics"} +{"id":"d7","text":"The Higgs boson, discovered at CERN in 2012, is a fundamental particle that gives other particles mass through interaction with the Higgs field.","label":"physics"} +{"id":"d8","text":"General relativity describes gravity as the curvature of spacetime caused by mass and energy, generalising Newton's law of universal gravitation.","label":"physics"} +{"id":"d9","text":"Thermodynamics describes the relationships between heat, work, temperature, and energy, underpinned by the four laws of thermodynamics.","label":"physics"} +{"id":"d10","text":"Wave-particle duality states that every quantum entity exhibits both wave and particle properties depending on how it is observed or measured.","label":"physics"} +{"id":"d11","text":"Gradient descent is an optimisation algorithm that iteratively adjusts model parameters in the direction that minimises a loss function.","label":"ml"} +{"id":"d12","text":"Transformer architectures use self-attention mechanisms to model relationships between all tokens in a sequence simultaneously.","label":"ml"} +{"id":"d13","text":"Overfitting occurs when a model learns the training data too well, including noise, leading to poor generalisation on unseen examples.","label":"ml"} +{"id":"d14","text":"Convolutional neural networks apply learnable filters across input data, making them effective for image recognition and spatial feature extraction.","label":"ml"} +{"id":"d15","text":"Reinforcement learning trains agents to maximise cumulative reward by taking actions in an environment and learning from feedback signals.","label":"ml"} +{"id":"d16","text":"The Roman Empire at its height controlled territories stretching from Britain in the northwest to Mesopotamia in the east.","label":"history"} +{"id":"d17","text":"The Industrial Revolution, beginning in Britain in the late 18th century, transformed manufacturing through mechanisation and the use of steam power.","label":"history"} +{"id":"d18","text":"The French Revolution of 1789 led to the abolition of the monarchy, the rise of Napoleon, and the spread of republican ideals across Europe.","label":"history"} +{"id":"d19","text":"The Silk Road was an ancient network of trade routes connecting China and East Asia with Central Asia, the Middle East, and Europe.","label":"history"} +{"id":"d20","text":"The Renaissance was a cultural and intellectual movement in Europe from the 14th to 17th centuries, reviving interest in classical Greek and Roman thought.","label":"history"} diff --git a/unit-circle-benchmark/data/qrels.jsonl b/unit-circle-benchmark/data/qrels.jsonl new file mode 100644 index 0000000..a696bd8 --- /dev/null +++ b/unit-circle-benchmark/data/qrels.jsonl @@ -0,0 +1,12 @@ +{"query_id":"q1","doc_id":"d1","relevance":1} +{"query_id":"q1","doc_id":"d4","relevance":1} +{"query_id":"q2","doc_id":"d7","relevance":1} +{"query_id":"q3","doc_id":"d13","relevance":1} +{"query_id":"q3","doc_id":"d11","relevance":1} +{"query_id":"q4","doc_id":"d18","relevance":1} +{"query_id":"q5","doc_id":"d6","relevance":1} +{"query_id":"q6","doc_id":"d3","relevance":1} +{"query_id":"q7","doc_id":"d11","relevance":1} +{"query_id":"q8","doc_id":"d19","relevance":1} +{"query_id":"q9","doc_id":"d14","relevance":1} +{"query_id":"q10","doc_id":"d17","relevance":1} diff --git a/unit-circle-benchmark/data/queries.jsonl b/unit-circle-benchmark/data/queries.jsonl new file mode 100644 index 0000000..1b59b92 --- /dev/null +++ b/unit-circle-benchmark/data/queries.jsonl @@ -0,0 +1,10 @@ +{"id":"q1","text":"How do cells produce energy?"} +{"id":"q2","text":"What is the role of the Higgs boson?"} +{"id":"q3","text":"How do neural networks avoid overfitting?"} +{"id":"q4","text":"What caused the French Revolution?"} +{"id":"q5","text":"How does quantum entanglement work?"} +{"id":"q6","text":"What is the process of DNA copying?"} +{"id":"q7","text":"How does gradient descent optimise a model?"} +{"id":"q8","text":"What trade routes connected ancient East and West?"} +{"id":"q9","text":"How do convolutional networks process images?"} +{"id":"q10","text":"What were the effects of industrialisation in Britain?"} diff --git a/unit-circle-benchmark/notebooks/sanity_checks.ipynb b/unit-circle-benchmark/notebooks/sanity_checks.ipynb new file mode 100644 index 0000000..124b74a --- /dev/null +++ b/unit-circle-benchmark/notebooks/sanity_checks.ipynb @@ -0,0 +1,185 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": ["# Unit-circle embedding benchmark — sanity checks\n", "\n", "Run this notebook after `python -m src.main` to visualise results and inspect the transform geometry."] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys, os\n", + "sys.path.insert(0, os.path.abspath(\"..\"))\n", + "\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from src.utils import load_json\n", + "from src.metrics import format_results_table, delta_vs_baseline" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": ["## 1. Load results"] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "runs = load_json(\"../outputs/metrics/nomic_results.json\")\n", + "print(format_results_table(runs))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": ["## 2. Deltas vs native baseline"] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "deltas = delta_vs_baseline(runs, baseline=\"native\")\n", + "for name, d in deltas.items():\n", + " print(f\"\\n--- {name} ---\")\n", + " for k, v in sorted(d.items()):\n", + " sign = \"+\" if v >= 0 else \"\"\n", + " print(f\" {k:<40} {sign}{v:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": ["## 3. MRR bar chart: all transforms"] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "names = list(runs.keys())\n", + "mrrs = [runs[n][\"retrieval\"][\"mrr\"] for n in names]\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 4))\n", + "bars = ax.bar(names, mrrs, color=[\"steelblue\" if \"uc\" in n else \"grey\" for n in names])\n", + "ax.axhline(runs[\"native\"][\"retrieval\"][\"mrr\"], color=\"red\", linestyle=\"--\", label=\"native\")\n", + "ax.set_ylabel(\"MRR\")\n", + "ax.set_title(\"MRR by transform (blue = unit-circle)\")\n", + "ax.legend()\n", + "plt.xticks(rotation=30, ha=\"right\")\n", + "plt.tight_layout()\n", + "os.makedirs(\"../outputs/plots\", exist_ok=True)\n", + "plt.savefig(\"../outputs/plots/mrr_bar.png\", dpi=150)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": ["## 4. Neighbor overlap: native vs uc_256 vs pca_256"] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "keys_of_interest = [\"native\", \"pca_256\", \"uc_256\", \"pca_64\", \"uc_64\"]\n", + "overlap_key = [k for k in runs[\"native\"][\"geometry\"] if \"overlap\" in k][0]\n", + "\n", + "labels = [k for k in keys_of_interest if k in runs]\n", + "overlaps = [runs[k][\"geometry\"][overlap_key] for k in labels]\n", + "\n", + "fig, ax = plt.subplots(figsize=(7, 4))\n", + "ax.bar(labels, overlaps, color=[\"steelblue\" if \"uc\" in l else \"grey\" for l in labels])\n", + "ax.set_ylabel(overlap_key)\n", + "ax.set_title(\"Neighbor overlap vs native space\")\n", + "plt.tight_layout()\n", + "plt.savefig(\"../outputs/plots/neighbor_overlap.png\", dpi=150)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": ["## 5. Direct geometry check: T1 transform on a toy vector"] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from src.transforms import _uc_map, l2_normalize\n", + "\n", + "# Construct a simple 4-dim test vector\n", + "x = np.array([[1.0, -1.0, 0.5, -0.5]], dtype=np.float32)\n", + "x_norm = l2_normalize(x)\n", + "mapped = _uc_map(x_norm, out_dim=4, mode=\"direct\")\n", + "\n", + "print(\"input (normalised):\", x_norm)\n", + "print(\"UC-mapped output :\", mapped)\n", + "print(\"\\ncos half :\", np.cos(np.pi * x_norm))\n", + "print(\"sin half :\", np.sin(np.pi * x_norm))\n", + "\n", + "# Visualise on unit circle\n", + "thetas = np.pi * x_norm[0]\n", + "fig, ax = plt.subplots(figsize=(5, 5))\n", + "circle = plt.Circle((0, 0), 1, fill=False, color=\"lightgrey\")\n", + "ax.add_patch(circle)\n", + "for i, theta in enumerate(thetas):\n", + " ax.plot([0, np.cos(theta)], [0, np.sin(theta)], label=f\"x_{i}={x_norm[0,i]:.2f}\")\n", + " ax.scatter(np.cos(theta), np.sin(theta), s=80, zorder=5)\n", + "ax.set_xlim(-1.3, 1.3)\n", + "ax.set_ylim(-1.3, 1.3)\n", + "ax.set_aspect(\"equal\")\n", + "ax.legend(fontsize=8)\n", + "ax.set_title(\"T1 direct angle map\")\n", + "plt.tight_layout()\n", + "plt.savefig(\"../outputs/plots/uc_geometry.png\", dpi=150)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": ["## 6. Success criteria summary"] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from src.metrics import check_success_criteria\n", + "\n", + "checks = check_success_criteria(runs, uc_key=\"uc_256\", pca_key=\"pca_256\")\n", + "print(\"Success criteria:\")\n", + "for criterion, passed in checks.items():\n", + " mark = \"PASS\" if passed else \"FAIL\"\n", + " print(f\" [{mark}] {criterion}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/unit-circle-benchmark/pyproject.toml b/unit-circle-benchmark/pyproject.toml new file mode 100644 index 0000000..b7c743e --- /dev/null +++ b/unit-circle-benchmark/pyproject.toml @@ -0,0 +1,28 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.backends.legacy:build" + +[project] +name = "unit-circle-benchmark" +version = "0.1.0" +description = "Benchmark: does a unit-circle coordinate remapping improve embedding usefulness?" +requires-python = ">=3.10" +dependencies = [ + "numpy", + "pandas", + "scikit-learn", + "scipy", + "sentence-transformers", + "transformers", + "torch", + "tqdm", + "orjson", + "matplotlib", +] + +[project.optional-dependencies] +openai = ["openai"] + +[tool.setuptools.packages.find] +where = ["."] +include = ["src*"] diff --git a/unit-circle-benchmark/requirements.txt b/unit-circle-benchmark/requirements.txt new file mode 100644 index 0000000..adc87e6 --- /dev/null +++ b/unit-circle-benchmark/requirements.txt @@ -0,0 +1,12 @@ +numpy +pandas +scikit-learn +scipy +sentence-transformers +transformers +torch +tqdm +orjson +matplotlib +# Optional — uncomment for OpenAI ceiling-check (Phase 3) +# openai diff --git a/unit-circle-benchmark/src/__init__.py b/unit-circle-benchmark/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/unit-circle-benchmark/src/benchmark.py b/unit-circle-benchmark/src/benchmark.py new file mode 100644 index 0000000..6ba5a46 --- /dev/null +++ b/unit-circle-benchmark/src/benchmark.py @@ -0,0 +1,153 @@ +# src/benchmark.py +import numpy as np +from sklearn.metrics import silhouette_score, normalized_mutual_info_score, adjusted_rand_score +from sklearn.cluster import KMeans +from scipy.stats import spearmanr +from typing import List, Dict, Any, Tuple, Optional + + +# --------------------------------------------------------------------------- +# Retrieval +# --------------------------------------------------------------------------- + +def run_retrieval( + query_vecs: np.ndarray, + doc_vecs: np.ndarray, + queries: List[Dict[str, Any]], + corpus: List[Dict[str, Any]], + qrels: List[Dict[str, Any]], + k_values: Tuple[int, ...] = (1, 5, 10), +) -> Dict[str, float]: + qrel_map: Dict[str, set] = {} + for row in qrels: + qrel_map.setdefault(row["query_id"], set()).add(row["doc_id"]) + + doc_ids = [d["id"] for d in corpus] + # cosine similarity (vectors are L2-normalised going in) + sims = query_vecs @ doc_vecs.T + + hits = {k: [] for k in k_values} + reciprocal_ranks = [] + ndcg_scores = {k: [] for k in k_values} + + for i, q in enumerate(queries): + ranked_indices = np.argsort(-sims[i]) + ranked_doc_ids = [doc_ids[j] for j in ranked_indices] + gold = qrel_map.get(q["id"], set()) + + # MRR + first_rr = 0.0 + for rank_idx, doc_id in enumerate(ranked_doc_ids, start=1): + if doc_id in gold: + first_rr = 1.0 / rank_idx + break + reciprocal_ranks.append(first_rr) + + # Recall@k + nDCG@k + for k in k_values: + topk_ids = ranked_doc_ids[:k] + hits[k].append(1.0 if gold & set(topk_ids) else 0.0) + ndcg_scores[k].append(_ndcg(gold, topk_ids, k)) + + result = {f"recall@{k}": float(np.mean(v)) for k, v in hits.items()} + result["mrr"] = float(np.mean(reciprocal_ranks)) + result.update({f"ndcg@{k}": float(np.mean(v)) for k, v in ndcg_scores.items()}) + return result + + +def _ndcg(gold: set, ranked: List[str], k: int) -> float: + dcg = sum( + 1.0 / np.log2(i + 2) + for i, doc_id in enumerate(ranked[:k]) + if doc_id in gold + ) + ideal = sum(1.0 / np.log2(i + 2) for i in range(min(len(gold), k))) + return float(dcg / ideal) if ideal > 0 else 0.0 + + +# --------------------------------------------------------------------------- +# Clustering +# --------------------------------------------------------------------------- + +def run_clustering( + doc_vecs: np.ndarray, + labels: Optional[List[Any]], +) -> Dict[str, float]: + if labels is None or any(x is None for x in labels): + return {} + + uniq = sorted(set(labels)) + if len(uniq) < 2: + return {} + + label_to_int = {x: i for i, x in enumerate(uniq)} + y = np.array([label_to_int[x] for x in labels]) + + km = KMeans(n_clusters=len(uniq), n_init=10, random_state=42) + pred = km.fit_predict(doc_vecs) + + return { + "silhouette": float(silhouette_score(doc_vecs, pred, sample_size=min(len(doc_vecs), 5000))), + "nmi": float(normalized_mutual_info_score(y, pred)), + "ari": float(adjusted_rand_score(y, pred)), + } + + +# --------------------------------------------------------------------------- +# Geometry / neighborhood preservation +# --------------------------------------------------------------------------- + +def run_geometry_suite( + native_vecs: np.ndarray, + transformed_vecs: np.ndarray, + k: int = 10, +) -> Dict[str, float]: + n = native_vecs.shape[0] + # Use cosine similarity (vectors should already be L2-normalised) + native_sims = native_vecs @ native_vecs.T + trans_sims = transformed_vecs @ transformed_vecs.T + + overlaps = [] + rank_corrs = [] + trustworthy = [] + continuous = [] + + for i in range(n): + native_order = np.argsort(-native_sims[i]) + trans_order = np.argsort(-trans_sims[i]) + + # exclude self (index 0 after sort) + n_nbrs = set(native_order[1: k + 1]) + t_nbrs = set(trans_order[1: k + 1]) + + overlaps.append(len(n_nbrs & t_nbrs) / k) + + # Spearman on the transformed similarities for the native top-k neighbours + n_nbr_list = list(n_nbrs) + n_vals = native_sims[i][n_nbr_list] + t_vals = trans_sims[i][n_nbr_list] + corr = spearmanr(n_vals, t_vals).statistic + rank_corrs.append(0.0 if (corr is None or np.isnan(corr)) else float(corr)) + + # Trustworthiness contribution: penalise items in t_nbrs but not n_nbrs + extra_in_t = t_nbrs - n_nbrs + for j in extra_in_t: + rank_in_native = int(np.where(native_order == j)[0][0]) + trustworthy.append(rank_in_native - k) + + # Continuity contribution: penalise items in n_nbrs but not t_nbrs + missing_in_t = n_nbrs - t_nbrs + for j in missing_in_t: + rank_in_trans = int(np.where(trans_order == j)[0][0]) + continuous.append(rank_in_trans - k) + + # Normalise trustworthiness / continuity to [0,1] range (simplified) + t_penalty = float(np.mean(trustworthy)) if trustworthy else 0.0 + c_penalty = float(np.mean(continuous)) if continuous else 0.0 + + return { + f"neighbor_overlap@{k}": float(np.mean(overlaps)), + "spearman_on_native_neighbors": float(np.mean(rank_corrs)), + "trustworthiness_penalty": t_penalty, + "continuity_penalty": c_penalty, + } diff --git a/unit-circle-benchmark/src/config.py b/unit-circle-benchmark/src/config.py new file mode 100644 index 0000000..77db3dc --- /dev/null +++ b/unit-circle-benchmark/src/config.py @@ -0,0 +1,57 @@ +# src/config.py +from dataclasses import dataclass, field +from typing import List, Optional + + +@dataclass +class EmbedConfig: + model_name: str = "nomic-ai/nomic-embed-text-v1.5" + doc_prefix: Optional[str] = "search_document: " + query_prefix: Optional[str] = "search_query: " + batch_size: int = 64 + + +# Named presets for the models tested in the benchmark +EMBED_PRESETS = { + "nomic": EmbedConfig( + model_name="nomic-ai/nomic-embed-text-v1.5", + doc_prefix="search_document: ", + query_prefix="search_query: ", + ), + "bge-m3": EmbedConfig( + model_name="BAAI/bge-m3", + doc_prefix=None, + query_prefix=None, + ), + # text-embedding-3-large requires the openai package; handled separately +} + + +@dataclass +class TransformConfig: + name: str + out_dim: int + mode: Optional[str] = None # used by unit_circle_transform + + +# Ablation ladder as described in the spec +ABLATION_LADDER: List[TransformConfig] = [ + TransformConfig("native", 768, None), + TransformConfig("pca_256", 256, None), + TransformConfig("uc_256", 256, "direct"), + TransformConfig("pca_64", 64, None), + TransformConfig("uc_64", 64, "direct"), + TransformConfig("uc_256_rank", 256, "rank"), + TransformConfig("uc_64_rank", 64, "rank"), +] + + +@dataclass +class BenchmarkConfig: + corpus_path: str = "data/corpus.jsonl" + queries_path: str = "data/queries.jsonl" + qrels_path: str = "data/qrels.jsonl" + output_dir: str = "outputs" + k_values: List[int] = field(default_factory=lambda: [1, 5, 10]) + geometry_k: int = 10 + embed: EmbedConfig = field(default_factory=EmbedConfig) diff --git a/unit-circle-benchmark/src/data.py b/unit-circle-benchmark/src/data.py new file mode 100644 index 0000000..2269999 --- /dev/null +++ b/unit-circle-benchmark/src/data.py @@ -0,0 +1,51 @@ +# src/data.py +import orjson +from pathlib import Path +from typing import List, Dict, Any + + +def _load_jsonl(path: str) -> List[Dict[str, Any]]: + records = [] + with open(path, "rb") as fh: + for line in fh: + line = line.strip() + if line: + records.append(orjson.loads(line)) + return records + + +def load_corpus(path: str = "data/corpus.jsonl") -> List[Dict[str, Any]]: + """Each record must have at minimum 'id' and 'text'. Optional: 'label'.""" + return _load_jsonl(path) + + +def load_queries(path: str = "data/queries.jsonl") -> List[Dict[str, Any]]: + """Each record must have 'id' and 'text'.""" + return _load_jsonl(path) + + +def load_qrels(path: str = "data/qrels.jsonl") -> List[Dict[str, Any]]: + """Each record must have 'query_id', 'doc_id', and 'relevance'.""" + return _load_jsonl(path) + + +def validate_corpus(corpus: List[Dict[str, Any]]) -> None: + for i, doc in enumerate(corpus): + if "id" not in doc or "text" not in doc: + raise ValueError(f"corpus record {i} missing 'id' or 'text': {doc}") + + +def validate_queries(queries: List[Dict[str, Any]]) -> None: + for i, q in enumerate(queries): + if "id" not in q or "text" not in q: + raise ValueError(f"query record {i} missing 'id' or 'text': {q}") + + +def validate_qrels(qrels: List[Dict[str, Any]], corpus: List[Dict[str, Any]], queries: List[Dict[str, Any]]) -> None: + doc_ids = {d["id"] for d in corpus} + query_ids = {q["id"] for q in queries} + for i, rel in enumerate(qrels): + if rel["query_id"] not in query_ids: + raise ValueError(f"qrel {i}: unknown query_id '{rel['query_id']}'") + if rel["doc_id"] not in doc_ids: + raise ValueError(f"qrel {i}: unknown doc_id '{rel['doc_id']}'") diff --git a/unit-circle-benchmark/src/embed.py b/unit-circle-benchmark/src/embed.py new file mode 100644 index 0000000..d455d3c --- /dev/null +++ b/unit-circle-benchmark/src/embed.py @@ -0,0 +1,47 @@ +# src/embed.py +import numpy as np +from typing import List, Optional +from sentence_transformers import SentenceTransformer + +_MODEL_CACHE = {} + + +def _get_model(model_name: str) -> SentenceTransformer: + if model_name not in _MODEL_CACHE: + _MODEL_CACHE[model_name] = SentenceTransformer(model_name, trust_remote_code=True) + return _MODEL_CACHE[model_name] + + +def embed_texts( + texts: List[str], + model_name: str = "nomic-ai/nomic-embed-text-v1.5", + prefix: Optional[str] = None, + batch_size: int = 64, +) -> np.ndarray: + """ + Encode texts with an optional task-prefix (e.g. 'search_query: '). + Returns L2-normalised float32 array of shape (N, D). + """ + model = _get_model(model_name) + inputs = texts if prefix is None else [f"{prefix}{t}" for t in texts] + vecs = model.encode( + inputs, + batch_size=batch_size, + normalize_embeddings=True, + show_progress_bar=True, + ) + return np.asarray(vecs, dtype=np.float32) + + +def embed_corpus_and_queries( + corpus_texts: List[str], + query_texts: List[str], + model_name: str = "nomic-ai/nomic-embed-text-v1.5", + doc_prefix: Optional[str] = "search_document: ", + query_prefix: Optional[str] = "search_query: ", + batch_size: int = 64, +): + """Convenience wrapper that respects Nomic-style task prefixes.""" + doc_vecs = embed_texts(corpus_texts, model_name, prefix=doc_prefix, batch_size=batch_size) + q_vecs = embed_texts(query_texts, model_name, prefix=query_prefix, batch_size=batch_size) + return doc_vecs, q_vecs diff --git a/unit-circle-benchmark/src/main.py b/unit-circle-benchmark/src/main.py new file mode 100644 index 0000000..d176458 --- /dev/null +++ b/unit-circle-benchmark/src/main.py @@ -0,0 +1,138 @@ +# src/main.py +""" +Unit-circle embedding benchmark — main runner. + +Phase 1 ablation ladder on nomic-embed-text-v1.5: + A0 native + A1 pca_256 / pca_64 + A3 uc_256 / uc_64 (direct angle map) + A3b uc_256_rank / uc_64_rank (rank angle map) + +Set MODEL_PRESET env var to switch to "bge-m3": + MODEL_PRESET=bge-m3 python -m src.main +""" + +import os +import numpy as np + +from src.config import EMBED_PRESETS, BenchmarkConfig +from src.data import load_corpus, load_queries, load_qrels, validate_corpus, validate_queries, validate_qrels +from src.embed import embed_corpus_and_queries +from src.transforms import ( + identity_transform, + pca_transform, + unit_circle_transform, + random_rotation_transform, + uc_broken_phase_transform, +) +from src.benchmark import run_retrieval, run_clustering, run_geometry_suite +from src.metrics import format_results_table, delta_vs_baseline, check_success_criteria +from src.utils import save_json, ensure_dirs, pretty_print_results + + +def main(): + preset_name = os.environ.get("MODEL_PRESET", "nomic") + cfg = BenchmarkConfig() + embed_cfg = EMBED_PRESETS.get(preset_name) + if embed_cfg is None: + raise ValueError(f"Unknown MODEL_PRESET '{preset_name}'. Choose from: {list(EMBED_PRESETS)}") + + print(f"[benchmark] model preset : {preset_name}") + print(f"[benchmark] model : {embed_cfg.model_name}") + + # ------------------------------------------------------------------ + # Load data + # ------------------------------------------------------------------ + corpus = load_corpus(cfg.corpus_path) + queries = load_queries(cfg.queries_path) + qrels = load_qrels(cfg.qrels_path) + + validate_corpus(corpus) + validate_queries(queries) + validate_qrels(qrels, corpus, queries) + + print(f"[benchmark] corpus : {len(corpus)} docs") + print(f"[benchmark] queries : {len(queries)}") + print(f"[benchmark] qrels : {len(qrels)}") + + doc_texts = [x["text"] for x in corpus] + query_texts = [x["text"] for x in queries] + labels = [x.get("label") for x in corpus] + + # ------------------------------------------------------------------ + # Embed (native, once) + # ------------------------------------------------------------------ + print("[benchmark] embedding corpus and queries …") + docs_native, queries_native = embed_corpus_and_queries( + doc_texts, + query_texts, + model_name=embed_cfg.model_name, + doc_prefix=embed_cfg.doc_prefix, + query_prefix=embed_cfg.query_prefix, + batch_size=embed_cfg.batch_size, + ) + print(f"[benchmark] native shape : {docs_native.shape}") + + # ------------------------------------------------------------------ + # Transform registry + # ------------------------------------------------------------------ + transforms = { + # A0 + "native": lambda d, q: (identity_transform(d), identity_transform(q)), + # A1 + "pca_256": lambda d, q: pca_transform(d, q, out_dim=256), + "pca_64": lambda d, q: pca_transform(d, q, out_dim=64), + # A2 + "rot_768": lambda d, q: random_rotation_transform(d, q), + # A3 + "uc_256": lambda d, q: unit_circle_transform(d, q, out_dim=256, mode="direct"), + "uc_64": lambda d, q: unit_circle_transform(d, q, out_dim=64, mode="direct"), + "uc_256_rank": lambda d, q: unit_circle_transform(d, q, out_dim=256, mode="rank"), + "uc_64_rank": lambda d, q: unit_circle_transform(d, q, out_dim=64, mode="rank"), + # A5 — negative control + "uc_256_broken": lambda d, q: uc_broken_phase_transform(d, q, out_dim=256, mode="direct"), + } + + # ------------------------------------------------------------------ + # Run ablations + # ------------------------------------------------------------------ + runs = {} + for name, fn in transforms.items(): + print(f"[benchmark] running transform: {name}") + d_vecs, q_vecs = fn(docs_native, queries_native) + + retrieval = run_retrieval(q_vecs, d_vecs, queries, corpus, qrels, k_values=cfg.k_values) + clustering = run_clustering(d_vecs, labels) + geometry = run_geometry_suite(docs_native, d_vecs, k=cfg.geometry_k) + + runs[name] = { + "retrieval": retrieval, + "clustering": clustering, + "geometry": geometry, + } + print(f" mrr={retrieval['mrr']:.4f} recall@10={retrieval.get('recall@10', 0):.4f}") + + # ------------------------------------------------------------------ + # Persist and report + # ------------------------------------------------------------------ + ensure_dirs(cfg.output_dir, f"{cfg.output_dir}/metrics", f"{cfg.output_dir}/runs") + out_path = f"{cfg.output_dir}/metrics/{preset_name}_results.json" + save_json(out_path, runs) + print(f"\n[benchmark] results saved → {out_path}\n") + + pretty_print_results(runs) + + deltas = delta_vs_baseline(runs, baseline="native") + save_json(f"{cfg.output_dir}/metrics/{preset_name}_deltas.json", deltas) + + checks = check_success_criteria(runs, uc_key="uc_256", pca_key="pca_256") + print("\n[benchmark] success criteria:") + for criterion, passed in checks.items(): + mark = "PASS" if passed else "FAIL" + print(f" {mark} {criterion}") + + save_json(f"{cfg.output_dir}/metrics/{preset_name}_success_checks.json", checks) + + +if __name__ == "__main__": + main() diff --git a/unit-circle-benchmark/src/metrics.py b/unit-circle-benchmark/src/metrics.py new file mode 100644 index 0000000..7024568 --- /dev/null +++ b/unit-circle-benchmark/src/metrics.py @@ -0,0 +1,98 @@ +# src/metrics.py +""" +Formatting and comparison helpers for benchmark results. +""" +import numpy as np +from typing import Dict, Any + + +def format_results_table(runs: Dict[str, Dict[str, Any]]) -> str: + """Return a plain-text table of all run results for quick inspection.""" + all_keys = sorted({ + f"{task}/{metric}" + for run in runs.values() + for task, scores in run.items() + for metric in (scores.keys() if isinstance(scores, dict) else []) + }) + + col_width = max(20, max((len(k) for k in all_keys), default=20)) + run_names = list(runs.keys()) + header = f"{'metric':<{col_width}}" + "".join(f" {n:>12}" for n in run_names) + sep = "-" * len(header) + rows = [header, sep] + + for key in all_keys: + task, metric = key.split("/", 1) + row = f"{key:<{col_width}}" + for name in run_names: + val = runs[name].get(task, {}).get(metric) + if val is None: + row += f" {'—':>12}" + else: + row += f" {val:>12.4f}" + rows.append(row) + + return "\n".join(rows) + + +def delta_vs_baseline( + runs: Dict[str, Dict[str, Any]], + baseline: str = "native", +) -> Dict[str, Dict[str, float]]: + """ + Return delta of each run relative to the baseline run. + Positive = better than baseline. + """ + base = runs.get(baseline, {}) + deltas = {} + for name, result in runs.items(): + if name == baseline: + continue + d = {} + for task, scores in result.items(): + if not isinstance(scores, dict): + continue + for metric, val in scores.items(): + base_val = base.get(task, {}).get(metric) + if base_val is not None: + d[f"{task}/{metric}"] = float(val) - float(base_val) + deltas[name] = d + return deltas + + +def check_success_criteria( + runs: Dict[str, Dict[str, Any]], + uc_key: str = "uc_256", + pca_key: str = "pca_256", +) -> Dict[str, bool]: + """ + Evaluate the spec's success criteria. + Returns a dict of named checks and whether they pass. + """ + def _get(run_name, task, metric): + return runs.get(run_name, {}).get(task, {}).get(metric) + + checks = {} + + # 1. unit-circle beats PCA on retrieval + uc_mrr = _get(uc_key, "retrieval", "mrr") + pca_mrr = _get(pca_key, "retrieval", "mrr") + checks["uc_beats_pca_retrieval_mrr"] = ( + uc_mrr is not None and pca_mrr is not None and uc_mrr > pca_mrr + ) + + # 2. unit-circle beats PCA on clustering + uc_sil = _get(uc_key, "clustering", "silhouette") + pca_sil = _get(pca_key, "clustering", "silhouette") + checks["uc_beats_pca_clustering_silhouette"] = ( + uc_sil is not None and pca_sil is not None and uc_sil > pca_sil + ) + + # 3. neighbor overlap does not collapse while retrieval improves + uc_overlap = _get(uc_key, "geometry", "neighbor_overlap@10") + native_overlap = _get("native", "geometry", "neighbor_overlap@10") + checks["neighbor_overlap_preserved"] = ( + uc_overlap is not None and native_overlap is not None and uc_overlap >= 0.5 * native_overlap + ) + + return checks diff --git a/unit-circle-benchmark/src/transforms.py b/unit-circle-benchmark/src/transforms.py new file mode 100644 index 0000000..e621a04 --- /dev/null +++ b/unit-circle-benchmark/src/transforms.py @@ -0,0 +1,151 @@ +# src/transforms.py +""" +Ablation ladder transforms. + +A0 identity_transform — native L2-normalised baseline +A1 pca_transform — PCA control (same target dim) +A2 random_rotation_transform — orthogonal rotation control +A3 unit_circle_transform — main hypothesis (modes: direct, rank) +A4 uc_then_inverse_transform — test whether gain survives back-projection +A5 uc_broken_phase_transform — negative control: destroy angular meaning +""" + +import numpy as np +from sklearn.decomposition import PCA +from typing import Tuple + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def l2_normalize(x: np.ndarray, eps: float = 1e-12) -> np.ndarray: + norms = np.linalg.norm(x, axis=1, keepdims=True) + return x / np.clip(norms, eps, None) + + +# --------------------------------------------------------------------------- +# A0 — native +# --------------------------------------------------------------------------- + +def identity_transform(x: np.ndarray) -> np.ndarray: + return l2_normalize(x) + + +# --------------------------------------------------------------------------- +# A1 — PCA control +# --------------------------------------------------------------------------- + +def pca_transform( + doc_vecs: np.ndarray, + query_vecs: np.ndarray, + out_dim: int = 256, +) -> Tuple[np.ndarray, np.ndarray]: + pca = PCA(n_components=out_dim, random_state=42) + d = pca.fit_transform(doc_vecs) + q = pca.transform(query_vecs) + return l2_normalize(d), l2_normalize(q) + + +# --------------------------------------------------------------------------- +# A2 — random orthogonal rotation +# --------------------------------------------------------------------------- + +def random_rotation_transform( + doc_vecs: np.ndarray, + query_vecs: np.ndarray, + seed: int = 42, +) -> Tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(seed) + dim = doc_vecs.shape[1] + # Gram-Schmidt on a random matrix gives a uniformly random orthogonal matrix + H = rng.standard_normal((dim, dim)) + Q, _ = np.linalg.qr(H) + d = doc_vecs @ Q + q = query_vecs @ Q + return l2_normalize(d), l2_normalize(q) + + +# --------------------------------------------------------------------------- +# A3 — unit-circle transform (main hypothesis) +# --------------------------------------------------------------------------- + +def unit_circle_transform( + doc_vecs: np.ndarray, + query_vecs: np.ndarray, + out_dim: int = 256, + mode: str = "direct", +) -> Tuple[np.ndarray, np.ndarray]: + d = _uc_map(doc_vecs, out_dim=out_dim, mode=mode) + q = _uc_map(query_vecs, out_dim=out_dim, mode=mode) + return l2_normalize(d), l2_normalize(q) + + +def _uc_map(x: np.ndarray, out_dim: int = 256, mode: str = "direct") -> np.ndarray: + # Truncate (or pad) to out_dim before mapping + if x.shape[1] >= out_dim: + x = x[:, :out_dim] + else: + pad = np.zeros((x.shape[0], out_dim - x.shape[1]), dtype=x.dtype) + x = np.concatenate([x, pad], axis=1) + + if mode == "direct": + # T1: theta_i = pi * x_i (x already in [-1,1] due to L2-norm) + theta = np.pi * np.clip(x, -1.0, 1.0) + z = np.concatenate([np.cos(theta), np.sin(theta)], axis=1) + + elif mode == "rank": + # T2: assign angles by coordinate rank, preserve magnitude separately + ranks = np.argsort(np.argsort(x, axis=1), axis=1).astype(np.float32) + theta = 2.0 * np.pi * ranks / max(1, x.shape[1] - 1) + mag = np.abs(x) + z = np.concatenate([mag * np.cos(theta), mag * np.sin(theta)], axis=1) + + else: + raise ValueError(f"Unknown unit-circle mode '{mode}'. Choose 'direct' or 'rank'.") + + return z.astype(np.float32) + + +# --------------------------------------------------------------------------- +# A4 — unit-circle then inverse-project back to Euclidean +# --------------------------------------------------------------------------- + +def uc_then_inverse_transform( + doc_vecs: np.ndarray, + query_vecs: np.ndarray, + out_dim: int = 256, + mode: str = "direct", +) -> Tuple[np.ndarray, np.ndarray]: + """Map to unit-circle space then project back via pseudo-inverse PCA.""" + d_uc, q_uc = unit_circle_transform(doc_vecs, query_vecs, out_dim=out_dim, mode=mode) + # Project back to original dimensionality via PCA fit on the UC-mapped docs + pca = PCA(n_components=doc_vecs.shape[1], random_state=42) + d = pca.fit_transform(d_uc) + q = pca.transform(q_uc) + return l2_normalize(d), l2_normalize(q) + + +# --------------------------------------------------------------------------- +# A5 — broken phase (negative control) +# --------------------------------------------------------------------------- + +def uc_broken_phase_transform( + doc_vecs: np.ndarray, + query_vecs: np.ndarray, + out_dim: int = 256, + mode: str = "direct", + seed: int = 42, +) -> Tuple[np.ndarray, np.ndarray]: + """Unit-circle map with shuffled angular assignments — destroys the theory.""" + rng = np.random.default_rng(seed) + + def _broken(x: np.ndarray) -> np.ndarray: + mapped = _uc_map(x, out_dim=out_dim, mode=mode) + # Shuffle the column assignment independently per sample + idx = rng.permutation(mapped.shape[1]) + return mapped[:, idx] + + d = _broken(doc_vecs) + q = _broken(query_vecs) + return l2_normalize(d), l2_normalize(q) diff --git a/unit-circle-benchmark/src/utils.py b/unit-circle-benchmark/src/utils.py new file mode 100644 index 0000000..a767ed1 --- /dev/null +++ b/unit-circle-benchmark/src/utils.py @@ -0,0 +1,25 @@ +# src/utils.py +import json +import orjson +from pathlib import Path +from typing import Any + + +def save_json(path: str, data: Any) -> None: + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(orjson.dumps(data, option=orjson.OPT_INDENT_2)) + + +def load_json(path: str) -> Any: + return orjson.loads(Path(path).read_bytes()) + + +def ensure_dirs(*paths: str) -> None: + for p in paths: + Path(p).mkdir(parents=True, exist_ok=True) + + +def pretty_print_results(runs: dict) -> None: + from src.metrics import format_results_table + print(format_results_table(runs))