embedder

  1#!/usr/bin/env python3
  2"""
  3embedder.py — swappable embedding backends for the feels vector search.
  4
  5The Atlas cosine search + material aggregation + margin guard don't care where
  6the vector comes from. This module puts one interface in front of the options:
  7
  8    FEELS_EMBEDDER=hand      10-D hand-engineered swirl features (default; current behavior)
  9    FEELS_EMBEDDER=pca       10-D PCA on raw 6-channel torque-delta trajectories,
 10                             fit on swirl-only train recordings. Unsupervised —
 11                             no hand tuning, no scales to calibrate.
 12    FEELS_EMBEDDER=chronos   2304-D zero-shot time-series foundation model
 13                             (amazon/chronos-t5-mini). Live-capable: Atlas picks
 14                             the vector index by embedding dim. Venue setup:
 15                               pip install torch chronos-forecasting
 16                               python -c "import atlas; atlas.ensure_vector_indexes('chronos')"
 17                               FEELS_EMBEDDER=chronos python teach_csv.py --force --require-swirl recordings/
 18                             then restart the dashboard with FEELS_EMBEDDER=chronos.
 19                             (First embed() call downloads the model from HuggingFace.)
 20
 21Contract: embed(df) -> list[float], L2-normalized. The pca backend reuses the
 22hand pipeline's preprocessing (INIT-median tare via _tare_with_qc, SWIRL window
 23else fallback dynamic phases) so the only difference is the representation.
 24
 25Switching embedders changes the vector SPACE, so the Atlas library must be
 26rebuilt under the same embedder or queries hit the wrong space:
 27
 28    FEELS_EMBEDDER=pca python teach_csv.py --force --require-swirl recordings/
 29
 30The live predict path refuses to query a library built with a different
 31embedder (explicit 409, not silent garbage). Taught signatures are tagged
 32with their embedder in properties.
 33
 34Offline bake-off on the committed train/test split (swirl-only):
 35
 36    python embedder.py --compare     # hand vs pca, atlas-style ranking
 37    python embedder.py --fit         # (re)fit the PCA cache from swirl train
 38
 392026-10-29 bake-off (swirl-only test, n=9, top-5 mean-cosine per material):
 40hand 22%, pca 33%, chronos 56%. Absolute numbers are modest for all three —
 41the 3.5 Hz signal is the bottleneck, not the features.
 42"""
 43
 44import argparse
 45import csv
 46import os
 47import re
 48import sys
 49
 50import numpy as np
 51
 52_HERE = os.path.dirname(os.path.abspath(__file__))
 53if _HERE not in sys.path:
 54    sys.path.insert(0, _HERE)
 55
 56EMBEDDER = os.getenv("FEELS_EMBEDDER", "hand").strip().lower()
 57if EMBEDDER not in ("hand", "pca", "chronos"):
 58    raise ValueError(f"FEELS_EMBEDDER={EMBEDDER!r}: choose hand|pca|chronos")
 59
 60DIM = {"hand": 10, "pca": 10, "chronos": 2304}
 61CACHE = os.path.join(_HERE, "pca_swirl.npz")
 62TAU_COLS = [f"tau{i}_nm" for i in range(1, 7)]
 63
 64_MATERIAL_STEMS = {
 65    "water": "water", "soap": "soap", "dish_soap": "soap",
 66    "candy": "candy", "solid": "solid", "screws": "screws",
 67    "screw": "screws", "empty": "empty", "empty_cup": "empty",
 68}
 69_CAPACITY_SUFFIX = re.compile(r"_(\d+(pct|g|ml|oz)|half|full|mixed|empty)$")
 70
 71
 72def material_of(label: str) -> str:
 73    stem = _CAPACITY_SUFFIX.sub("", label.strip().lower())
 74    return _MATERIAL_STEMS.get(stem, stem)
 75
 76
 77def _l2(v):
 78    v = np.asarray(v, dtype=float)
 79    n = np.linalg.norm(v)
 80    return v / n if n > 1e-12 else v
 81
 82
 83def _window_deltas(df):
 84    """(T, 6) torque deltas over the dynamic window, hand-pipeline preprocessing."""
 85    from feature_extractor import _tare_with_qc, SWIRL_PHASE, FALLBACK_DYN_PHASES
 86    tares, _ = _tare_with_qc(df)
 87    idx = df.index[df["phase"] == SWIRL_PHASE].values
 88    if len(idx) < 4:
 89        idx = df.index[df["phase"].isin(FALLBACK_DYN_PHASES)].values
 90    if len(idx) < 4:
 91        raise ValueError("no dynamic window")
 92    return np.stack([df[c].values[idx] - tares[i + 1] for i, c in enumerate(TAU_COLS)], axis=1)
 93
 94
 95def _resample(mat, T):
 96    n, ch = mat.shape
 97    if n == T:
 98        return mat
 99    x_old = np.linspace(0, 1, n)
100    x_new = np.linspace(0, 1, T)
101    return np.stack([np.interp(x_new, x_old, mat[:, c]) for c in range(ch)], axis=1)
102
103
104def _swirl_train_files():
105    manifest = os.path.join(_HERE, "recordings", "split_manifest.csv")
106    if not os.path.isfile(manifest):
107        raise RuntimeError("recordings/split_manifest.csv not found — can't fit PCA")
108    files = []
109    with open(manifest) as f:
110        for r in csv.DictReader(f):
111            if r["split"].strip() == "train" and r["motion_path"].strip() == "swirl":
112                files.append(os.path.join(_HERE, "recordings", r["file"].strip()))
113    if not files:
114        raise RuntimeError("no swirl train recordings in manifest")
115    return files
116
117
118def fit_pca(cache_path=CACHE):
119    """Fit PCA-10 on swirl-only train recordings; persist components to .npz."""
120    import pandas as pd
121    from sklearn.decomposition import PCA
122    mats = []
123    for p in _swirl_train_files():
124        try:
125            df = pd.read_csv(p)
126            mats.append(_window_deltas(df))
127        except Exception as e:
128            print(f"  skip {os.path.basename(p)}: {e}")
129    T = int(np.median([m.shape[0] for m in mats]))
130    X = np.stack([_resample(m, T).ravel() for m in mats])
131    pca = PCA(n_components=10, random_state=0).fit(X)
132    np.savez(cache_path, components=pca.components_, mean=pca.mean_,
133             T=np.array(T), n_train=np.array(len(mats)))
134    print(f"fit PCA on {len(mats)} swirl train recordings "
135          f"(T={T}, explained var {pca.explained_variance_ratio_.sum():.3f}) -> {cache_path}")
136    return cache_path
137
138
139def _load_pca():
140    if not os.path.isfile(CACHE):
141        print("PCA cache missing — fitting from swirl train recordings...")
142        fit_pca()
143    z = np.load(CACHE)
144    return z["components"], z["mean"], int(z["T"])
145
146
147_chronos_pipe = None
148
149
150def _chronos_embed(deltas):
151    global _chronos_pipe
152    try:
153        import torch
154        from chronos import ChronosPipeline
155    except ImportError:
156        raise RuntimeError(
157            "FEELS_EMBEDDER=chronos needs: pip install torch chronos-forecasting "
158            "(plus a HuggingFace model download on first use)")
159    if _chronos_pipe is None:
160        print("loading amazon/chronos-t5-mini ...", flush=True)
161        _chronos_pipe = ChronosPipeline.from_pretrained(
162            "amazon/chronos-t5-mini", device_map="cpu", torch_dtype=torch.float32)
163    ctx = torch.tensor(deltas.T, dtype=torch.float32)  # (6, T)
164    with torch.no_grad():
165        emb, _ = _chronos_pipe.embed(ctx)  # (6, T', d)
166    pooled = emb.float().mean(dim=1).numpy()  # mean over time -> (6, d)
167    return _l2(pooled.ravel())
168
169
170def embed(df):
171    """10-D (hand/pca) or 2304-D (chronos) L2-normalized embedding for one recording."""
172    if EMBEDDER == "hand":
173        from feature_extractor import extract_features
174        return [float(x) for x in extract_features(df)["e"]]
175    if EMBEDDER == "pca":
176        comp, mean, T = _load_pca()
177        d = _window_deltas(df)
178        v = (_resample(d, T).ravel()[None, :] - mean) @ comp.T
179        return [float(x) for x in _l2(v[0])]
180    if EMBEDDER == "chronos":
181        return [float(x) for x in _chronos_embed(_window_deltas(df))]
182    raise AssertionError("unreachable")
183
184
185def _compare():
186    """Offline bake-off: hand vs pca on the swirl-only test split, atlas-style ranking."""
187    import pandas as pd
188    manifest = os.path.join(_HERE, "recordings", "split_manifest.csv")
189    rows = []
190    with open(manifest) as f:
191        for r in csv.DictReader(f):
192            if r["motion_path"].strip() == "swirl":
193                rows.append(r)
194    recs = {}
195    for r in rows:
196        try:
197            df = pd.read_csv(os.path.join(_HERE, "recordings", r["file"].strip()))
198            from feature_extractor import extract_features
199            e_hand = _l2(np.array(extract_features(df)["e"]))
200            comp, mean, T = _load_pca()
201            d = _window_deltas(df)
202            e_pca = _l2(((_resample(d, T).ravel()[None, :] - mean) @ comp.T)[0])
203            recs[r["file"]] = {"material": material_of(r["label"]),
204                               "split": r["split"].strip(),
205                               "hand": e_hand, "pca": e_pca}
206        except Exception as e:
207            print("  skip", r["file"].strip(), str(e)[:80])
208    train = [v for v in recs.values() if v["split"] == "train"]
209    test = [v for v in recs.values() if v["split"] == "test"]
210    for name, key in [("hand", "hand"), ("pca", "pca")]:
211        lib = np.stack([v[key] for v in train])
212        mats = [v["material"] for v in train]
213        correct = 0
214        for v in test:
215            sims = lib @ v[key]
216            order = np.argsort(-sims)[:5]
217            scores = {}
218            for j in order:
219                scores.setdefault(mats[j], []).append(float(sims[j]))
220            ranked = sorted(scores.items(), key=lambda kv: sum(kv[1]) / len(kv[1]), reverse=True)
221            correct += (ranked[0][0] == v["material"])
222        print(f"{name}: {correct}/{len(test)} = {correct / len(test):.3f}  (swirl-only test)")
223
224
225def main():
226    ap = argparse.ArgumentParser(description="Swappable embedding backends for feels vector search.")
227    ap.add_argument("--fit", action="store_true", help="(re)fit the PCA cache from swirl train recordings")
228    ap.add_argument("--compare", action="store_true", help="offline bake-off: hand vs pca on swirl-only test")
229    args = ap.parse_args()
230    if args.fit:
231        fit_pca()
232    elif args.compare:
233        _compare()
234    else:
235        print(f"FEELS_EMBEDDER={EMBEDDER}  dim={DIM[EMBEDDER]}")
236
237
238if __name__ == "__main__":
239    main()