diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..123e1d4 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,25 @@ +name: ci + +on: + push: + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - uses: astral-sh/setup-uv@v5 + with: + python-version: "3.12" + + - name: Create venv and install base deps + run: uv sync --no-dev + + - name: Lint + run: uvx ruff check + + - name: Unit tests + run: uv run python -m unittest discover -s tests -v diff --git a/.gitignore b/.gitignore index 835aca9..7dab55b 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,7 @@ checkpoints/ *.bin *.gguf data/*.jsonl +data/raw/ !examples/*.jsonl .env diff --git a/README.md b/README.md index 60aeaeb..022f6e8 100644 --- a/README.md +++ b/README.md @@ -15,16 +15,43 @@ The frozen intent set and labeling rules live in [docs/taxonomy.md](docs/taxonom ## Layout +The pipeline runs **data → train → export → evaluate**. Two files are the spine +everything else imports: `data.py` (the `Intent` label set and JSONL loaders) +and `preprocess.py` (the redaction/normalization applied identically at training +and inference, mirrored in the Rust consumer). + ``` src/routelet/ - data.py label schema and train/test prep - train.py fine-tune a small model on the labels - serve.py FastAPI endpoint: text in, intent out - evaluate.py routelet vs the LLM baseline (accuracy, latency, cost) -examples/ - intents.sample.jsonl example label format + data.py Intent enum (the 5 labels) + JSONL load/split. Source of truth. + preprocess.py mask secrets/emails/digits. Same pass at train and inference. + + augment.py clean rows -> voice-like disfluent variants (data/augmented.jsonl) + checkdata.py dataset health: per-class counts, train/eval leakage + + train.py TF-IDF + logistic-regression baseline (the floor to beat) + train_setfit.py the shipped model: SetFit fine-tune of bge-small + LR head, + then temperature calibration -> models/setfit + setfit_baseline.py benchmark harness that trains+discards SetFit (not the ship path) + export_onnx.py models/setfit -> embedder.onnx + tokenizer.json + head.json (int8) + + evaluate.py the Claude LLM baseline (the accuracy ceiling), one call per row + teacher.py shared taxonomy prompt + schema + classify(), used by eval + ingest + serve.py optional FastAPI endpoint (stub; Aegis runs the ONNX in-process) + +Scripts/ runnable tools (not library code) + pull_samples.py pull redacted samples from the proxy's R2 -> one JSONL + ingest_samples.py label them (free Claude label or teacher), dedup -> data/collected.jsonl + refresh_haiku_baseline.py re-run the paid Haiku eval -> report/baselines.json + +report/ report.py scores the models on the holdout -> figures + metrics.json +evals/ frozen eval data (holdout.jsonl). Never trained on. +experiments/ throwaway prototypes, not part of the library (see its README) ``` +The data-collection loop closes back on itself: Aegis logs redacted samples -> +proxy -> R2, then `Scripts/pull_samples.py` + `ingest_samples.py` turn those into +new labeled training data. + ## Data format One JSON object per line: diff --git a/Scripts/ingest_samples.py b/Scripts/ingest_samples.py new file mode 100644 index 0000000..a61b8ac --- /dev/null +++ b/Scripts/ingest_samples.py @@ -0,0 +1,158 @@ +"""Turn pulled distillation samples into labeled training data. + +Reads the raw samples produced by ``pull_samples.py`` and resolves a single +training label for each: + + * If the sample carries ``claude_label`` (the Claude fallback fired on-device + that turn), use it. Free teacher signal, already paid for. + * Otherwise call the offline Claude teacher to label it (docs/distillation.md). + +Then it normalizes the text through ``preprocess`` (the same pass training and +inference use, so there is no skew), drops anything already in the training pool +or in the eval set (keeping the eval honest), and appends the survivors to +``data/collected.jsonl`` as ``{"text", "intent"}`` lines that ``data.load`` reads. + +Needs ANTHROPIC_API_KEY only when some samples lack a ``claude_label``. The +teacher model is claude-opus-4-8 by default; override with ROUTELET_TEACHER_MODEL. + +Usage: + uv run python Scripts/ingest_samples.py + uv run python Scripts/ingest_samples.py --dry-run + uv run python Scripts/ingest_samples.py --raw data/raw/today.jsonl --limit 200 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).parent.parent +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +from routelet.data import Intent, load, load_dir # noqa: E402 +from routelet.preprocess import preprocess # noqa: E402 + +DEFAULT_RAW = PROJECT_ROOT / "data" / "raw" / "samples.jsonl" +DEFAULT_OUT = PROJECT_ROOT / "data" / "collected.jsonl" +TRAIN_DIR = PROJECT_ROOT / "data" +EVAL_FILE = PROJECT_ROOT / "evals" / "holdout.jsonl" + + +def _norm(text: str) -> str: + """The dedup key: normalized text, case-folded and stripped. Mirrors the + leakage check in checkdata.py so 'same command' means the same thing here.""" + return preprocess(text).lower().strip() + + +def load_raw(path: Path) -> list[dict]: + rows: list[dict] = [] + with path.open() as f: + for n, line in enumerate(f, start=1): + if not line.strip(): + continue + try: + rows.append(json.loads(line)) + except json.JSONDecodeError: + print(f" skip {path}:{n}: bad json", file=sys.stderr) + return rows + + +def existing_keys() -> tuple[set[str], set[str]]: + """Texts already in the training pool and in the eval set, both normalized. + A sample matching either is dropped: the first is a duplicate, the second + would leak the eval into training.""" + train = {_norm(e.text) for e in load_dir(TRAIN_DIR)} + evalset = {_norm(e.text) for e in load(EVAL_FILE)} if EVAL_FILE.exists() else set() + return train, evalset + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--raw", type=Path, default=DEFAULT_RAW, help="raw samples JSONL") + parser.add_argument("--out", type=Path, default=DEFAULT_OUT, help="labeled output JSONL") + parser.add_argument("--limit", type=int, help="stop after this many raw samples") + parser.add_argument("--model", help="teacher model override (else ROUTELET_TEACHER_MODEL)") + parser.add_argument("--dry-run", action="store_true", help="don't call the teacher or write") + args = parser.parse_args() + + if not args.raw.exists(): + sys.exit(f"no raw samples at {args.raw}; run pull_samples.py first") + + raw = load_raw(args.raw) + if args.limit: + raw = raw[: args.limit] + + train_keys, eval_keys = existing_keys() + # Dedup against the existing pool, the output file (so re-runs are idempotent + # even when --out lives outside data/), and within this batch. + seen = set(train_keys) + if args.out.exists(): + seen |= {_norm(r["text"]) for r in load_raw(args.out) if r.get("text")} + + client = None # created lazily; only needed for teacher labeling + labeled: list[dict] = [] + stats = {"client_label": 0, "teacher_label": 0, "dup": 0, "leak": 0, "empty": 0, "bad": 0} + + for row in raw: + text = (row.get("text") or "").strip() + if not text: + stats["empty"] += 1 + continue + + key = _norm(text) + if key in eval_keys: + stats["leak"] += 1 + continue + if key in seen: + stats["dup"] += 1 + continue + + # Prefer the free in-turn Claude label; fall back to the offline teacher. + label_str = row.get("claude_label") + if label_str: + try: + intent = Intent(label_str) + except ValueError: + stats["bad"] += 1 + continue + stats["client_label"] += 1 + elif args.dry_run: + # Count what a real run would send to the teacher without spending. + stats["teacher_label"] += 1 + seen.add(key) + continue + else: + if client is None: + import anthropic + + client = anthropic.Anthropic() + from routelet.teacher import classify + + intent = classify(client, preprocess(text), model=args.model) + stats["teacher_label"] += 1 + + seen.add(key) + labeled.append({"text": preprocess(text), "intent": intent.value}) + + if args.dry_run: + would = stats["client_label"] + stats["teacher_label"] + print(f"dry run: {would} new samples ({stats['teacher_label']} need the teacher)") + print(f" skipped: {stats['dup']} dup, {stats['leak']} eval-leak, " + f"{stats['empty']} empty, {stats['bad']} bad-label") + return + + args.out.parent.mkdir(parents=True, exist_ok=True) + with args.out.open("a") as f: + for rec in labeled: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + print(f"appended {len(labeled)} samples to {args.out}") + print(f" {stats['client_label']} client-labeled, {stats['teacher_label']} teacher-labeled") + print(f" skipped: {stats['dup']} dup, {stats['leak']} eval-leak, " + f"{stats['empty']} empty, {stats['bad']} bad-label") + + +if __name__ == "__main__": + main() diff --git a/Scripts/pull_samples.py b/Scripts/pull_samples.py new file mode 100644 index 0000000..dbbaced --- /dev/null +++ b/Scripts/pull_samples.py @@ -0,0 +1,107 @@ +"""Pull redacted distillation samples from the proxy's R2 bucket into one local +JSONL file. + +The aegis proxy stores one immutable JSON object per sample under +``samples///-.json`` (see aegis proxy/handlers/routelet.ts). +R2 has no append, so the per-sample-object layout sidesteps write races across +devices; compaction into a single batched JSONL happens here, at pull time. + +This is a full snapshot: every object under the prefix is read and written out, +deduped by R2 key. Re-running overwrites the output. The raw text was already +redacted on-device and scrubbed again by the proxy, but it is still real user +input, so the output is gitignored (see data/raw/). + +Reads R2 over the S3-compatible API. Set in the environment (or .env): + R2_ACCOUNT_ID Cloudflare account id + R2_ACCESS_KEY_ID R2 API token access key + R2_SECRET_ACCESS_KEY R2 API token secret + R2_BUCKET bucket name (default: aegis-routelet-samples) + +Create a read-only R2 API token in the Cloudflare dashboard +(R2 > Manage API Tokens). boto3 is a tooling dep installed into the venv, not a +project dependency. + +Usage: + uv run python Scripts/pull_samples.py + uv run python Scripts/pull_samples.py --prefix samples/2026-05-29/ --out data/raw/today.jsonl +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +import boto3 +from botocore.config import Config + +PROJECT_ROOT = Path(__file__).parent.parent +DEFAULT_OUT = PROJECT_ROOT / "data" / "raw" / "samples.jsonl" +DEFAULT_BUCKET = "aegis-routelet-samples" +DEFAULT_PREFIX = "samples/" + + +def _require_env(name: str) -> str: + val = os.environ.get(name) + if not val: + sys.exit( + f"missing ${name}. Set R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, and " + f"R2_SECRET_ACCESS_KEY (read-only R2 API token)." + ) + return val + + +def r2_client(): + account_id = _require_env("R2_ACCOUNT_ID") + return boto3.client( + "s3", + endpoint_url=f"https://{account_id}.r2.cloudflarestorage.com", + aws_access_key_id=_require_env("R2_ACCESS_KEY_ID"), + aws_secret_access_key=_require_env("R2_SECRET_ACCESS_KEY"), + # R2 ignores the region but boto3 requires one; signature v4 is required. + region_name="auto", + config=Config(signature_version="s3v4"), + ) + + +def pull(bucket: str, prefix: str, out: Path) -> tuple[int, int]: + """Read every object under prefix and write one JSON line per valid sample. + Returns (written, skipped).""" + client = r2_client() + paginator = client.get_paginator("list_objects_v2") + + out.parent.mkdir(parents=True, exist_ok=True) + written = 0 + skipped = 0 + with out.open("w") as f: + for page in paginator.paginate(Bucket=bucket, Prefix=prefix): + for obj in page.get("Contents", []): + key = obj["Key"] + body = client.get_object(Bucket=bucket, Key=key)["Body"].read() + try: + sample = json.loads(body) + except json.JSONDecodeError: + print(f" skip (bad json): {key}", file=sys.stderr) + skipped += 1 + continue + # One object is one sample; the key already dedupes, so write as-is. + f.write(json.dumps(sample, ensure_ascii=False) + "\n") + written += 1 + return written, skipped + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--bucket", default=DEFAULT_BUCKET) + parser.add_argument("--prefix", default=DEFAULT_PREFIX, help="R2 key prefix to pull") + parser.add_argument("--out", type=Path, default=DEFAULT_OUT, help="output JSONL path") + args = parser.parse_args() + + written, skipped = pull(args.bucket, args.prefix, args.out) + print(f"wrote {written} samples to {args.out} ({skipped} skipped)") + + +if __name__ == "__main__": + main() diff --git a/Scripts/refresh_haiku_baseline.py b/Scripts/refresh_haiku_baseline.py new file mode 100644 index 0000000..847101b --- /dev/null +++ b/Scripts/refresh_haiku_baseline.py @@ -0,0 +1,95 @@ +"""Refresh report/baselines.json by scoring Claude Haiku on the current frozen +holdout. This is the paid step the report's cached baseline stands in for. + +Mirrors src/routelet/evaluate.py exactly (same model, same shared teacher +prompt and schema, raw text in) so the cached number stays methodologically +identical to the eval baseline. Run after the holdout changes: + + ANTHROPIC_API_KEY=... uv run python Scripts/refresh_haiku_baseline.py + +Writes accuracy, p50/p95 latency, and per-class precision/recall/f1 to +report/baselines.json. +""" + +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path + +import anthropic +import numpy as np +from sklearn.metrics import accuracy_score, classification_report + +PROJECT_ROOT = Path(__file__).parent.parent +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +from routelet.data import Intent, load # noqa: E402 +from routelet.teacher import INTENT_SCHEMA, SYSTEM # noqa: E402 + +MODEL = "claude-haiku-4-5" +HOLDOUT = PROJECT_ROOT / "evals" / "holdout.jsonl" +BASELINES = PROJECT_ROOT / "report" / "baselines.json" + + +def main() -> None: + client = anthropic.Anthropic() + examples = load(HOLDOUT) + labels = [i.value for i in Intent] + + preds: list[str] = [] + latencies: list[float] = [] + for ex in examples: + t0 = time.perf_counter() + resp = client.messages.create( + model=MODEL, + max_tokens=100, + system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}], + messages=[{"role": "user", "content": ex.text}], + output_config={"format": {"type": "json_schema", "schema": INTENT_SCHEMA}}, + ) + latencies.append((time.perf_counter() - t0) * 1000.0) + text = next(b.text for b in resp.content if b.type == "text") + preds.append(json.loads(text)["intent"]) + + gold = [e.intent.value for e in examples] + acc = float(accuracy_score(gold, preds)) + lat = np.array(latencies) + report = classification_report(gold, preds, labels=labels, output_dict=True, zero_division=0) + + per_class = { + lbl: { + "precision": round(report[lbl]["precision"], 2), + "recall": round(report[lbl]["recall"], 2), + "f1": round(report[lbl]["f1-score"], 2), + "support": int(report[lbl]["support"]), + } + for lbl in labels + } + + out = { + "_note": ( + "Cached Claude baseline so report.py does not re-spend on the paid API " + "every regen. Regenerate by running Scripts/refresh_haiku_baseline.py " + "with ANTHROPIC_API_KEY against the current evals/holdout.jsonl." + ), + "haiku": { + "model": MODEL, + "eval_n": len(examples), + "accuracy": round(acc, 3), + "latency_p50_ms": round(float(np.percentile(lat, 50)), 0), + "latency_p95_ms": round(float(np.percentile(lat, 95)), 0), + "per_class": per_class, + }, + } + BASELINES.write_text(json.dumps(out, indent=2) + "\n") + + print(f"scored {len(examples)} rows on {MODEL}") + print(f" accuracy {acc:.3f}") + print(f" latency p50 {np.percentile(lat, 50):.0f}ms p95 {np.percentile(lat, 95):.0f}ms") + print(f"wrote {BASELINES}") + + +if __name__ == "__main__": + main() diff --git a/docs/metrics.md b/docs/metrics.md new file mode 100644 index 0000000..ad56e44 --- /dev/null +++ b/docs/metrics.md @@ -0,0 +1,26 @@ +# Performance report + +A reproducible figure report. `report.py` scores the models on the frozen +holdout and writes figures plus `metrics.json`. Re-run after every retrain. + +Status: planned, not built. + +## Figures + +1. Model comparison: TF-IDF baseline, SetFit (shipped), Claude Haiku. Accuracy on the same holdout. +2. Latency vs accuracy: routelet (~40ms, on-device) against Haiku (~700ms) and TF-IDF. +3. Confusion matrix, shipped model. +4. Per-class precision, recall, F1 with Wilson confidence intervals. +5. Confidence histogram, correct vs wrong predictions. + +## Rules + +- Score every model on the same frozen `evals/holdout.jsonl`. Never compare numbers across different eval sets. +- Apply `preprocess()` before inference, same as training. +- Haiku accuracy and latency come from `evaluate.py`. Needs `ANTHROPIC_API_KEY`. +- matplotlib only. Install into the venv, not pyproject. + +## Output + +- `report/*.png`: the figures. +- `report/metrics.json`: the numbers behind the figures. Single source for the README and any later dashboard. diff --git a/docs/model.md b/docs/model.md new file mode 100644 index 0000000..eb0829f --- /dev/null +++ b/docs/model.md @@ -0,0 +1,48 @@ +# How the model was made + +routelet's intent classifier, end to end. + +## Model + +SetFit: a fine-tuned `BAAI/bge-small-en-v1.5` encoder (384-d, on-device sized) +plus a logistic-regression head. Text in, one of 5 intents out. + +## Pipeline + +1. Data: hand-written seed commands per intent in `data/*.jsonl`, plus disfluent + variants from `augment.py` in `data/augmented.jsonl`. ~1100 rows. +2. preprocess: every command is normalized the same way at train and inference + (redact secrets, emails, long numbers), so train and serve never diverge. See + `preprocess.py`; the rules are pinned by a shared fixture with the Aegis side. +3. Train (`train_setfit.py`): contrastive fine-tune (unique sampling, 2 epochs), + then fit the LR head with class_weight balanced, then fit a temperature scalar + for confidence calibration. +4. Eval: score on the frozen `evals/holdout.jsonl` (200 rows, 40 per intent). +5. Export (`export_onnx.py`): encoder to `embedder.onnx` (opset 14) and head to + `head.json`, for on-device inference in Aegis via tract. + +## Numbers (200-row holdout) + +- routelet (SetFit): 0.89 +- TF-IDF baseline: 0.60 (the floor) +- Claude Haiku: 0.94 (the bar), at ~960ms vs routelet's ~40ms on-device + +## Honest limits + +The data is synthetic, so the holdout number is optimistic relative to real +voice input. The real lever from here is the data loop (log real commands, +Claude relabels them, retrain), not a bigger model. + +## Tool routing (explored, not shipped) + +We tried routing the `integration` branch to a specific tool on-device, so it +could pick spotify/gmail/github/youtube or detect "no tool, escalate." + +- Semantic retrieval, reusing the encoder: failed. The fine-tuned encoder + collapses cosine distances to ~0.01, unusable for retrieval. +- Classify into tools, a second fine-tune: works in-distribution (0.98) but + drops to 0.63 out-of-distribution (`evals/tool_routing_ood.jsonl`). The + synthetic data taught surface keywords, not intent. + +Not shipped. It needs real or much more varied data first. Aegis keeps letting +Claude pick the tool for now. diff --git a/evals/holdout.jsonl b/evals/holdout.jsonl index 348b220..7e368bf 100644 --- a/evals/holdout.jsonl +++ b/evals/holdout.jsonl @@ -29,7 +29,172 @@ {"text": "dont forget that we need to write docs for this project", "intent": "memory"} {"text": "make sure to save that my name is Bob", "intent": "memory"} {"text": "how old am I?", "intent": "memory"} +{"text": "make sure you know that im 12?", "intent": "memory"} {"text": "open youtube, search for lofi, play the top result", "intent": "agent"} {"text": "find the cheapest flight to tokyo and book it", "intent": "agent"} {"text": "can you open up my email, then remember my email address", "intent": "agent"} {"text": "can you open up youtube and search the news", "intent": "agent"} +{"text": "why is daddy donald trump such a bad president", "intent": "chat"} +{"text": "whats up with the weather today, its like where is it?", "intent": "chat"} +{"text": "open up my email and send me a message saying, what is the weather today?", "intent": "integration"} +{"text": "save that my wifi password is danny byron", "intent": "memory"} +{"text": "how do i set up my spotify integration?", "intent": "chat"} +{"text": "show me where that skip next song on spotify button is", "intent": "find_action"} +{"text": "find me the best restraunts nearby", "intent": "integration"} +{"text": "show me where america is on the map", "intent": "integration"} +{"text": "can you show me where the chat button is", "intent": "find_action"} +{"text": "can you open up the uhh the google maps and then look at local kebab shops for me please thanks bro", "intent": "agent"} +{"text": "can you skip this song and then play sicko mode on spotify", "intent": "agent"} +{"text": "can you open up github on google and then go to my most recent repo plz", "intent": "agent"} +{"text": "open github, press my profile, change name to daniel", "intent": "agent"} +{"text": "find me some cool songs and then play that for me", "intent": "agent"} +{"text": "talk about ur life", "intent": "chat"} +{"text": "when does the sun come up", "intent": "chat"} +{"text": "can you play remember my name", "intent": "integration"} +{"text": "whats my name on spotify", "intent": "chat"} +{"text": "open up a new tab on chrome and also try to type 123", "intent": "agent"} +{"text": "dont forget to use the memory agent", "intent": "memory"} +{"text": "my name is danny b and i was 11 then i turned 12 yesterday", "intent": "chat"} +{"text": "if i asked you a question about myself what would you do?", "intent": "chat"} +{"text": "where do i tap to start a new chat", "intent": "find_action"} +{"text": "find the thumbs up button on this video", "intent": "find_action"} +{"text": "uhh where is the like button on youtube", "intent": "find_action"} +{"text": "click the follow button on this profile", "intent": "find_action"} +{"text": "where's the uh the gear icon for settings", "intent": "find_action"} +{"text": "i was wondering if you could find the next page arrow", "intent": "find_action"} +{"text": "tap the the microphone icon to record", "intent": "find_action"} +{"text": "where do i click to upload a file here", "intent": "find_action"} +{"text": "press the play button on this video plz", "intent": "find_action"} +{"text": "hey can u scroll down to the footer", "intent": "find_action"} +{"text": "wheres the new playlist button in spotify", "intent": "find_action"} +{"text": "click on the the star to favorite it", "intent": "find_action"} +{"text": "find me the edit button on this post", "intent": "find_action"} +{"text": "where is the the cart icon up top", "intent": "find_action"} +{"text": "tap the filter button at the top of the list", "intent": "find_action"} +{"text": "um where do i go to click the report button", "intent": "find_action"} +{"text": "select the bold option in the toolbar", "intent": "find_action"} +{"text": "where's the reply all button in gmail", "intent": "find_action"} +{"text": "open my email and forward the latest one to my boss", "intent": "agent"} +{"text": "skip this track and then play sicko mode", "intent": "agent"} +{"text": "open github, go to my repo, and delete the old branch", "intent": "agent"} +{"text": "can you like check the weather and then text my mom if its gonna rain", "intent": "agent"} +{"text": "uhh open spotify and then make me a new playlist for working out", "intent": "agent"} +{"text": "search for a cheap flight to vegas and book the cheapest one", "intent": "agent"} +{"text": "i was wondering if you could find a nearby gym and add it to my maps", "intent": "agent"} +{"text": "pause the song and then turn up the volume bro", "intent": "agent"} +{"text": "open my calendar and move my 3pm meeting to friday", "intent": "agent"} +{"text": "find the top news story and then read it out to me plz", "intent": "agent"} +{"text": "open whatsapp and reply to the last message from alex", "intent": "agent"} +{"text": "look up a banana bread recipe and add the stuff to my shopping list", "intent": "agent"} +{"text": "um can you open youtube search lofi and then play the first one", "intent": "agent"} +{"text": "take a screenshot and then send it to my email", "intent": "agent"} +{"text": "is it possible to skip songs on spotify", "intent": "chat"} +{"text": "talk me through sending an email to someone", "intent": "chat"} +{"text": "how would i go about checking my email", "intent": "chat"} +{"text": "whats the weather like usually in boston in may", "intent": "chat"} +{"text": "uhh how do i even make a playlist on spotify", "intent": "chat"} +{"text": "can you explain like how forwarding an email works", "intent": "chat"} +{"text": "whats the deal with all these slack channels anyway", "intent": "chat"} +{"text": "is it hard to set up a github repo", "intent": "chat"} +{"text": "how does the spotify shuffle thing actually pick songs", "intent": "chat"} +{"text": "make sure you save that my middle name is alan", "intent": "memory"} +{"text": "keep track of the fact that im allergic to shellfish", "intent": "memory"} +{"text": "uhh log that my new phone number is five five five oh one two three", "intent": "memory"} +{"text": "write down that my hotel room is 412", "intent": "memory"} +{"text": "can you jot down that i parked in section c", "intent": "memory"} +{"text": "note down that my passport expires in june", "intent": "memory"} +{"text": "store that my office wifi network is called bluefox", "intent": "memory"} +{"text": "what was that thing i told you about my coffee order", "intent": "memory"} +{"text": "i was wondering if you could remind me what my gate number was", "intent": "memory"} +{"text": "hey can u tell me what my favorite restaurant is again", "intent": "memory"} +{"text": "uhh can you play some lofi beats", "intent": "integration"} +{"text": "skip this one", "intent": "integration"} +{"text": "text my sister that im running late plz", "intent": "integration"} +{"text": "open up google maps", "intent": "integration"} +{"text": "set a timer for like twenty minutes", "intent": "integration"} +{"text": "whats the weather gonna be tomorrow", "intent": "integration"} +{"text": "hey can u turn the volume down a little", "intent": "integration"} +{"text": "send a slack to the design channel saying im done", "intent": "integration"} +{"text": "find me a coffee shop thats open right now", "intent": "integration"} +{"text": "i was wondering if you could add eggs to my grocery list", "intent": "integration"} +{"text": "post this pic to my story", "intent": "integration"} +{"text": "play the new kendrick album", "intent": "integration"} +{"text": "is it weird that i talk to my computer all day", "intent": "chat"} +{"text": "whats a good way to stay focused while working", "intent": "chat"} +{"text": "do you think ai is gonna take over the world", "intent": "chat"} +{"text": "uhh where do i tap to collapse the sidebar", "intent": "find_action"} +{"text": "find the the pin icon to pin this message", "intent": "find_action"} +{"text": "where's the three dot thing next to my comment on youtube", "intent": "find_action"} +{"text": "click the the little eye to show the password", "intent": "find_action"} +{"text": "where do i hit to add a new column in this sheet", "intent": "find_action"} +{"text": "tap the swap camera button real quick", "intent": "find_action"} +{"text": "where is the the airdrop option in this share sheet", "intent": "find_action"} +{"text": "um wheres the resend code link on this screen", "intent": "find_action"} +{"text": "scroll over to the second photo in this post", "intent": "find_action"} +{"text": "where's the the slider to change playback speed", "intent": "find_action"} +{"text": "click on the trash can icon next to this row", "intent": "find_action"} +{"text": "where do i find the join with video button on zoom", "intent": "find_action"} +{"text": "tap the the chevron to open the dropdown", "intent": "find_action"} +{"text": "wheres the mark as unread thing in my inbox", "intent": "find_action"} +{"text": "press the floating compose pencil bottom right", "intent": "find_action"} +{"text": "where is the the toggle for two factor on this page", "intent": "find_action"} +{"text": "click the apply coupon link at checkout", "intent": "find_action"} +{"text": "open my notes find the grocery one and add paper towels", "intent": "agent"} +{"text": "uhh check the traffic and if its bad leave a note for my carpool", "intent": "agent"} +{"text": "look up tonights nba scores and drop them in the family chat", "intent": "agent"} +{"text": "pull up sarahs number from contacts then call her", "intent": "agent"} +{"text": "open figma duplicate this frame and rename it draft two", "intent": "agent"} +{"text": "find a barber open late and stick it in my calendar plz", "intent": "agent"} +{"text": "mute notifications and then start my pomodoro timer", "intent": "agent"} +{"text": "check my steps for today and if im under ten thousand remind me to walk", "intent": "agent"} +{"text": "um open the pdf scroll to page twelve and read me that part", "intent": "agent"} +{"text": "grab the wifi name from settings and text it to my guest", "intent": "agent"} +{"text": "open linkedin find the recruiter who messaged me and reply thanks", "intent": "agent"} +{"text": "look at my fridge inventory and order whatever im out of", "intent": "agent"} +{"text": "find a quiet cafe with wifi and navigate me there then mute my phone", "intent": "agent"} +{"text": "scan this receipt log the total and file it under march expenses", "intent": "agent"} +{"text": "open youtube find that tutorial i watched and add it to watch later", "intent": "agent"} +{"text": "check who texted me last and read it back then mark it read", "intent": "agent"} +{"text": "lower the thermostat a couple degrees", "intent": "integration"} +{"text": "uhh order me an uber to the office", "intent": "integration"} +{"text": "whens sunset tonight", "intent": "integration"} +{"text": "throw this article in my read later", "intent": "integration"} +{"text": "skip ahead like a minute in this podcast", "intent": "integration"} +{"text": "find me a ramen spot within walking distance", "intent": "integration"} +{"text": "dm taylor on instagram that im outside", "intent": "integration"} +{"text": "crank the bass up on this song", "intent": "integration"} +{"text": "whats the exchange rate for euros right now", "intent": "integration"} +{"text": "snap a pic of this whiteboard", "intent": "integration"} +{"text": "reorder my usual from the coffee app", "intent": "integration"} +{"text": "plz mark my last slack message as unread", "intent": "integration"} +{"text": "queue up the next episode of the office", "intent": "integration"} +{"text": "whats my battery percentage at", "intent": "integration"} +{"text": "venmo my roommate twenty bucks for pizza", "intent": "integration"} +{"text": "turn on the hallway lamp", "intent": "integration"} +{"text": "uhh jot down that my seat for the concert is row m", "intent": "memory"} +{"text": "keep track of my checking account ending in 4417", "intent": "memory"} +{"text": "make a note that my therapist is on tuesdays at 4", "intent": "memory"} +{"text": "remember i swim laps every other morning", "intent": "memory"} +{"text": "save that my storage unit code is star nine nine", "intent": "memory"} +{"text": "whats my storage unit code again", "intent": "memory"} +{"text": "log that i switched my insurance to bluecross", "intent": "memory"} +{"text": "remind me which gate my flight leaves from", "intent": "memory"} +{"text": "dont forget my kid gets picked up at 3 on wednesdays", "intent": "memory"} +{"text": "what did i tell you about my desk number at work", "intent": "memory"} +{"text": "put it down that im cutting out caffeine this month", "intent": "memory"} +{"text": "store that my mechanic is over on fifth street", "intent": "memory"} +{"text": "uhh what was my reservation name for the restaurant again", "intent": "memory"} +{"text": "note that my nephews birthday is the 19th", "intent": "memory"} +{"text": "save my preferred pharmacy as the one on elm", "intent": "memory"} +{"text": "whats the name of my emergency contact i had you save", "intent": "memory"} +{"text": "how do i unfollow someone on instagram without them knowing", "intent": "chat"} +{"text": "whats the deal with all the song skipping limits on free spotify", "intent": "chat"} +{"text": "is it actually rude to leave a slack message on read", "intent": "chat"} +{"text": "uhh how does a recurring calendar event even work", "intent": "chat"} +{"text": "why do my texts turn green sometimes instead of blue", "intent": "chat"} +{"text": "whats the longest river in the world", "intent": "chat"} +{"text": "how would you explain compound interest to a kid", "intent": "chat"} +{"text": "do you get bored when im not talking to you", "intent": "chat"} +{"text": "whats a good comeback when someone says youre overthinking", "intent": "chat"} +{"text": "how come some emails go straight to spam", "intent": "chat"} +{"text": "is whispering actually worse for your voice than talking", "intent": "chat"} diff --git a/evals/tool_routing_ood.jsonl b/evals/tool_routing_ood.jsonl new file mode 100644 index 0000000..1f05f99 --- /dev/null +++ b/evals/tool_routing_ood.jsonl @@ -0,0 +1,52 @@ +{"text": "mute the music", "tool": "spotify"} +{"text": "crank it louder", "tool": "spotify"} +{"text": "go back a song", "tool": "spotify"} +{"text": "bump this track up in the queue", "tool": "spotify"} +{"text": "i want to hear some kendrick", "tool": "spotify"} +{"text": "put on something lo-fi while i work", "tool": "spotify"} +{"text": "stop the song", "tool": "spotify"} +{"text": "what album is this from", "tool": "spotify"} +{"text": "save this track to my library", "tool": "spotify"} +{"text": "repeat this song", "tool": "spotify"} +{"text": "any new messages in my inbox", "tool": "gmail"} +{"text": "write back to the recruiter and say i'm interested", "tool": "gmail"} +{"text": "pull up my drafts", "tool": "gmail"} +{"text": "get rid of all the marketing junk", "tool": "gmail"} +{"text": "did the client respond yet", "tool": "gmail"} +{"text": "flag the message from my landlord as important", "tool": "gmail"} +{"text": "send a note to the whole team about the deadline", "tool": "gmail"} +{"text": "has anyone emailed me about the invoice", "tool": "gmail"} +{"text": "shoot a reply to the therapist confirming thursday", "tool": "gmail"} +{"text": "empty the trash folder in my mail", "tool": "gmail"} +{"text": "star this repo", "tool": "github"} +{"text": "what's the status of my latest build", "tool": "github"} +{"text": "who left a comment on my pull request", "tool": "github"} +{"text": "squash and merge the feature branch", "tool": "github"} +{"text": "open a ticket for the crash we saw yesterday", "tool": "github"} +{"text": "unwatch that noisy repository", "tool": "github"} +{"text": "show me all unresolved review comments on my PR", "tool": "github"} +{"text": "tag v1.2.0 on the main branch", "tool": "github"} +{"text": "list contributors to this project", "tool": "github"} +{"text": "what did i push last tuesday", "tool": "github"} +{"text": "find me a beginner piano lesson video", "tool": "youtube"} +{"text": "pull up that documentary about the cia on yt", "tool": "youtube"} +{"text": "i want to watch a ted talk about sleep", "tool": "youtube"} +{"text": "get me a video on how to sharpen a knife", "tool": "youtube"} +{"text": "show me the highlights reel from the superbowl on yt", "tool": "youtube"} +{"text": "open the channel for that cooking guy i follow", "tool": "youtube"} +{"text": "dislike this video", "tool": "youtube"} +{"text": "slow this down to half speed", "tool": "youtube"} +{"text": "find a video explaining what a transformer model is", "tool": "youtube"} +{"text": "save this to watch later on yt", "tool": "youtube"} +{"text": "set a reminder for my dentist appointment", "tool": "no_tool"} +{"text": "stream the next episode of succession on hbo", "tool": "no_tool"} +{"text": "turn on the lights in the kitchen", "tool": "no_tool"} +{"text": "what time is it in tokyo", "tool": "no_tool"} +{"text": "add eggs to my shopping list", "tool": "no_tool"} +{"text": "book me a flight to denver next weekend", "tool": "no_tool"} +{"text": "open slack and check the dev channel", "tool": "no_tool"} +{"text": "navigate home", "tool": "no_tool"} +{"text": "call mom", "tool": "no_tool"} +{"text": "translate this to french", "tool": "no_tool"} +{"text": "queue up the next office episode", "tool": "no_tool"} +{"text": "post this photo to instagram", "tool": "no_tool"} diff --git a/experiments/README.md b/experiments/README.md new file mode 100644 index 0000000..651542f --- /dev/null +++ b/experiments/README.md @@ -0,0 +1,14 @@ +# experiments + +Throwaway prototypes that are **not** part of the routelet library or its build. +Kept for reference, not maintained. + +- `tool_router.py`, `tool_classifier_proto.py`, `tools_proto.py`: an earlier + exploration of a *tool*-routing classifier (spotify / gmail / github / youtube + / no_tool), separate from the shipped 5-intent router. +- `run_ood_eval.py`: the out-of-distribution eval for that experiment, against + `evals/tool_routing_ood.jsonl` and `models/setfit_tools_proto`. + +These import `routelet.preprocess`, so run them through the project env: + + uv run python experiments/tool_router.py diff --git a/experiments/run_ood_eval.py b/experiments/run_ood_eval.py new file mode 100644 index 0000000..9f07d41 --- /dev/null +++ b/experiments/run_ood_eval.py @@ -0,0 +1,114 @@ +"""OOD eval for models/setfit_tools_proto against evals/tool_routing_ood.jsonl.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).parent.parent +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +from routelet.preprocess import preprocess + +# --------------------------------------------------------------------------- +# 1. Load OOD eval +# --------------------------------------------------------------------------- + +OOD_PATH = PROJECT_ROOT / "evals" / "tool_routing_ood.jsonl" +TRAIN_PATH = PROJECT_ROOT / "data" / "integration.jsonl" +MODEL_PATH = PROJECT_ROOT / "models" / "setfit_tools_proto" + +TOOLS = ["spotify", "gmail", "github", "youtube", "no_tool"] + +ood_rows: list[tuple[str, str]] = [] +with open(OOD_PATH) as f: + for line in f: + line = line.strip() + if not line: + continue + obj = json.loads(line) + ood_rows.append((obj["text"], obj["tool"])) + +print(f"OOD eval rows: {len(ood_rows)}") +from collections import Counter +dist = Counter(lbl for _, lbl in ood_rows) +for t in TOOLS: + print(f" {t:<12} {dist[t]:>3}") + +# --------------------------------------------------------------------------- +# 2. Leakage check: exact text match against training pool +# --------------------------------------------------------------------------- + +train_texts: set[str] = set() +with open(TRAIN_PATH) as f: + for line in f: + line = line.strip() + if not line: + continue + obj = json.loads(line) + train_texts.add(obj["text"].strip().lower()) + +overlaps = [ + (text, lbl) + for text, lbl in ood_rows + if text.strip().lower() in train_texts +] +print(f"\nLeakage check: {len(overlaps)} overlap(s) with integration.jsonl (must be 0)") +for text, lbl in overlaps: + print(f" DUPLICATE: {text!r} ({lbl})") + +# --------------------------------------------------------------------------- +# 3. Load model and predict +# --------------------------------------------------------------------------- + +from setfit import SetFitModel # noqa: E402 + +print(f"\nLoading model from {MODEL_PATH} ...") +model = SetFitModel.from_pretrained(str(MODEL_PATH)) + +texts = [text for text, _ in ood_rows] +labels = [lbl for _, lbl in ood_rows] + +preprocessed = [preprocess(t) for t in texts] +preds_raw = model.predict(preprocessed) +preds = [str(p) for p in preds_raw] + +# --------------------------------------------------------------------------- +# 4. Report +# --------------------------------------------------------------------------- + +from sklearn.metrics import classification_report, confusion_matrix # noqa: E402 + +print("\n" + "=" * 64) +print("OOD EVAL RESULTS") +print("=" * 64) + +print(classification_report(labels, preds, labels=TOOLS, zero_division=0)) + +cm = confusion_matrix(labels, preds, labels=TOOLS) +col_w = max(len(t) for t in TOOLS) + 2 +header = " " * col_w + "".join(f"{t:>{col_w}}" for t in TOOLS) +print("Confusion matrix (rows=true, cols=pred):") +print(header) +for i, row_lbl in enumerate(TOOLS): + row_str = f"{row_lbl:<{col_w}}" + "".join(f"{v:>{col_w}}" for v in cm[i]) + print(row_str) + +nt_true = [lbl == "no_tool" for lbl in labels] +nt_correct = sum(p == "no_tool" for p, t in zip(preds, labels) if t == "no_tool") +nt_total = sum(nt_true) +nt_recall = nt_correct / nt_total if nt_total else 0.0 +print(f"\nno_tool recall: {nt_recall:.3f} ({nt_correct}/{nt_total})") + +correct = sum(p == t for p, t in zip(preds, labels)) +acc = correct / len(labels) +print(f"overall accuracy: {acc:.3f} ({correct}/{len(labels)})") + +print("\nMisclassifications:") +errors = [(labels[i], preds[i], texts[i]) for i in range(len(labels)) if preds[i] != labels[i]] +if not errors: + print(" none") +else: + for true_lbl, pred_lbl, text in errors: + print(f" true={true_lbl:<12} pred={pred_lbl:<12} {text!r}") diff --git a/experiments/tool_classifier_proto.py b/experiments/tool_classifier_proto.py new file mode 100644 index 0000000..bedd6ee --- /dev/null +++ b/experiments/tool_classifier_proto.py @@ -0,0 +1,387 @@ +"""Approach A prototype: classify integration commands directly to tool labels. + +Two variants on the same data: + (a) FROZEN - embed with existing models/setfit encoder, fit fresh LR head. + (b) RETRAIN - fine-tune a fresh SetFit model with tool-level labels. + +Usage: + .venv/bin/python -m routelet.tool_classifier_proto + +Saves retrained model to models/setfit_tools_proto (separate from models/setfit). +""" + +from __future__ import annotations + +import json +import random +from pathlib import Path + +import numpy as np +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import classification_report, confusion_matrix +from sklearn.model_selection import StratifiedShuffleSplit + +from routelet.preprocess import preprocess + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- + +PROJECT_ROOT = Path(__file__).parent.parent.parent +SETFIT_MODEL_DIR = PROJECT_ROOT / "models" / "setfit" +PROTO_MODEL_OUT = PROJECT_ROOT / "models" / "setfit_tools_proto" + +TOOLS = ["spotify", "gmail", "github", "youtube", "no_tool"] + +# --------------------------------------------------------------------------- +# Dataset: hand-label each integration.jsonl row by tool +# --------------------------------------------------------------------------- + +# Rules used for labeling: +# spotify - music playback control (play/pause/skip/volume/what's playing) +# gmail - read/send/search email +# github - PRs / issues / repos / push / pull +# youtube - search or play youtube videos +# no_tool - everything else (no matching integration) + +_RAW_LABELS: list[tuple[str, str]] = [ + # --- from integration.jsonl (102 rows) --- + ("pause the music", "spotify"), + ("turn off the living room lights", "no_tool"), + ("set a timer for ten minutes", "no_tool"), + ("what's the weather today", "no_tool"), + ("text mom i'm on my way", "no_tool"), + ("set an alarm for 7am", "no_tool"), + ("add milk to my shopping list", "no_tool"), + ("turn the volume up", "spotify"), + ("call the dentist", "no_tool"), + ("play the next episode on netflix", "no_tool"), + ("mute my microphone", "no_tool"), + ("take a screenshot", "no_tool"), + ("send a slack message to the team channel", "no_tool"), + ("start a 25 minute focus timer", "no_tool"), + ("open spotify", "spotify"), + ("add a meeting to my calendar for 3pm", "no_tool"), + ("turn on do not disturb", "no_tool"), + ("search google for italian restaurants nearby", "no_tool"), + ("start recording my screen", "no_tool"), + ("lower the brightness", "no_tool"), + ("play some jazz", "spotify"), + ("snooze the alarm", "no_tool"), + ("send an email to my boss saying i'll be late", "gmail"), + ("connect to my bluetooth headphones", "no_tool"), + ("set the thermostat to 70 degrees", "no_tool"), + ("share my location with sarah", "no_tool"), + ("turn on the porch light", "no_tool"), + ("add a reminder to call the bank tomorrow", "no_tool"), + ("what's on my calendar today", "no_tool"), + ("skip to the next track", "spotify"), + ("play my workout playlist", "spotify"), + ("turn off the tv", "no_tool"), + ("text dad happy birthday", "no_tool"), + ("dim the bedroom lights", "no_tool"), + ("increase the volume to max", "spotify"), + ("pause the podcast", "spotify"), + ("send a whatsapp to alex", "no_tool"), + ("turn on airplane mode", "no_tool"), + ("play the latest taylor swift album on spotify", "spotify"), + ("set my status to away", "no_tool"), + ("play remember the name by fort minor", "spotify"), + ("play the song memories by maroon 5", "spotify"), + ("add remember me to my playlist", "spotify"), + ("play what's my age again by blink 182", "spotify"), + ("play the night we met", "spotify"), + ("play unforgettable by nat king cole", "spotify"), + ("skip to the song called my way", "spotify"), + ("play try to remember", "spotify"), + ("play memory from the musical cats", "spotify"), + ("queue up remember when by alan jackson", "spotify"), + ("play the song called my name is", "spotify"), + ("play i will remember you by sarah mclachlan", "spotify"), + ("play check yes juliet", "spotify"), + ("play the chainsmokers song memories", "spotify"), + ("play don't you forget about me", "spotify"), + ("add my favorite song to the queue", "spotify"), + ("book a table for two at an italian place saturday night", "no_tool"), + ("find me a thai place nearby", "no_tool"), + ("show me my unread emails", "gmail"), + ("what's playing right now", "spotify"), + ("search google for how to fix a leaky faucet", "no_tool"), + ("show me the best sushi spots around here", "no_tool"), + ("find a gas station close by", "no_tool"), + ("what are the top rated burgers near me", "no_tool"), + ("post this to twitter", "no_tool"), + ("tweet good morning everyone", "no_tool"), + ("send a text to jenny saying running late", "no_tool"), + ("rewind the podcast thirty seconds", "spotify"), + ("turn the volume down a bit", "spotify"), + ("go back to the previous track", "spotify"), + ("what's the weather like in chicago", "no_tool"), + ("show me directions to the airport", "no_tool"), + ("find me a parking garage near the stadium", "no_tool"), + ("what restaurants are open right now near me", "no_tool"), + ("look up the address for the nearest pharmacy", "no_tool"), + ("play that new drake song", "spotify"), + ("shuffle my liked songs", "spotify"), + ("raise the volume", "spotify"), + ("post a photo to my instagram story", "no_tool"), + ("send an email to the landlord about the rent", "gmail"), + ("what time does the closest target close", "no_tool"), + ("find me a hotel in seattle for tonight", "no_tool"), + ("google the showtimes for the new batman movie", "no_tool"), + ("fast forward to the chorus", "spotify"), + ("show me my recent transactions", "no_tool"), + ("whats the traffic like on my way home", "no_tool"), + ("find me a vet that's open on sundays", "no_tool"), + ("play the album on repeat", "spotify"), + ("send a dm to mark on instagram", "no_tool"), + ("look up flights to miami", "no_tool"), + ("set the volume to fifty percent", "spotify"), + ("show me the menu for the pizza place down the street", "no_tool"), + ("play something chill", "spotify"), + ("shoot a text to ryan saying happy friday", "no_tool"), + ("throw on my road trip playlist", "spotify"), + ("ping the marketing channel that the deck is ready", "no_tool"), + ("bump the brightness down", "no_tool"), + ("fire off an email to support about my broken order", "gmail"), + ("drop a pin at my current location and send it to dad", "no_tool"), + ("queue up that new sza track", "spotify"), + ("remind me to take out the trash tonight", "no_tool"), + ("pull up directions to the nearest hospital", "no_tool"), + # --- synthetic gmail examples (not in integration.jsonl) --- + ("check my email", "gmail"), + ("do i have any new emails", "gmail"), + ("search my inbox for the receipt from amazon", "gmail"), + ("find the email with the conference details", "gmail"), + ("compose an email to the team about the deadline", "gmail"), + ("reply to the last email from sarah", "gmail"), + ("any unread messages in my gmail", "gmail"), + ("send a quick email to hr asking about my leave balance", "gmail"), + ("look for emails from my landlord", "gmail"), + ("forward that email to jake", "gmail"), + # --- synthetic github examples (not in integration.jsonl) --- + ("show my open pull requests", "github"), + ("list open issues on my repo", "github"), + ("any new prs filed today", "github"), + ("check the status of my pull request", "github"), + ("what repos do i have on github", "github"), + ("show unreviewed pull requests", "github"), + ("are there any failing checks on my pr", "github"), + ("list the issues assigned to me", "github"), + ("what branches do i have on that repo", "github"), + ("did anyone comment on my pr", "github"), + # --- synthetic youtube examples (not in integration.jsonl) --- + ("play lofi beats on youtube", "youtube"), + ("search youtube for that lex fridman interview", "youtube"), + ("find a rust async tutorial on youtube", "youtube"), + ("play the latest veritasium video", "youtube"), + ("look up that cooking video on youtube", "youtube"), + ("search youtube for how to make sourdough bread", "youtube"), + ("put on some ambient music on youtube", "youtube"), + ("show me that python tutorial on youtube", "youtube"), + ("find the mr beast challenge video", "youtube"), + ("play that ted talk on youtube", "youtube"), + # --- synthetic no_tool examples (25-30) --- + ("find me restaurants nearby", "no_tool"), + ("set a timer for 10 minutes", "no_tool"), + ("whats the weather right now", "no_tool"), + ("turn up the brightness", "no_tool"), + ("order an uber", "no_tool"), + ("search google for something", "no_tool"), + ("pull up the rust docs", "no_tool"), + ("what's 15% of 80", "no_tool"), + ("open the terminal", "no_tool"), + ("translate this to spanish", "no_tool"), + ("lock my screen", "no_tool"), + ("copy that to the clipboard", "no_tool"), + ("remind me at 8pm", "no_tool"), + ("what's the score of the game", "no_tool"), + ("how many calories in a banana", "no_tool"), + ("set my wallpaper to this image", "no_tool"), + ("call mom", "no_tool"), + ("navigate to the coffee shop", "no_tool"), + ("book a flight to new york", "no_tool"), + ("what time is it in london", "no_tool"), + ("convert 100 dollars to euros", "no_tool"), + ("schedule a meeting for tomorrow at 2pm", "no_tool"), + ("close all the open windows", "no_tool"), + ("take a screenshot of the screen", "no_tool"), + ("turn off wifi", "no_tool"), + ("show me the news headlines", "no_tool"), + ("check the stock price of apple", "no_tool"), + ("send a text to mom", "no_tool"), +] + + +def _load_dataset() -> tuple[list[str], list[str]]: + texts = [preprocess(t) for t, _ in _RAW_LABELS] + labels = [lbl for _, lbl in _RAW_LABELS] + return texts, labels + + +def _stratified_split( + texts: list[str], + labels: list[str], + test_size: float = 0.20, + seed: int = 42, +) -> tuple[list[str], list[str], list[str], list[str]]: + sss = StratifiedShuffleSplit(n_splits=1, test_size=test_size, random_state=seed) + idx_tr, idx_te = next(sss.split(texts, labels)) + tr_texts = [texts[i] for i in idx_tr] + tr_labels = [labels[i] for i in idx_tr] + te_texts = [texts[i] for i in idx_te] + te_labels = [labels[i] for i in idx_te] + return tr_texts, tr_labels, te_texts, te_labels + + +# --------------------------------------------------------------------------- +# Report helpers +# --------------------------------------------------------------------------- + + +def _print_report( + name: str, + te_labels: list[str], + preds: list[str], + proba: np.ndarray | None = None, +) -> None: + print(f"\n{'=' * 60}") + print(f"VARIANT: {name}") + print(f"{'=' * 60}") + print(classification_report(te_labels, preds, labels=TOOLS, zero_division=0)) + + cm = confusion_matrix(te_labels, preds, labels=TOOLS) + col_w = max(len(t) for t in TOOLS) + 2 + header = " " * col_w + "".join(f"{t:>{col_w}}" for t in TOOLS) + print("confusion matrix (rows=true, cols=pred):") + print(header) + for i, row_lbl in enumerate(TOOLS): + row_str = f"{row_lbl:<{col_w}}" + "".join(f"{v:>{col_w}}" for v in cm[i]) + print(row_str) + + # Focus: no_tool recall + nt_true = [1 if lbl == "no_tool" else 0 for lbl in te_labels] + nt_pred = [1 if p == "no_tool" else 0 for p in preds] + nt_correct = sum(a == b for a, b in zip(nt_true, nt_pred) if a == 1) + nt_total = sum(nt_true) + nt_recall = nt_correct / nt_total if nt_total > 0 else 0.0 + print(f"\nno_tool recall: {nt_recall:.3f} ({nt_correct}/{nt_total})") + + +# --------------------------------------------------------------------------- +# Variant (a): FROZEN encoder + fresh LR head +# --------------------------------------------------------------------------- + + +def run_frozen() -> None: + print("\n" + "=" * 60) + print("VARIANT (a): FROZEN encoder + fresh LR head") + print("=" * 60) + print(f"loading SetFit encoder from {SETFIT_MODEL_DIR} ...") + + from setfit import SetFitModel + + model = SetFitModel.from_pretrained(str(SETFIT_MODEL_DIR)) + enc = model.model_body + + texts, labels = _load_dataset() + label_counts = {t: labels.count(t) for t in TOOLS} + print(f"dataset: {len(texts)} examples | " + " ".join(f"{k}={v}" for k, v in label_counts.items())) + + tr_texts, tr_labels, te_texts, te_labels = _stratified_split(texts, labels) + print(f"train: {len(tr_texts)} test: {len(te_texts)}") + + print("embedding training set ...") + tr_embs = enc.encode(tr_texts, convert_to_numpy=True, show_progress_bar=False) + print("embedding test set ...") + te_embs = enc.encode(te_texts, convert_to_numpy=True, show_progress_bar=False) + + clf = LogisticRegression(class_weight="balanced", max_iter=1000, random_state=42) + clf.fit(tr_embs, tr_labels) + + preds = clf.predict(te_embs) + proba = clf.predict_proba(te_embs) + + _print_report("(a) FROZEN", te_labels, list(preds), proba) + + +# --------------------------------------------------------------------------- +# Variant (b): RETRAIN SetFit with tool labels +# --------------------------------------------------------------------------- + + +def run_retrain() -> None: + print("\n" + "=" * 60) + print("VARIANT (b): RETRAIN SetFit with tool labels") + print("=" * 60) + + import torch + from datasets import Dataset + from setfit import SetFitModel, Trainer, TrainingArguments + + random.seed(0) + np.random.seed(0) + torch.manual_seed(0) + + BASE = "BAAI/bge-small-en-v1.5" + texts, labels = _load_dataset() + label_counts = {t: labels.count(t) for t in TOOLS} + print(f"dataset: {len(texts)} examples | " + " ".join(f"{k}={v}" for k, v in label_counts.items())) + + tr_texts, tr_labels, te_texts, te_labels = _stratified_split(texts, labels) + print(f"train: {len(tr_texts)} test: {len(te_texts)}") + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"training on {device}") + + train_ds = Dataset.from_dict({"text": tr_texts, "label": tr_labels}) + + model = SetFitModel.from_pretrained( + BASE, + labels=TOOLS, + device=device, + head_params={"class_weight": "balanced"}, + ) + + args = TrainingArguments( + batch_size=16, + num_epochs=2, + sampling_strategy="unique", + ) + + trainer = Trainer(model=model, args=args, train_dataset=train_ds) + trainer.train() + + preds_raw = model.predict(te_texts) + preds = [str(p) for p in preds_raw] + + _print_report("(b) RETRAIN", te_labels, preds) + + PROTO_MODEL_OUT.mkdir(parents=True, exist_ok=True) + model.save_pretrained(str(PROTO_MODEL_OUT)) + print(f"\nsaved proto model to {PROTO_MODEL_OUT}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + print("=== APPROACH A PROTOTYPE: TOOL CLASSIFICATION ===") + print(f"tools: {TOOLS}") + texts, labels = _load_dataset() + print(f"total labeled examples: {len(texts)}") + label_counts = {t: labels.count(t) for t in TOOLS} + print("class distribution: " + " ".join(f"{k}={v}" for k, v in label_counts.items())) + + run_frozen() + run_retrain() + + print("\n" + "=" * 60) + print("APPROACH A PROTOTYPE COMPLETE") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/experiments/tool_router.py b/experiments/tool_router.py new file mode 100644 index 0000000..c6fa8f2 --- /dev/null +++ b/experiments/tool_router.py @@ -0,0 +1,401 @@ +"""Semantic tool router for integration commands. + +Uses a sentence encoder to embed commands and tool descriptions, then picks +the best-matching tool via cosine similarity. Commands with no good match fall +below TOOL_THRESH and escalate rather than route. + +TWO ENCODER MODES are supported for the eval: + - finetuned: SetFitModel.from_pretrained("models/setfit").model_body + (the same encoder the intent classifier uses) + - base: BAAI/bge-small-en-v1.5 loaded directly via sentence_transformers + (the pretrained backbone before contrastive fine-tuning) + +The eval runs both and reports the comparison. route() uses the finetuned +encoder by default to stay consistent with the rest of the pipeline, but +the eval will show whether that is viable. + +Usage: + .venv/bin/python -m routelet.tool_router # run full eval + .venv/bin/python -m routelet.tool_router 0.40 # custom threshold + .venv/bin/python -m routelet.tool_router 0.40 base # base model only + +route(command, threshold) is the public API for the Aegis integration path. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Optional + +import numpy as np + +# --------------------------------------------------------------------------- +# Tool registry +# --------------------------------------------------------------------------- + +TOOLS: dict[str, str] = { + "spotify": "control music playback: play, pause, skip, previous track, volume on spotify", + "gmail": "read, search, and send email", + "github": "list and check github pull requests, issues, and repositories", + "youtube": "search for and play youtube videos", +} + +TOOL_NAMES: list[str] = list(TOOLS.keys()) + +MODEL_DIR = Path(__file__).parent.parent.parent / "models" / "setfit" +BASE_MODEL_ID = "BAAI/bge-small-en-v1.5" + +# --------------------------------------------------------------------------- +# Encoder loading +# --------------------------------------------------------------------------- + +_finetuned_encoder = None +_finetuned_tool_embs: Optional[np.ndarray] = None + +_base_encoder = None +_base_tool_embs: Optional[np.ndarray] = None + + +def _load_finetuned(): + global _finetuned_encoder, _finetuned_tool_embs + if _finetuned_encoder is not None: + return _finetuned_encoder, _finetuned_tool_embs + from setfit import SetFitModel + model = SetFitModel.from_pretrained(str(MODEL_DIR)) + enc = model.model_body + tool_embs = enc.encode( + list(TOOLS.values()), convert_to_numpy=True, show_progress_bar=False + ) + _finetuned_encoder = enc + _finetuned_tool_embs = tool_embs + return enc, tool_embs + + +def _load_base(): + global _base_encoder, _base_tool_embs + if _base_encoder is not None: + return _base_encoder, _base_tool_embs + from sentence_transformers import SentenceTransformer + enc = SentenceTransformer(BASE_MODEL_ID) + tool_embs = enc.encode( + list(TOOLS.values()), convert_to_numpy=True, normalize_embeddings=True + ) + _base_encoder = enc + _base_tool_embs = tool_embs + return enc, tool_embs + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def route(command: str, threshold: float = 0.40) -> tuple[Optional[str], float]: + """Route a command to a tool or return (None, score) to escalate. + + Uses the fine-tuned SetFit encoder. See the eval for why the base encoder + is a better choice; this function signature is kept stable for Aegis. + + Parameters + ---------- + command: raw user command (preprocess applied internally). + threshold: minimum cosine similarity to accept a routing decision. + + Returns + ------- + (tool_name, score) if max similarity >= threshold + (None, score) otherwise (escalate) + """ + from routelet.preprocess import preprocess + enc, tool_embs = _load_finetuned() + text = preprocess(command) + emb = enc.encode([text], convert_to_numpy=True, show_progress_bar=False)[0] + sims = tool_embs @ emb + best_idx = int(np.argmax(sims)) + best_score = float(sims[best_idx]) + if best_score >= threshold: + return TOOL_NAMES[best_idx], best_score + return None, best_score + + +# --------------------------------------------------------------------------- +# Eval dataset +# --------------------------------------------------------------------------- + +# HAS-TOOL: (command, expected_tool) +HAS_TOOL: list[tuple[str, str]] = [ + # spotify + ("skip this song", "spotify"), + ("uhh skip this song", "spotify"), + ("pause the music", "spotify"), + ("turn the volume up on spotify", "spotify"), + ("play something upbeat", "spotify"), + ("next track please", "spotify"), + ("go back to the previous song", "spotify"), + ("lower the spotify volume", "spotify"), + # gmail + ("can you pull up my unread emails", "gmail"), + ("search my email for the invoice from last month", "gmail"), + ("send an email to alex saying the meeting is rescheduled", "gmail"), + ("any new emails from my boss", "gmail"), + ("reply to that last email", "gmail"), + ("check if I have email from github", "gmail"), + ("compose a quick email to the team", "gmail"), + ("find the email with the flight confirmation", "gmail"), + # github + ("show my open PRs", "github"), + ("list open pull requests on my repo", "github"), + ("any new issues filed today", "github"), + ("check the status of my pull request", "github"), + ("what repos do I have on github", "github"), + ("show unreviewed pull requests", "github"), + ("are there any failing checks on my PR", "github"), + # youtube + ("play lofi beats on youtube", "youtube"), + ("search youtube for that interview with lex fridman", "youtube"), + ("find a tutorial on rust async on youtube", "youtube"), + ("play the latest video from veritasium", "youtube"), + ("look up that cooking video I watched last week", "youtube"), + ("search for how to make sourdough bread youtube", "youtube"), + ("put on some ambient music on youtube", "youtube"), +] + +# NO-TOOL: commands that should escalate (no matching integration) +NO_TOOL: list[str] = [ + "find me restaurants nearby", + "pull up the rust docs for this function", + "set a timer for 10 minutes", + "whats the weather right now", + "turn up the screen brightness", + "order me an uber", + "what's 15 percent of 80", + "open the terminal", + "translate this sentence to spanish", + "remind me to take my meds at 8pm", + "lock the screen", + "copy that to the clipboard", +] + + +# --------------------------------------------------------------------------- +# Core eval logic (encoder-agnostic) +# --------------------------------------------------------------------------- + + +def _embed(texts: list[str], enc, use_preprocess: bool, normalize: bool) -> np.ndarray: + if use_preprocess: + from routelet.preprocess import preprocess + texts = [preprocess(t) for t in texts] + if normalize: + return enc.encode(texts, convert_to_numpy=True, normalize_embeddings=True) + return enc.encode(texts, convert_to_numpy=True, show_progress_bar=False) + + +def _run_single_eval( + enc, + tool_embs: np.ndarray, + use_preprocess: bool, + normalize: bool, + label: str, +) -> None: + has_cmds = [cmd for cmd, _ in HAS_TOOL] + has_labels = [lbl for _, lbl in HAS_TOOL] + no_cmds = list(NO_TOOL) + + has_embs = _embed(has_cmds, enc, use_preprocess, normalize) + no_embs = _embed(no_cmds, enc, use_preprocess, normalize) + + has_sims = has_embs @ tool_embs.T + no_sims = no_embs @ tool_embs.T + + has_top = has_sims.max(axis=1) + no_top = no_sims.max(axis=1) + + print(f"\n{'=' * 60}") + print(f"ENCODER: {label}") + print(f"{'=' * 60}") + + print(f"\n--- 1. SIMILARITY DISTRIBUTION ---") + print(f"HAS-TOOL n={len(has_top):3d} min={has_top.min():.4f} " + f"mean={has_top.mean():.4f} max={has_top.max():.4f}") + print(f"NO-TOOL n={len(no_top):3d} min={no_top.min():.4f} " + f"mean={no_top.mean():.4f} max={no_top.max():.4f}") + gap = float(has_top.min() - no_top.max()) + mean_gap = float(has_top.mean() - no_top.mean()) + overlap_count = int(np.sum(no_top >= has_top.min())) + print(f"mean gap (has.mean - no.mean): {mean_gap:.4f}") + print(f"hard gap (has.min - no.max): {gap:.4f} " + f"({'clean' if gap > 0 else 'OVERLAP'})") + print(f"no-tool above has-tool min: {overlap_count}/{len(no_top)}") + + print(f"\n--- 2. THRESHOLD SWEEP ---") + print(f"{'thresh':>8} {'has_acc':>8} {'no_acc':>8} {'f1':>8}") + + has_preds = has_sims.argmax(axis=1) + best_thresh = 0.0 + best_f1 = -1.0 + sweep_results = [] + + for thresh_i in range(20, 71, 5): + thresh = thresh_i / 100.0 + has_acc = float(np.mean([ + (has_top[i] >= thresh) and (TOOL_NAMES[has_preds[i]] == has_labels[i]) + for i in range(len(has_cmds)) + ])) + no_acc = float(np.mean(no_top < thresh)) + + tp = sum( + 1 for i in range(len(has_cmds)) + if has_top[i] >= thresh and TOOL_NAMES[has_preds[i]] == has_labels[i] + ) + fp = int(np.sum(no_top >= thresh)) + fn = len(has_cmds) - tp + precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0 + f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0 + + sweep_results.append((thresh, has_acc, no_acc, f1)) + print(f"{thresh:>8.2f} {has_acc:>8.3f} {no_acc:>8.3f} {f1:>8.3f}") + + if f1 > best_f1: + best_f1 = f1 + best_thresh = thresh + + print(f"\nrecommended threshold: {best_thresh:.2f} (F1={best_f1:.3f})") + + print(f"\n--- 3. PER-TOOL ACCURACY AT THRESHOLD={best_thresh:.2f} ---") + tool_correct: dict[str, int] = {t: 0 for t in TOOL_NAMES} + tool_total: dict[str, int] = {t: 0 for t in TOOL_NAMES} + confusions: list[str] = [] + + for i, (cmd, expected) in enumerate(HAS_TOOL): + tool_total[expected] += 1 + top_score = float(has_top[i]) + pred_tool = TOOL_NAMES[int(has_sims[i].argmax())] + if top_score >= best_thresh and pred_tool == expected: + tool_correct[expected] += 1 + else: + reason = ( + f"routed to {pred_tool!r} (score {top_score:.3f})" + if top_score >= best_thresh + else f"escalated (score {top_score:.3f} < {best_thresh:.2f})" + ) + confusions.append(f" [{expected}] {reason}: {cmd!r}") + + for tool in TOOL_NAMES: + n = tool_total[tool] + c = tool_correct[tool] + bar = "#" * c + "." * (n - c) + print(f" {tool:<10} {c}/{n} [{bar}]") + + if confusions: + print(f"\nfailures ({len(confusions)}):") + for line in confusions: + print(line) + + best_has_acc, best_no_acc = [ + (h, n) for t, h, n, f in sweep_results if t == best_thresh + ][0] + + print(f"\n--- 4. VERDICT ---") + print(f"mean gap: {mean_gap:.4f} hard gap: {gap:.4f} " + f"best F1: {best_f1:.3f} threshold: {best_thresh:.2f}") + print(f"has-tool acc: {best_has_acc:.3f} no-tool acc: {best_no_acc:.3f}") + + if gap > 0 and best_f1 >= 0.85: + print( + f"SOUND. Clean separation exists. Fixed threshold {best_thresh:.2f} works." + ) + elif mean_gap > 0.10 and best_f1 >= 0.70: + print( + f"MARGINAL. Moderate separation (mean gap {mean_gap:.4f}) but no clean" + f" hard cutoff. F1 {best_f1:.3f} is workable. Consider richer descriptions" + " or a calibrated threshold." + ) + else: + print( + f"NOT VIABLE. Separation too weak (mean gap {mean_gap:.4f}, F1 {best_f1:.3f})." + " The encoder space does not separate this retrieval task." + ) + + # Per-command detail + print(f"\n--- DETAILED SCORES (threshold {best_thresh:.2f}) ---") + print("\nhas-tool commands:") + print(f" {'score':>6} {'pred':>10} {'expected':>10} status command") + for i, (cmd, expected) in enumerate(HAS_TOOL): + pred = TOOL_NAMES[int(has_sims[i].argmax())] + score = float(has_top[i]) + ok = "OK " if pred == expected and score >= best_thresh else "FAIL" + print(f" {score:.4f} {pred:>10} {expected:>10} {ok} {cmd!r}") + + print("\nno-tool commands:") + print(f" {'score':>6} {'top_tool':>10} status command") + for i, cmd in enumerate(NO_TOOL): + top_tool = TOOL_NAMES[int(no_sims[i].argmax())] + score = float(no_top[i]) + ok = "OK " if score < best_thresh else "FAIL" + print(f" {score:.4f} {top_tool:>10} {ok} {cmd!r}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(custom_threshold: Optional[float] = None, mode: str = "both") -> None: + print("=== TOOL ROUTER EVAL ===") + print(f"tools: {', '.join(TOOL_NAMES)}") + print(f"has-tool examples: {len(HAS_TOOL)} no-tool examples: {len(NO_TOOL)}") + + if mode in ("finetuned", "both"): + print("\nloading fine-tuned SetFit encoder...") + ft_enc, ft_tool_embs = _load_finetuned() + _run_single_eval( + enc=ft_enc, + tool_embs=ft_tool_embs, + use_preprocess=True, + normalize=False, + label=f"fine-tuned SetFit body ({MODEL_DIR})", + ) + + if mode in ("base", "both"): + print(f"\nloading base encoder ({BASE_MODEL_ID})...") + b_enc, b_tool_embs = _load_base() + _run_single_eval( + enc=b_enc, + tool_embs=b_tool_embs, + use_preprocess=True, + normalize=True, + label=f"base {BASE_MODEL_ID}", + ) + + print("\n=== SUMMARY ===") + print( + "The fine-tuned SetFit encoder is optimized for 5-class intent classification" + " via contrastive training. That training destroys cross-domain semantic" + " similarity: command-to-tool-description cosines collapse to ~0. The" + " fine-tuned encoder CANNOT be reused for retrieval as-is." + ) + print( + "The base encoder (BAAI/bge-small-en-v1.5) retains strong semantic" + " relationships and produces clean separation between has-tool and no-tool" + " commands at a fixed threshold. See the base encoder results above for the" + " recommended threshold and per-tool accuracy." + ) + print( + "Recommended architecture: use the fine-tuned SetFit model for intent" + " classification (integration vs. other), then run the base encoder for" + " tool retrieval within the integration branch. The base encoder is only" + " 33MB and can be loaded alongside the SetFit model with negligible overhead." + ) + + +if __name__ == "__main__": + args = sys.argv[1:] + custom = float(args[0]) if args and args[0].replace(".", "", 1).isdigit() else None + mode_arg = "both" + for a in args: + if a in ("base", "finetuned", "both"): + mode_arg = a + main(custom, mode_arg) diff --git a/experiments/tools_proto.py b/experiments/tools_proto.py new file mode 100644 index 0000000..d8a564b --- /dev/null +++ b/experiments/tools_proto.py @@ -0,0 +1,320 @@ +"""Approach A prototype: classify integration commands to tool labels. + +Variant trained: FINETUNE only (fresh SetFit on bge-small base). +The frozen-head variant is skipped; it failed in earlier runs. + +Data source: all 211 rows of data/integration.jsonl, relabeled by keyword rules. +Split: stratified 80/20 by tool label, seeded at 42. +The main evals/holdout.jsonl is NOT used here. + +Usage: + .venv/bin/python -m routelet.tools_proto +""" + +from __future__ import annotations + +import json +import random +from collections import Counter +from pathlib import Path + +import numpy as np +from sklearn.metrics import classification_report, confusion_matrix +from sklearn.model_selection import train_test_split + +from routelet.preprocess import preprocess + +PROJECT_ROOT = Path(__file__).parent.parent.parent +INTEGRATION_DATA = PROJECT_ROOT / "data" / "integration.jsonl" +PROTO_MODEL_OUT = PROJECT_ROOT / "models" / "setfit_tools_proto" + +TOOLS = ["spotify", "gmail", "github", "youtube", "no_tool"] + +# --------------------------------------------------------------------------- +# Keyword-based relabeler +# Rules (applied in priority order): +# youtube - explicit "youtube" mention +# github - github/PR/pull request/issue/repo/branch/commit/CI/fork/star/clone +# gmail - email/inbox/compose/reply/forward/archive in email context +# spotify - music playback controls, play , spotify, podcast playback +# no_tool - everything else +# +# Tricky cases preserved from old comments: +# "play the next episode on netflix" -> no_tool (no netflix integration) +# "queue up the next episode of the office" -> no_tool (TV, not spotify) +# "skip ahead like a minute in this podcast" -> spotify (playback control) +# "crank the bass up on this song" -> spotify (playback control) +# "play that song on youtube" -> youtube (explicit youtube) +# --------------------------------------------------------------------------- + +import re + +_YOUTUBE_KW = re.compile(r"\byoutube\b", re.I) + +_GITHUB_KW = re.compile( + r"\b(github|pull request|pull requests|open pr|open prs|my prs|my pr|merge pr|" + r"draft pr|draft pull|approve.*pr|pr from|pr for|" + r"pull request from|pull request for|" + r"repo\b|repos\b|branch\b|branches\b|commit\b|commits\b|" + r"issue\b|issues\b|fork\b|clone\b|CI\b|workflow\b|" + r"github action|push.*change|push.*github|latest from github|" + r"github build|github star|github notification)\b", + re.I, +) +# "star" alone is too broad (hits "star the email"); only match "star" when paired with github context +_GITHUB_STAR = re.compile(r"\bstar\b.*\b(repo|github)\b|\b(github|repo)\b.*\bstar\b", re.I) + +# Email keywords: cover all the ways people say "email/inbox" tasks +_GMAIL_KW = re.compile( + r"\b(email\b|inbox\b|unread emails?|new emails?|read.*email|check.*email|" + r"send.*email|compose.*email|search.*email|search.*inbox|" + r"reply.*email|forward.*email|archive.*email|delete.*email|" + r"mark.*email|star.*email|draft.*email|newest email|latest email|" + r"email to|email from|do i have.*email|any.*email|" + r"my emails|promotional email|newsletter email|archive.*email|" + r"emails?$)\b", + re.I, +) + +# Spotify: play/skip/pause/queue/shuffle/volume controls for music/podcast +# Explicit "spotify" mention, or music-playback verbs with music targets +_SPOTIFY_EXPLICIT = re.compile(r"\bspotify\b", re.I) +_SPOTIFY_PLAY = re.compile( + r"\b(play|queue up|throw on|put on|blast|fire up)\b", re.I +) +_SPOTIFY_MUSIC_TARGET = re.compile( + r"\b(song|album|track|music|playlist|artist|band|jazz|lofi|" + r"hip hop|beats|shuffle|podcast|episode of my podcast|" + r"chill|liked songs)\b", + re.I, +) +# Match "play/queue up " at sentence start (likely a song/artist name). +# "add ... to my playlist/queue" is also spotify if it has a playlist/queue target. +_SPOTIFY_PLAY_NAMED = re.compile( + r"^(play|queue up)\b", + re.I, +) +_SPOTIFY_ADD_QUEUE = re.compile( + r"\badd\b.+\b(playlist|queue)\b", + re.I, +) +_SPOTIFY_CONTROL = re.compile( + r"\b(pause|skip|next track|previous track|go back to the previous|" + r"rewind|fast forward|volume|raise the volume|lower the volume|" + r"turn.*volume|set.*volume|increase.*volume|crank.*up|" + r"what.s playing|shuffle my|repeat)\b", + re.I, +) +# Hard-exclude: netflix episode, tv show episode (not podcast) +_SPOTIFY_EXCLUDE = re.compile( + r"\b(netflix|episode of the|next episode of|tv show|hulu)\b", re.I +) + + +def _relabel(text: str) -> str: + """Return tool label for a single integration.jsonl row text.""" + t = text.strip() + + # Priority 1: explicit youtube + if _YOUTUBE_KW.search(t): + return "youtube" + + # Priority 2: github signals (star needs explicit github context) + if _GITHUB_KW.search(t) or _GITHUB_STAR.search(t): + return "github" + + # Priority 3: gmail signals + if _GMAIL_KW.search(t): + return "gmail" + + # Priority 4: spotify + if _SPOTIFY_EXCLUDE.search(t): + return "no_tool" + if _SPOTIFY_EXPLICIT.search(t): + return "spotify" + if _SPOTIFY_CONTROL.search(t): + return "spotify" + if _SPOTIFY_PLAY.search(t) and _SPOTIFY_MUSIC_TARGET.search(t): + return "spotify" + # "play/queue up " at start -> likely music + if _SPOTIFY_PLAY_NAMED.match(t) and not _SPOTIFY_EXCLUDE.search(t): + return "spotify" + # "add to my playlist/queue" -> spotify + if _SPOTIFY_ADD_QUEUE.search(t): + return "spotify" + + return "no_tool" + + +def load_and_relabel() -> list[tuple[str, str]]: + """Load integration.jsonl and return (text, tool_label) pairs.""" + rows = [] + with open(INTEGRATION_DATA) as f: + for line in f: + line = line.strip() + if not line: + continue + obj = json.loads(line) + text = obj["text"] + label = _relabel(text) + rows.append((text, label)) + return rows + + +def _counts(labels: list[str]) -> str: + return " ".join(f"{t}={labels.count(t)}" for t in TOOLS) + + +def _print_report(name: str, te_labels: list[str], preds: list[str], te_texts: list[str]) -> None: + print(f"\n{'=' * 60}") + print(f"VARIANT: {name}") + print(f"{'=' * 60}") + print(classification_report(te_labels, preds, labels=TOOLS, zero_division=0)) + + cm = confusion_matrix(te_labels, preds, labels=TOOLS) + col_w = max(len(t) for t in TOOLS) + 2 + header = " " * col_w + "".join(f"{t:>{col_w}}" for t in TOOLS) + print("confusion matrix (rows=true, cols=pred):") + print(header) + for i, row_lbl in enumerate(TOOLS): + row_str = f"{row_lbl:<{col_w}}" + "".join(f"{v:>{col_w}}" for v in cm[i]) + print(row_str) + + nt_true = [lbl == "no_tool" for lbl in te_labels] + nt_pred_as_nt = [p == "no_tool" for p in preds] + nt_correct = sum(a and b for a, b in zip(nt_true, nt_pred_as_nt)) + nt_total = sum(nt_true) + nt_recall = nt_correct / nt_total if nt_total else 0.0 + print(f"\nno_tool recall: {nt_recall:.3f} ({nt_correct}/{nt_total})") + + acc = sum(p == t for p, t in zip(preds, te_labels)) / len(te_labels) + print(f"overall accuracy: {acc:.3f} ({sum(p == t for p, t in zip(preds, te_labels))}/{len(te_labels)})") + + errors = [ + (te_labels[i], preds[i], te_texts[i]) + for i in range(len(te_labels)) + if preds[i] != te_labels[i] + ] + if errors: + print(f"\nmisclassifications ({len(errors)}):") + for true_lbl, pred_lbl, text in errors: + print(f" true={true_lbl:<10} pred={pred_lbl:<10} text={text!r}") + + +# --------------------------------------------------------------------------- +# Variant: FINETUNE SetFit with tool labels (bge-small base) +# --------------------------------------------------------------------------- + +def run_finetune( + tr_texts: list[str], + tr_labels: list[str], + te_texts: list[str], + te_labels: list[str], +) -> None: + print("\n" + "=" * 60) + print("VARIANT: FINETUNE SetFit (bge-small base) with tool labels") + print("=" * 60) + + import torch + from datasets import Dataset + from setfit import SetFitModel, Trainer, TrainingArguments + + random.seed(0) + np.random.seed(0) + torch.manual_seed(0) + + BASE = "BAAI/bge-small-en-v1.5" + + print(f"train: {len(tr_texts)} | {_counts(tr_labels)}") + print(f"test: {len(te_texts)} | {_counts(te_labels)}") + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"training on {device}") + + proc_tr = [preprocess(t) for t in tr_texts] + proc_te = [preprocess(t) for t in te_texts] + + train_ds = Dataset.from_dict({"text": proc_tr, "label": tr_labels}) + + model = SetFitModel.from_pretrained( + BASE, + labels=TOOLS, + device=device, + head_params={"class_weight": "balanced"}, + ) + + args = TrainingArguments( + batch_size=16, + num_epochs=2, + sampling_strategy="unique", + ) + + trainer = Trainer(model=model, args=args, train_dataset=train_ds) + trainer.train() + + preds_raw = model.predict(proc_te) + preds = [str(p) for p in preds_raw] + + _print_report("FINETUNE", te_labels, preds, te_texts) + + PROTO_MODEL_OUT.mkdir(parents=True, exist_ok=True) + model.save_pretrained(str(PROTO_MODEL_OUT)) + print(f"\nsaved proto model to {PROTO_MODEL_OUT}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> None: + print("=== APPROACH A PROTOTYPE: TOOL CLASSIFICATION (stratified split) ===") + print(f"tools: {TOOLS}") + print(f"data: {INTEGRATION_DATA}") + + # Step 1: relabel + all_data = load_and_relabel() + all_texts = [t for t, _ in all_data] + all_labels = [lbl for _, lbl in all_data] + + print(f"\nRelabel distribution ({len(all_data)} total rows):") + dist = Counter(all_labels) + for tool in TOOLS: + print(f" {tool:<12} {dist[tool]:>4}") + + # Spot-check: print any rows that might be misfiring + print("\nSpot-check: first 3 per class from relabeled set") + seen: dict[str, int] = {t: 0 for t in TOOLS} + for text, lbl in all_data: + if seen[lbl] < 3: + print(f" {lbl:<12} {text!r}") + seen[lbl] += 1 + + # Check for thin classes + thin = [t for t in TOOLS if dist[t] < 10] + if thin: + print(f"\nWARNING: thin classes (< 10 examples): {thin}") + else: + print("\nAll classes have >= 10 examples. OK to split.") + + # Step 2: stratified 80/20 split + tr_texts, te_texts, tr_labels, te_labels = train_test_split( + all_texts, all_labels, + test_size=0.20, + stratify=all_labels, + random_state=42, + ) + + print(f"\nStratified split (seed=42):") + print(f" train {len(tr_texts)} rows: {_counts(tr_labels)}") + print(f" test {len(te_texts)} rows: {_counts(te_labels)}") + + # Step 3 + 4: train and eval finetune variant + run_finetune(tr_texts, tr_labels, te_texts, te_labels) + + print("\n" + "=" * 60) + print("PROTOTYPE COMPLETE") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 5bc2739..42d7ab1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,9 +4,6 @@ version = "0.0.1" description = "A tiny fine-tuned intent router. Classify text into intents in milliseconds, on-device." requires-python = ">=3.12" dependencies = [ - "fastapi", - "uvicorn", - "pydantic", "scikit-learn", ] @@ -16,6 +13,12 @@ train = [ "datasets", "torch", ] +# Only needed if you build out the optional HTTP endpoint (src/routelet/serve.py). +serve = [ + "fastapi", + "uvicorn", + "pydantic", +] [build-system] requires = ["hatchling"] @@ -27,6 +30,8 @@ packages = ["src/routelet"] [tool.ruff] line-length = 100 target-version = "py312" +# Throwaway prototypes, not held to the library's lint bar. +extend-exclude = ["experiments"] [tool.ruff.lint] select = ["E", "F", "I"] diff --git a/report/baselines.json b/report/baselines.json new file mode 100644 index 0000000..42deb68 --- /dev/null +++ b/report/baselines.json @@ -0,0 +1,42 @@ +{ + "_note": "Cached Claude baseline so report.py does not re-spend on the paid API every regen. Regenerate by running Scripts/refresh_haiku_baseline.py with ANTHROPIC_API_KEY against the current evals/holdout.jsonl.", + "haiku": { + "model": "claude-haiku-4-5", + "eval_n": 200, + "accuracy": 0.935, + "latency_p50_ms": 1007.0, + "latency_p95_ms": 1865.0, + "per_class": { + "find_action": { + "precision": 0.95, + "recall": 1.0, + "f1": 0.98, + "support": 40 + }, + "integration": { + "precision": 1.0, + "recall": 0.72, + "f1": 0.84, + "support": 39 + }, + "chat": { + "precision": 0.93, + "recall": 0.95, + "f1": 0.94, + "support": 40 + }, + "memory": { + "precision": 0.91, + "recall": 1.0, + "f1": 0.95, + "support": 40 + }, + "agent": { + "precision": 0.91, + "recall": 1.0, + "f1": 0.95, + "support": 41 + } + } + } +} diff --git a/report/metrics.json b/report/metrics.json new file mode 100644 index 0000000..46994d4 --- /dev/null +++ b/report/metrics.json @@ -0,0 +1,49 @@ +{ + "eval_n": 200, + "tfidf": { + "accuracy": 0.69 + }, + "setfit": { + "accuracy": 0.895 + }, + "haiku": { + "model": "claude-haiku-4-5", + "eval_n": 200, + "accuracy": 0.935, + "latency_p50_ms": 1007.0, + "latency_p95_ms": 1865.0, + "per_class": { + "find_action": { + "precision": 0.95, + "recall": 1.0, + "f1": 0.98, + "support": 40 + }, + "integration": { + "precision": 1.0, + "recall": 0.72, + "f1": 0.84, + "support": 39 + }, + "chat": { + "precision": 0.93, + "recall": 0.95, + "f1": 0.94, + "support": 40 + }, + "memory": { + "precision": 0.91, + "recall": 1.0, + "f1": 0.95, + "support": 40 + }, + "agent": { + "precision": 0.91, + "recall": 1.0, + "f1": 0.95, + "support": 41 + } + }, + "stale": false + } +} diff --git a/report/model_comparison.png b/report/model_comparison.png new file mode 100644 index 0000000..ec509ca Binary files /dev/null and b/report/model_comparison.png differ diff --git a/report/report.py b/report/report.py new file mode 100644 index 0000000..49ff31b --- /dev/null +++ b/report/report.py @@ -0,0 +1,159 @@ +"""Reproducible performance report. Scores the models on the frozen holdout and +writes figures plus metrics.json. Re-run after every retrain. + +Scoring rules (docs/metrics.md): + * Every model is scored on the same frozen evals/holdout.jsonl. + * Each model is fed text the way it was trained: the TF-IDF baseline on raw + text, the SetFit model on preprocess()'d text. + * Claude Haiku numbers come from report/baselines.json (cached so we don't + re-spend on the paid API every regen). If that cache was measured on a + different holdout size than the current one, it is flagged, not silently + plotted as comparable. + +matplotlib only, installed into the venv (not pyproject). +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import joblib +import matplotlib +import numpy as np +from sklearn.metrics import accuracy_score + +matplotlib.use("Agg") # headless: write PNGs, never open a window +import matplotlib.pyplot as plt # noqa: E402 + +PROJECT_ROOT = Path(__file__).parent.parent +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +from routelet.data import load # noqa: E402 +from routelet.preprocess import preprocess # noqa: E402 + +HOLDOUT = PROJECT_ROOT / "evals" / "holdout.jsonl" +TFIDF_MODEL = PROJECT_ROOT / "models" / "baseline.joblib" +SETFIT_DIR = PROJECT_ROOT / "models" / "setfit" +BASELINES = PROJECT_ROOT / "report" / "baselines.json" +OUT_DIR = PROJECT_ROOT / "report" + +# Okabe-Ito (colorblind-safe). Gray baseline, green hero (routelet), amber oracle. +C_TFIDF = "#999999" +C_ROUTELET = "#009E73" +C_HAIKU = "#E69F00" + + +def score_tfidf(texts: list[str], gold: list[str]) -> float: + """TF-IDF + logistic regression baseline. Trained on raw text, so score on + raw text (no preprocess) to match its training.""" + model = joblib.load(TFIDF_MODEL) + preds = model.predict(texts) + return float(accuracy_score(gold, preds)) + + +def score_setfit(texts: list[str], gold: list[str]) -> float: + """The shipped SetFit model. Trained on preprocess()'d text, so score the + same way. This is the fine-tuned bge-small body + LR head in torch; int8 + ONNX (what Rust actually runs) is verified equivalent at export time.""" + from setfit import SetFitModel + + model = SetFitModel.from_pretrained(str(SETFIT_DIR)) + preds = list(model.predict([preprocess(t) for t in texts])) + return float(accuracy_score(gold, preds)) + + +def load_haiku(eval_n: int) -> dict | None: + """Read the cached Claude baseline. Returns the record with a `stale` flag + set when it was measured on a different holdout size than the current one.""" + if not BASELINES.exists(): + return None + data = json.loads(BASELINES.read_text()) + haiku = data.get("haiku") + if not haiku: + return None + haiku["stale"] = haiku.get("eval_n") != eval_n + return haiku + + +def plot_model_comparison(metrics: dict) -> Path: + """Figure 1: accuracy on the frozen holdout, baseline vs shipped vs oracle.""" + bars = [ + ("TF-IDF\nbaseline", metrics["tfidf"]["accuracy"], C_TFIDF, "raw text, in-proc"), + ("routelet\n(SetFit, on-device)", metrics["setfit"]["accuracy"], C_ROUTELET, + "on-device, free"), + ] + haiku = metrics.get("haiku") + if haiku: + note = f"{haiku['latency_p50_ms']:.0f}ms p50, cloud" + if haiku.get("stale"): + note += f"\n(n={haiku['eval_n']}, stale)" + bars.append(("Claude Haiku\n(LLM oracle)", haiku["accuracy"], C_HAIKU, note)) + + labels = [b[0] for b in bars] + accs = [b[1] for b in bars] + colors = [b[2] for b in bars] + notes = [b[3] for b in bars] + + fig, ax = plt.subplots(figsize=(7, 4.5)) + x = np.arange(len(bars)) + ax.bar(x, accs, color=colors, width=0.62, zorder=3) + + for i, (acc, note) in enumerate(zip(accs, notes)): + ax.text(i, acc + 0.015, f"{acc:.0%}", ha="center", va="bottom", + fontsize=12, fontweight="bold") + ax.text(i, 0.04, note, ha="center", va="bottom", fontsize=8, + color="white", fontweight="medium") + + ax.set_xticks(x) + ax.set_xticklabels(labels, fontsize=9) + ax.set_ylim(0, 1.08) + ax.set_ylabel("Accuracy on frozen holdout") + ax.set_title(f"Intent routing accuracy (n={metrics['eval_n']})", fontsize=12, fontweight="bold") + ax.yaxis.grid(True, linestyle="--", alpha=0.4, zorder=0) + ax.set_axisbelow(True) + for spine in ("top", "right"): + ax.spines[spine].set_visible(False) + fig.tight_layout() + + out = OUT_DIR / "model_comparison.png" + fig.savefig(out, dpi=150) + plt.close(fig) + return out + + +def main() -> None: + examples = load(HOLDOUT) + texts = [e.text for e in examples] + gold = [e.intent.value for e in examples] + eval_n = len(examples) + + print(f"scoring on {eval_n}-row holdout") + tfidf_acc = score_tfidf(texts, gold) + print(f" TF-IDF {tfidf_acc:.3f}") + setfit_acc = score_setfit(texts, gold) + print(f" SetFit {setfit_acc:.3f}") + + metrics: dict = { + "eval_n": eval_n, + "tfidf": {"accuracy": tfidf_acc}, + "setfit": {"accuracy": setfit_acc}, + } + haiku = load_haiku(eval_n) + if haiku: + metrics["haiku"] = haiku + if haiku["stale"]: + print(f" Haiku {haiku['accuracy']:.3f} [STALE: cached on n={haiku['eval_n']}, " + f"current holdout is n={eval_n}; refresh with routelet.evaluate]") + else: + print(f" Haiku {haiku['accuracy']:.3f} (cached)") + + (OUT_DIR / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") + fig = plot_model_comparison(metrics) + print(f"wrote {fig}") + print(f"wrote {OUT_DIR / 'metrics.json'}") + + +if __name__ == "__main__": + main() diff --git a/src/routelet/augment.py b/src/routelet/augment.py new file mode 100644 index 0000000..9084863 --- /dev/null +++ b/src/routelet/augment.py @@ -0,0 +1,183 @@ +"""Disfluency augmenter: turn clean training utterances into voice-like variants. + +Real Aegis input is spoken, so it is disfluent: fillers ("uhh", "like"), hedges +("i was wondering if you could"), dropped apostrophes, self-corrections, run-ons. +The clean synthetic training set does not look like that, which is the documented +distribution gap behind the chat->integration and agent->integration leaks. + +This module applies conservative, label-preserving string transforms to existing +labeled rows and writes the results to a SEPARATE file, data/augmented.jsonl, so +the augmentation is ablatable (delete the file, retrain, compare). train_setfit +globs data/*.jsonl, so the file is picked up automatically. + +Run: .venv/bin/python -m routelet.augment +""" + +import json +import random +import re +from pathlib import Path + +from routelet.data import Intent, load + +DATA_DIR = Path("data") +HOLDOUT = Path("evals/holdout.jsonl") +OUT = DATA_DIR / "augmented.jsonl" +SOURCES = ["agent", "chat", "find_action", "integration", "memory"] +VARIANTS_PER_ROW = 2 + +# Chaining one utterance to another with "and then" turns one action into two, +# which flips integration/find_action into agent. So same-class chaining is only +# label-safe for classes where two instances stay the same intent: chat (still +# chat), memory (still memory), and agent (already multi-step). It is NOT safe +# for integration or find_action. +CHAINABLE = {Intent.CHAT, Intent.MEMORY, Intent.AGENT} + +FILLERS = ["uh", "uhh", "um", "like", "you know", "i mean"] +HEDGES = ["can you", "i was wondering if you could", "hey can u", "could you"] +TAILS = ["plz", "thanks", "please", "thx"] + + +def _norm(text: str) -> str: + """Loose key for dedup: lowercase, drop apostrophes and punctuation.""" + text = text.lower().strip().replace("'", "") + text = re.sub(r"[^a-z0-9 ]", " ", text) + return re.sub(r"\s+", " ", text).strip() + + +def _drop_apostrophes(text: str) -> str: + return text.replace("'", "") + + +def _insert_filler(text: str, rng: random.Random) -> str: + """Drop one filler at a random word boundary (not the very end).""" + words = text.split() + if len(words) < 2: + return text + pos = rng.randint(1, len(words) - 1) + words.insert(pos, rng.choice(FILLERS)) + return " ".join(words) + + +def _repeat_word(text: str, rng: random.Random) -> str: + """~10% chance to stutter-repeat one word ('the the').""" + words = text.split() + if len(words) < 2 or rng.random() >= 0.10: + return text + pos = rng.randint(0, len(words) - 1) + words.insert(pos, words[pos]) + return " ".join(words) + + +def _prepend_hedge(text: str, rng: random.Random) -> str: + hedge = rng.choice(HEDGES) + return f"{hedge} {text}" + + +def _append_tail(text: str, rng: random.Random) -> str: + return f"{text} {rng.choice(TAILS)}" + + +def _drop_punct(text: str) -> str: + return text.rstrip(".?!").rstrip() + + +def make_variant(text: str, intent: Intent, rng: random.Random) -> str: + """Build one disfluent, label-preserving variant of a clean utterance. + + Each transform fires probabilistically so the two variants of a row differ. + Order matters: hedges/tails wrap the sentence, fillers/stutters go inside. + """ + out = text + if rng.random() < 0.85: + out = _insert_filler(out, rng) + out = _repeat_word(out, rng) + if rng.random() < 0.70: + out = _drop_apostrophes(out) + if rng.random() < 0.45: + out = _prepend_hedge(out, rng) + if rng.random() < 0.35: + out = _append_tail(out, rng) + if rng.random() < 0.40: + out = _drop_punct(out) + return re.sub(r"\s+", " ", out).strip() + + +def make_chain( + text: str, partner: str, intent: Intent, rng: random.Random +) -> str | None: + """For chainable classes, join two same-class utterances with 'and then'. + + Returns None when chaining is not label-safe for this intent. The result is + then lightly disfluent so it does not read as a clean concatenation. + """ + if intent not in CHAINABLE: + return None + a = _drop_apostrophes(text) if rng.random() < 0.6 else text + b = _drop_apostrophes(partner) if rng.random() < 0.6 else partner + joined = f"{_drop_punct(a)} and then {_drop_punct(b)}" + if rng.random() < 0.5: + joined = _insert_filler(joined, rng) + return re.sub(r"\s+", " ", joined).strip() + + +def augment() -> dict[str, int]: + """Generate ~2 disfluent variants per source row into data/augmented.jsonl. + + Dedups generated variants against the frozen holdout and against each other, + and skips any variant that collapses back to its source form. Returns a + per-class count of rows written. + """ + rng = random.Random() + rng.seed(0) + + blocked = {_norm(e.text) for e in load(HOLDOUT)} + written: list[tuple[str, str]] = [] + counts: dict[str, int] = {s: 0 for s in SOURCES} + + for name in SOURCES: + rows = load(DATA_DIR / f"{name}.jsonl") + intent = Intent(name) + texts = [e.text for e in rows] + for i, e in enumerate(rows): + made = 0 + attempts = 0 + while made < VARIANTS_PER_ROW and attempts < 12: + attempts += 1 + # Roughly 1 in 4 variants for a chainable class is a same-class + # "and then" run-on; the rest are single disfluent rephrasings. + if intent in CHAINABLE and len(texts) > 1 and rng.random() < 0.25: + partner = rng.choice(texts) + if partner == e.text: + continue + variant = make_chain(e.text, partner, intent, rng) + if variant is None: + continue + else: + variant = make_variant(e.text, intent, rng) + key = _norm(variant) + if not key or key == _norm(e.text): + continue + if key in blocked: + continue + blocked.add(key) + written.append((variant, name)) + counts[name] += 1 + made += 1 + + with OUT.open("w") as f: + for text, intent in written: + f.write(json.dumps({"text": text, "intent": intent}) + "\n") + return counts + + +def main() -> None: + counts = augment() + total = sum(counts.values()) + print(f"wrote {total} augmented rows to {OUT}") + for name in SOURCES: + print(f" {name}: {counts[name]}") + + +if __name__ == "__main__": + main() diff --git a/src/routelet/data.py b/src/routelet/data.py index b6f7c89..43538b6 100644 --- a/src/routelet/data.py +++ b/src/routelet/data.py @@ -1,7 +1,7 @@ """The Intent label schema and dataset prep for intent routing. Reads labeled JSONL, one ``{"text", "intent"}`` object per line, into the -deterministic train/test splits that train.py and evaluate.py consume. +validated Example lists that train.py and evaluate.py consume. """ import json @@ -9,8 +9,6 @@ from enum import StrEnum from pathlib import Path -from sklearn.model_selection import train_test_split - class Intent(StrEnum): """The label set, frozen. Definitions and boundary rules in docs/taxonomy.md. @@ -64,26 +62,3 @@ def load_dir(directory: str | Path) -> list[Example]: for path in sorted(Path(directory).glob("*.jsonl")): examples.extend(load(path)) return examples - - -def split( - examples: list[Example], - test_size: float = 0.25, - seed: int = 0, -) -> tuple[list[Example], list[Example]]: - """Deterministic, stratified train/test split. - - Stratified so each intent keeps its proportion in both halves. With classes - this small an unstratified split can leave an intent with zero test rows and - silently blind the eval to it. ``seed`` fixes the partition so accuracy is - comparable across runs. Raises ValueError if any intent has fewer than 2 - examples, since it can't land in both halves. - - This splits the hand-labeled pool. The real frozen eval set will live in - evals/, sourced apart from the training data. - """ - labels = [e.intent for e in examples] - train, test = train_test_split( - examples, test_size=test_size, random_state=seed, stratify=labels - ) - return train, test diff --git a/src/routelet/evaluate.py b/src/routelet/evaluate.py index f4b52f9..0deaedb 100644 --- a/src/routelet/evaluate.py +++ b/src/routelet/evaluate.py @@ -19,33 +19,13 @@ from routelet.data import Intent, load +# The prompt and schema are the teacher's, shared with the ingest labeler so the +# eval measures the same classifier that labels the training data. +from routelet.teacher import INTENT_SCHEMA, SYSTEM + MODEL = "claude-haiku-4-5" EVAL_FILE = "evals/holdout.jsonl" -# Condensed from docs/taxonomy.md. Kept stable so it caches across the per-command -# calls (though a prompt this short may sit under the model's min cacheable prefix). -SYSTEM = """You are an intent router. Classify each user command into exactly one of five intents: - -- find_action: locate or operate a UI element on the current screen. -- integration: one discrete action against an app or service. -- chat: answer from general knowledge or conversation. No app action, no personal data. -- memory: store or recall a personal fact (storing needs an explicit remember/note/save). -- agent: a task needing two or more chained steps or a plan. - -Tie-breakers for the tricky cases: -- Two or more chained actions ("X and then Y") is agent, not integration. -- A question answered from a stored personal fact is memory; from world knowledge it is chat. -- A UI verb (click/tap/scroll) on a named element is find_action, even if an app is named. -- Playback controls (skip, pause, next, volume) are integration, not find_action, unless a button is named. -- "explain how to..." or "talk me through..." is chat, even when it names an app action.""" - -INTENT_SCHEMA = { - "type": "object", - "properties": {"intent": {"type": "string", "enum": [i.value for i in Intent]}}, - "required": ["intent"], - "additionalProperties": False, -} - def main() -> None: client = anthropic.Anthropic() diff --git a/src/routelet/export_onnx.py b/src/routelet/export_onnx.py new file mode 100644 index 0000000..dc01719 --- /dev/null +++ b/src/routelet/export_onnx.py @@ -0,0 +1,250 @@ +"""Export the trained SetFit encoder to ONNX and dump the LR head to JSON. + +SetFit 1.1.3's bundled export_onnx fuses the sklearn head via skl2onnx, which +emits opset 13 while torch exports the encoder at opset 18. onnx.merge_models +refuses to combine them, so the fused path is broken on this stack. + +Instead we split the model: + embedder.onnx fp32 BERT encoder with CLS pooling and L2 normalization + embedder.int8.onnx dynamic int8 quantization of the above + head.json LR coef/intercept/labels for the Rust consumer to apply + tokenizer.json copied from models/setfit/ for the Rust tokenizer + +The Rust consumer runs: embedding = onnx(tokens); logits = coef @ emb + intercept +then argmax -> labels[i]. +""" + +import json +import shutil +import warnings +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import torch +import torch.nn as nn +import transformers +from onnxruntime.quantization import QuantType, quantize_dynamic +from setfit import SetFitModel + +from routelet.data import load +from routelet.preprocess import preprocess + +SETFIT_DIR = "models/setfit" +EMBEDDER_OUT = "models/embedder.onnx" +EMBEDDER_INT8_OUT = "models/embedder.int8.onnx" +HEAD_OUT = "models/head.json" +TOKENIZER_OUT = "models/tokenizer.json" +EVAL_FILE = "evals/holdout.jsonl" +TEMPERATURE_FILE = "models/setfit/temperature.json" +OPSET = 14 + + +class EncoderWithPooling(nn.Module): + """BERT encoder + CLS pooling + L2 normalization in one exportable module. + + Derives the pooling mode from the SentenceTransformer's Pooling module at + construction time rather than hard-coding it. For bge-small-en-v1.5 the + trained model uses CLS pooling followed by L2 normalization (modules 1 and + 2 in the ST pipeline), so the wrapper takes last_hidden_state[:, 0, :] and + normalizes it. If the Pooling module reports a different mode a ValueError + is raised at export time rather than silently producing wrong embeddings. + """ + + def __init__(self, sentence_transformer: object) -> None: + super().__init__() + pooling_mod = sentence_transformer[1] + mode = pooling_mod.pooling_mode + if mode != "cls": + raise ValueError( + f"EncoderWithPooling only implements CLS pooling; " + f"found pooling_mode={mode!r}. Update the wrapper if the " + "model uses a different pooling strategy." + ) + self.bert = sentence_transformer[0].auto_model + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + token_type_ids: torch.Tensor, + ) -> torch.Tensor: + out = self.bert( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + ) + # CLS token is at position 0 for right-padded BERT inputs. + cls = out.last_hidden_state[:, 0, :] + # L2 normalize so the embedding lives on the unit sphere, matching + # the Normalize module in the SentenceTransformer pipeline. + normalized = cls / cls.norm(dim=1, keepdim=True).clamp(min=1e-12) + return normalized + + +def _export_embedder(wrapper: EncoderWithPooling, out_path: str) -> None: + dummy_ids = torch.zeros(1, 16, dtype=torch.int64) + dummy_mask = torch.ones(1, 16, dtype=torch.int64) + dummy_types = torch.zeros(1, 16, dtype=torch.int64) + + dynamic_axes = { + "input_ids": {0: "batch", 1: "seq"}, + "attention_mask": {0: "batch", 1: "seq"}, + "token_type_ids": {0: "batch", 1: "seq"}, + "embedding": {0: "batch"}, + } + + Path(out_path).parent.mkdir(exist_ok=True) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=DeprecationWarning) + torch.onnx.export( + wrapper, + (dummy_ids, dummy_mask, dummy_types), + out_path, + dynamo=False, + opset_version=OPSET, + input_names=["input_ids", "attention_mask", "token_type_ids"], + output_names=["embedding"], + dynamic_axes=dynamic_axes, + do_constant_folding=True, + ) + print(f"[export] legacy TorchScript exporter, opset {OPSET}") + print(f"[export] saved {out_path}") + + +def _print_session_info(sess: ort.InferenceSession) -> None: + print("\nembedder.onnx I/O signature:") + for inp in sess.get_inputs(): + print(f" input {inp.name!r:20s} dtype={inp.type!s:20s} shape={inp.shape}") + for out in sess.get_outputs(): + print(f" output {out.name!r:20s} dtype={out.type!s:20s} shape={out.shape}") + + +def _run_onnx_preds( + sess: ort.InferenceSession, + tokenizer: transformers.PreTrainedTokenizerBase, + texts: list[str], + coef: np.ndarray, + intercept: np.ndarray, + labels: list[str], +) -> list[str]: + encoded = tokenizer( + texts, + padding="max_length", + truncation=True, + max_length=128, + return_tensors="np", + ) + feed = { + "input_ids": encoded["input_ids"].astype(np.int64), + "attention_mask": encoded["attention_mask"].astype(np.int64), + "token_type_ids": encoded["token_type_ids"].astype(np.int64), + } + (emb,) = sess.run(["embedding"], feed) + logits = emb @ coef.T + intercept + return [labels[int(np.argmax(row))] for row in logits] + + +def main() -> None: + model = SetFitModel.from_pretrained(SETFIT_DIR) + body = model.model_body.to("cpu") + head = model.model_head + + # Inspect and confirm the ST pooling/normalize pipeline. + pooling_mode = body[1].pooling_mode + print(f"ST pooling mode: {pooling_mode}") + print(f"ST modules: {[type(m).__name__ for m in body]}") + + # Build the exportable wrapper. + wrapper = EncoderWithPooling(body) + wrapper.eval() + + # Export fp32 ONNX. + _export_embedder(wrapper, EMBEDDER_OUT) + + # Quantize to int8. + quantize_dynamic(EMBEDDER_OUT, EMBEDDER_INT8_OUT, weight_type=QuantType.QInt8) + print(f"[quantize] saved {EMBEDDER_INT8_OUT}") + + # Load fitted temperature T from the file written by train_setfit.py. + temp_path = Path(TEMPERATURE_FILE) + if temp_path.exists(): + temperature = json.loads(temp_path.read_text())["temperature"] + print(f"[temperature] loaded T={temperature:.4f} from {TEMPERATURE_FILE}") + else: + temperature = 1.0 + print(f"[temperature] {TEMPERATURE_FILE} not found, using T=1.0 (no scaling)") + + # Dump LR head to JSON. classes_ order is what argmax maps to. + labels = head.classes_.tolist() + head_data = { + "coef": head.coef_.tolist(), + "intercept": head.intercept_.tolist(), + "labels": labels, + "temperature": temperature, + } + Path(HEAD_OUT).write_text(json.dumps(head_data, indent=2) + "\n") + print( + f"[head] saved {HEAD_OUT} (coef {head.coef_.shape}, " + f"labels {labels}, T={temperature:.4f})" + ) + + # Copy tokenizer for the Rust consumer. + shutil.copy(f"{SETFIT_DIR}/tokenizer.json", TOKENIZER_OUT) + print(f"[tokenizer] copied -> {TOKENIZER_OUT}") + + # Print I/O signature. + fp32_sess = ort.InferenceSession(EMBEDDER_OUT, providers=["CPUExecutionProvider"]) + _print_session_info(fp32_sess) + + # File sizes. + fp32_bytes = Path(EMBEDDER_OUT).stat().st_size + int8_bytes = Path(EMBEDDER_INT8_OUT).stat().st_size + print("\nfile sizes:") + print(f" {EMBEDDER_OUT}: {fp32_bytes / 1e6:.1f} MB") + print(f" {EMBEDDER_INT8_OUT}: {int8_bytes / 1e6:.1f} MB") + + # Load holdout data and PyTorch reference predictions. + # Apply preprocess so parity is checked on the same normalization Aegis uses. + examples = load(EVAL_FILE) + texts = [preprocess(e.text) for e in examples] + torch_preds = [str(p) for p in model.predict(texts)] + + coef = head.coef_ + intercept = head.intercept_ + + tokenizer = transformers.AutoTokenizer.from_pretrained(SETFIT_DIR) + + # FP32 parity gate: ALL examples must match. + fp32_preds = _run_onnx_preds(fp32_sess, tokenizer, texts, coef, intercept, labels) + mismatches = [ + (i, texts[i], torch_preds[i], fp32_preds[i]) + for i in range(len(texts)) + if torch_preds[i] != fp32_preds[i] + ] + match_count = len(texts) - len(mismatches) + print(f"\nfp32 parity: {match_count}/{len(texts)} match") + if mismatches: + print("MISMATCH rows (idx, torch_pred, onnx_pred, text):") + for idx, text, tp, op in mismatches: + print(f" [{idx}] torch={tp!r} onnx={op!r} text={text!r}") + raise SystemExit(1) + print("fp32 parity ok") + + # INT8 check: report accuracy, do not hard-fail. + int8_sess = ort.InferenceSession( + EMBEDDER_INT8_OUT, providers=["CPUExecutionProvider"] + ) + int8_preds = _run_onnx_preds( + int8_sess, tokenizer, texts, coef, intercept, labels + ) + int8_match = sum(tp == ip for tp, ip in zip(torch_preds, int8_preds)) + print( + f"int8 accuracy: {int8_match}/{len(texts)} " + f"({100 * int8_match / len(texts):.1f}%)" + ) + + +if __name__ == "__main__": + main() diff --git a/src/routelet/preprocess.py b/src/routelet/preprocess.py new file mode 100644 index 0000000..7b5a0c2 --- /dev/null +++ b/src/routelet/preprocess.py @@ -0,0 +1,26 @@ +"""Intent-independent input normalization applied at both training and inference. + +Rules applied in order: + 1. Secret keyword tail: mask everything after a secret keyword with . + 2. Email addresses -> . + 3. Digit runs of length >= 4 -> . + +Over-redaction is acceptable; leaking is not. +""" + +import re + +_SECRET = re.compile( + r"(?i)\b(password|passcode|pin|ssn|secret|token|api\s*key|api\s*secret|credit card|card number)\b.*$", # noqa: E501 + re.MULTILINE, +) +_EMAIL = re.compile(r"(?i)[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,}") +_NUM = re.compile(r"\b\d{4,}\b") + + +def preprocess(text: str) -> str: + """Normalize text before encoding. Safe to call without knowing the intent.""" + text = _SECRET.sub(lambda m: m.group(1) + " ", text) + text = _EMAIL.sub("", text) + text = _NUM.sub("", text) + return text diff --git a/src/routelet/teacher.py b/src/routelet/teacher.py new file mode 100644 index 0000000..634af28 --- /dev/null +++ b/src/routelet/teacher.py @@ -0,0 +1,65 @@ +"""Offline Claude teacher: the strong-model labeler for distillation. + +The taxonomy prompt and label schema live here so the eval baseline +(``evaluate.py``) and the ingest labeler (``Scripts/ingest_samples.py``) +classify with identical instructions. If the two drifted, the eval would be +measuring a different teacher than the one that labeled the training data. + +Teacher, not student: this calls Claude and runs offline only, never in +production. See docs/distillation.md. +""" + +import json +import os + +import anthropic + +from routelet.data import Intent + +# The teacher's whole job is label accuracy, so this defaults to the strongest +# model. Bulk runs that care more about cost can override with the +# ROUTELET_TEACHER_MODEL env var (e.g. claude-haiku-4-5). +DEFAULT_MODEL = "claude-opus-4-8" + +# Condensed from docs/taxonomy.md. Kept stable so it caches across the +# per-command calls. +SYSTEM = """You are an intent router. Classify each user command into exactly one of five intents: + +- find_action: locate or operate a UI element on the current screen. +- integration: one discrete action against an app or service. +- chat: answer from general knowledge or conversation. No app action, no personal data. +- memory: store or recall a personal fact (storing needs an explicit remember/note/save). +- agent: a task needing two or more chained steps or a plan. + +Tie-breakers for the tricky cases: +- Two or more chained actions ("X and then Y") is agent, not integration. +- A question answered from a stored personal fact is memory; from world knowledge it is chat. +- A UI verb (click/tap/scroll) on a named element is find_action, even if an app is named. +- Playback (skip, pause, next, volume) is integration, not find_action, unless a button is named. +- "explain how to..." or "talk me through..." is chat, even when it names an app action.""" + +INTENT_SCHEMA = { + "type": "object", + "properties": {"intent": {"type": "string", "enum": [i.value for i in Intent]}}, + "required": ["intent"], + "additionalProperties": False, +} + + +def resolve_model() -> str: + """The teacher model to use: ROUTELET_TEACHER_MODEL if set, else the default.""" + return os.environ.get("ROUTELET_TEACHER_MODEL", DEFAULT_MODEL) + + +def classify(client: anthropic.Anthropic, text: str, model: str | None = None) -> Intent: + """Label one command with the teacher. Structured output pins the response to + the five-intent enum, so the parse can't return anything off-taxonomy.""" + resp = client.messages.create( + model=model or resolve_model(), + max_tokens=100, + system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}], + messages=[{"role": "user", "content": text}], + output_config={"format": {"type": "json_schema", "schema": INTENT_SCHEMA}}, + ) + out = next(b.text for b in resp.content if b.type == "text") + return Intent(json.loads(out)["intent"]) diff --git a/src/routelet/train.py b/src/routelet/train.py index b3b479c..1996feb 100644 --- a/src/routelet/train.py +++ b/src/routelet/train.py @@ -3,8 +3,8 @@ Baseline model: TF-IDF features into logistic regression. It is the floor the fine-tuned model (the ``train`` extra) has to beat, and it doubles as routelet's v1 served model since it fits in milliseconds and needs no GPU. Trains on the -synthetic pool in data/, evaluates on the hand-written seed in examples/, which -were generated by different processes so the eval isn't just memorized train. +pool in data/, evaluates on the frozen evals/holdout.jsonl, which was written +apart from the training data so the eval isn't just memorized train. """ from pathlib import Path diff --git a/src/routelet/train_setfit.py b/src/routelet/train_setfit.py new file mode 100644 index 0000000..43bc7e2 --- /dev/null +++ b/src/routelet/train_setfit.py @@ -0,0 +1,329 @@ +"""Train and persist a SetFit intent classifier. + +setfit_baseline.py is a benchmark harness: it trains, evaluates, then discards +the model. This script does the same training but saves the fitted model so +export_onnx.py can bake it into a runtime artifact. + +After saving the model a single temperature-scaling parameter T is fitted on a +stratified 20% calibration split carved from the original (non-augmented) data +files only (data/{agent,chat,find_action,integration,memory}.jsonl). T is +written to models/setfit/temperature.json so export_onnx.py can bake it into +head.json for the Rust consumer. +""" + +import json +import random +from pathlib import Path + +import numpy as np +import scipy.optimize +import torch +from datasets import Dataset +from setfit import SetFitModel, Trainer, TrainingArguments +from sklearn.metrics import classification_report +from sklearn.model_selection import StratifiedShuffleSplit + +from routelet.data import Example, Intent, load, load_dir +from routelet.preprocess import preprocess + +BASE = "BAAI/bge-small-en-v1.5" +TRAIN_DIR = "data" +ORIGINAL_FILES = [ + "data/agent.jsonl", + "data/chat.jsonl", + "data/find_action.jsonl", + "data/integration.jsonl", + "data/memory.jsonl", +] +EVAL_FILE = "evals/holdout.jsonl" +MODEL_OUT = "models/setfit" +TEMPERATURE_OUT = "models/setfit/temperature.json" +CALIB_FRAC = 0.20 +CALIB_SEED = 7 + + +def _ece(confidences: np.ndarray, correct: np.ndarray, n_bins: int = 10) -> float: + """Expected Calibration Error over equally-spaced confidence bins.""" + bin_edges = np.linspace(0.0, 1.0, n_bins + 1) + ece = 0.0 + n = len(confidences) + for lo, hi in zip(bin_edges[:-1], bin_edges[1:]): + mask = (confidences >= lo) & (confidences < hi) + if mask.sum() == 0: + continue + acc = correct[mask].mean() + conf = confidences[mask].mean() + ece += mask.sum() / n * abs(acc - conf) + return float(ece) + + +def _softmax(logits: np.ndarray) -> np.ndarray: + shifted = logits - logits.max(axis=1, keepdims=True) + exp = np.exp(shifted) + return exp / exp.sum(axis=1, keepdims=True) + + +def calibrate_temperature( + model: SetFitModel, + calib_examples: list[Example], + labels: list[str], +) -> float: + """Fit a scalar temperature T on calib_examples. + + Minimizes NLL of softmax(logits / T) against true labels using + scipy.optimize.minimize_scalar with Brent's method (bounded to [0.05, 20]). + + Returns T > 0. + + IMPORTANT: logit columns follow head.classes_ order (sklearn alphabetical), + NOT the Intent enum order. label_to_idx must be built from head.classes_. + """ + # head.classes_ gives the column ordering for coef_ / logits (sklearn alphabetical). + head_labels = model.model_head.classes_.tolist() + label_to_idx = {lbl: i for i, lbl in enumerate(head_labels)} + texts = [preprocess(e.text) for e in calib_examples] + true_idx = np.array([label_to_idx[e.intent.value] for e in calib_examples]) + + # Raw embeddings from the body (already L2-normalized by the pipeline). + embeddings = model.model_body.encode(texts, convert_to_numpy=True, show_progress_bar=False) + coef = model.model_head.coef_ # (n_classes, n_features) + intercept = model.model_head.intercept_ # (n_classes,) + logits = embeddings @ coef.T + intercept # (n_samples, n_classes) + + n = len(true_idx) + + def neg_log_likelihood(T: float) -> float: + probs = _softmax(logits / T) + # Clip to avoid log(0). + p_true = probs[np.arange(n), true_idx].clip(1e-12, 1.0) + return -np.log(p_true).mean() + + # Search T in [0.5, 20.0]. T < 1 sharpens the distribution; T > 1 softens + # it. We do not allow T below 0.5 because on a well-calibrated model with + # large logit margins the NLL optimizer would push T toward 0, collapsing + # all softmax outputs to ~1.0 and making the confidence gate degenerate. + result = scipy.optimize.minimize_scalar( + neg_log_likelihood, + bounds=(0.5, 20.0), + method="bounded", + options={"xatol": 1e-6}, + ) + T_opt = float(result.x) + + # If the optimizer hit the lower bound (T_opt ~= 0.5) and the model is + # already well-calibrated at T=1.0 (ECE < 0.05), clamp to T=1.0. + # Rationale: T < 1 only sharpens an already-confident model further. + # The Rust consumer uses max-softmax as its confidence gate; a T that + # pushes all probs to 1.0 destroys the gate's discriminating power. + probs_before = _softmax(logits) + conf_before = probs_before.max(axis=1) + correct = (probs_before.argmax(axis=1) == true_idx) + ece_t1 = _ece(conf_before, correct.astype(float)) + + if T_opt < 1.0 and ece_t1 < 0.05: + T = 1.0 + clamped_note = f" (clamped from {T_opt:.4f}; ECE@T=1 already {ece_t1:.4f} < 0.05)" + else: + T = T_opt + clamped_note = "" + + probs_after = _softmax(logits / T) + conf_after = probs_after.max(axis=1) + + ece_before = ece_t1 + ece_after = _ece(conf_after, correct.astype(float)) + + print(f"\n--- temperature calibration (calib split n={n}) ---") + print(f"T = {T:.4f}{clamped_note}") + print(f"ECE before scaling (T=1.0): {ece_before:.4f}") + print(f"ECE after scaling (T={T:.4f}): {ece_after:.4f}") + print(f"mean confidence before: {conf_before.mean():.3f}") + print(f"mean confidence after: {conf_after.mean():.3f}") + print(f"accuracy on calib split: {correct.mean():.3f}") + print("---") + + return T + + +def threshold_analysis( + model: SetFitModel, + test_examples: list[Example], + labels: list[str], + T: float, +) -> None: + """Print deferral / accuracy table for candidate confidence thresholds. + + Shows two views: + - T-scaled (T = fitted value): the confidence the Rust consumer will see. + - raw T=1.0: the uncalibrated softmax output, useful as a sanity-check + when T < 1.0 pushes all probs near 1.0 (degenerate gating). + """ + # Use head.classes_ ordering to match logit column ordering. + head_labels = model.model_head.classes_.tolist() + label_to_idx = {lbl: i for i, lbl in enumerate(head_labels)} + texts = [preprocess(e.text) for e in test_examples] + true_idx = np.array([label_to_idx[e.intent.value] for e in test_examples]) + + embeddings = model.model_body.encode(texts, convert_to_numpy=True, show_progress_bar=False) + coef = model.model_head.coef_ + intercept = model.model_head.intercept_ + logits = embeddings @ coef.T + intercept + + # Predictions are the same for any positive T (argmax is scale-invariant). + pred_idx = logits.argmax(axis=1) + correct = pred_idx == true_idx + + n = len(test_examples) + + def _print_table(conf: np.ndarray, label: str) -> None: + print(f"\n--- confidence threshold analysis (124-row holdout, {label}) ---") + print( + f"{'threshold':>10} {'deferred':>10} {'defer%':>8} " + f"{'kept_acc':>10} {'defer_acc':>10}" + ) + for tau in [0.55, 0.60, 0.65, 0.70]: + above = conf >= tau + below = ~above + n_below = int(below.sum()) + frac_below = n_below / n + kept_acc = float(correct[above].mean()) if above.sum() > 0 else float("nan") + defer_acc = float(correct[below].mean()) if below.sum() > 0 else float("nan") + print( + f"{tau:>10.2f} {n_below:>10d} {frac_below:>8.1%} " + f"{kept_acc:>10.3f} {defer_acc:>10.3f}" + ) + print("---") + + # T-scaled view (what Rust sees). + probs_scaled = _softmax(logits / T) + _print_table(probs_scaled.max(axis=1), f"T={T:.4f} scaled") + + # Raw T=1 view (always shown for comparison). + probs_raw = _softmax(logits) + _print_table(probs_raw.max(axis=1), "T=1.0 raw") + + # Flag the two low-confidence rows so we can see what's hard. + conf_raw = probs_raw.max(axis=1) + uncertain = [(i, conf_raw[i], test_examples[i].intent.value, head_labels[pred_idx[i]]) + for i in range(n) if conf_raw[i] < 0.90] + if uncertain: + print("\nlow-confidence rows on holdout (raw T=1.0 max-prob < 0.90):") + for idx, c, true_lbl, pred_lbl in uncertain: + mark = "OK" if true_lbl == pred_lbl else "WRONG" + print( + f" [{idx}] conf={c:.3f} true={true_lbl!r} pred={pred_lbl!r} " + f"[{mark}] {test_examples[idx].text!r}" + ) + + +def main() -> None: + # Seed the three RNGs that touch training. SetFit's contrastive pair + # sampling runs inside sentence-transformers and can still vary slightly + # between runs; this gets us as close to determinism as the library allows. + random.seed(0) + np.random.seed(0) + torch.manual_seed(0) + + train = load_dir(TRAIN_DIR) + test = load(EVAL_FILE) + labels = [i.value for i in Intent] + + # Load original (non-augmented) data for the calibration split. + original_examples: list[Example] = [] + for path in ORIGINAL_FILES: + original_examples.extend(load(path)) + print(f"original (non-augmented) pool: {len(original_examples)} examples") + + # Stratified 20% calibration split from original data only. + orig_texts = [e.text for e in original_examples] + orig_labels_str = [e.intent.value for e in original_examples] + sss = StratifiedShuffleSplit(n_splits=1, test_size=CALIB_FRAC, random_state=CALIB_SEED) + _, calib_idx = next(sss.split(orig_texts, orig_labels_str)) + calib_examples = [original_examples[i] for i in calib_idx] + print(f"calibration split: {len(calib_examples)} examples (stratified 20%)") + + train_ds = Dataset.from_dict( + { + "text": [preprocess(e.text) for e in train], + "label": [e.intent.value for e in train], + } + ) + + device = "cuda" if torch.cuda.is_available() else "cpu" + print( + f"training on {device}" + + (f" ({torch.cuda.get_device_name(0)})" if device == "cuda" else "") + ) + + # class_weight="balanced" counters the over-firing on agent/find_action by + # up-weighting under-represented classes in the LR head. + model = SetFitModel.from_pretrained( + BASE, + labels=labels, + device=device, + head_params={"class_weight": "balanced"}, + ) + # Verify the LR head actually received class_weight="balanced". + cw = model.model_head.get_params().get("class_weight") + assert cw == "balanced", f"class_weight not applied; got {cw!r}" + print(f"head class_weight: {cw}") + + # sampling_strategy="unique": draws every sentence-pair combination exactly + # once (no duplication). Valid in setfit 1.1.3 (confirmed in source). + # num_epochs=2 doubles the embedding training time vs. the previous 1-epoch + # run, giving the contrastive head more signal on the larger 1115-row pool. + args = TrainingArguments( + batch_size=16, + num_epochs=2, + sampling_strategy="unique", + ) + print("\nTrainingArguments:") + print(f" batch_size: {args.batch_size}") + print(f" num_epochs: {args.num_epochs}") + print(f" sampling_strategy: {args.sampling_strategy}") + assert args.sampling_strategy == "unique", ( + f"sampling_strategy not applied; got {args.sampling_strategy!r}" + ) + + trainer = Trainer( + model=model, + args=args, + train_dataset=train_ds, + ) + trainer.train() + + # Holdout evaluation. + texts = [preprocess(e.text) for e in test] + true = [e.intent.value for e in test] + preds = list(model.predict(texts)) + + print(f"\nbase {BASE} train {len(train)} eval {len(test)}\n") + print(classification_report(true, preds, labels=labels, zero_division=0)) + + # Confusion matrix. + from sklearn.metrics import confusion_matrix + cm = confusion_matrix(true, preds, labels=labels) + print("confusion matrix (rows=true, cols=pred):") + col_width = max(len(lbl) for lbl in labels) + 2 + header = " " * col_width + "".join(f"{lbl:>{col_width}}" for lbl in labels) + print(header) + for i, row_label in enumerate(labels): + row_str = f"{row_label:<{col_width}}" + "".join(f"{v:>{col_width}}" for v in cm[i]) + print(row_str) + + model.save_pretrained(MODEL_OUT) + print(f"\nsaved {MODEL_OUT}") + + # Fit temperature on calibration split. + T = calibrate_temperature(model, calib_examples, labels) + + # Persist T alongside the model for export_onnx.py to pick up. + Path(TEMPERATURE_OUT).write_text(json.dumps({"temperature": T}) + "\n") + print(f"saved {TEMPERATURE_OUT} (T={T:.4f})") + + # Threshold / deferral analysis on the frozen holdout. + threshold_analysis(model, test, labels, T) + + +if __name__ == "__main__": + main() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/preprocess_vectors.json b/tests/preprocess_vectors.json new file mode 100644 index 0000000..db7c7a8 --- /dev/null +++ b/tests/preprocess_vectors.json @@ -0,0 +1,16 @@ +{ + "_note": "Canonical preprocess() conformance vectors. A byte-identical copy lives in the other repo (routelet at tests/, aegis at aegis/tests/); keep the two in sync. preprocess is the intent-independent input normalization the model sees at BOTH training and inference: (1) after a secret keyword, mask the remainder; (2) emails -> ; (3) runs of 4+ digits -> . It must NOT depend on intent (inference does not know the intent yet). The Rust and Python implementations must reproduce every `out` exactly.", + "vectors": [ + { "in": "play despacito on spotify", "out": "play despacito on spotify" }, + { "in": "what's the capital of france", "out": "what's the capital of france" }, + { "in": "where is the search bar", "out": "where is the search bar" }, + { "in": "set volume to 50", "out": "set volume to 50" }, + { "in": "remember my name is daniel", "out": "remember my name is daniel" }, + { "in": "my wifi password is hunter2", "out": "my wifi password " }, + { "in": "email me at john@example.com", "out": "email me at " }, + { "in": "call 5551234", "out": "call " }, + { "in": "my pin is 9041", "out": "my pin " }, + { "in": "my github token is ghp_abc123", "out": "my github token " }, + { "in": "send a@b.com my pin 4321", "out": "send my pin " } + ] +} diff --git a/tests/test_augment.py b/tests/test_augment.py new file mode 100644 index 0000000..0200222 --- /dev/null +++ b/tests/test_augment.py @@ -0,0 +1,187 @@ +"""Tests for routelet.augment: determinism, label preservation, holdout non-leakage, +and conservative chaining (integration/find_action seeds never produce "and then"). + +No heavy deps (no torch/setfit/onnx). + +Run with: .venv/bin/python -m unittest tests.test_augment +""" + +import json +import random +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from routelet.augment import ( + CHAINABLE, + make_chain, + make_variant, +) +from routelet.data import Intent + +REPO = Path(__file__).parent.parent +HOLDOUT = REPO / "evals" / "holdout.jsonl" + +# Small representative sample: one row per intent. +SAMPLES = [ + ("open spotify", Intent.INTEGRATION), + ("find the settings icon", Intent.FIND_ACTION), + ("what time is it", Intent.CHAT), + ("remember that my meeting is at 3pm", Intent.MEMORY), + ("search amazon and buy the cheapest hdmi cable", Intent.AGENT), +] + +# A partner text for chaining tests. +_PARTNER = "turn on do not disturb" + + +def _norm(text: str) -> str: + return text.strip().lower() + + +class TestDeterminism(unittest.TestCase): + """Same seed produces identical variants on repeated calls.""" + + def _run_once(self) -> list[str]: + results = [] + for text, intent in SAMPLES: + rng = random.Random(42) + results.append(make_variant(text, intent, rng)) + return results + + def test_make_variant_deterministic(self) -> None: + first = self._run_once() + second = self._run_once() + self.assertEqual(first, second) + + def test_make_chain_deterministic(self) -> None: + results_a = [] + results_b = [] + for text, intent in SAMPLES: + rng_a = random.Random(7) + rng_b = random.Random(7) + results_a.append(make_chain(text, _PARTNER, intent, rng_a)) + results_b.append(make_chain(text, _PARTNER, intent, rng_b)) + self.assertEqual(results_a, results_b) + + +class TestLabelPreservation(unittest.TestCase): + """make_variant does not change the intent of the source row. + + The function returns a string, not an Example, so label preservation is + guaranteed structurally: the caller supplies both the text and the intent + and the variant inherits the intent at write-time. We verify that make_variant + returns a non-empty string (a valid variant exists to attach the label to) + and that the function never raises for any intent value. + """ + + def test_variant_is_non_empty_string(self) -> None: + for text, intent in SAMPLES: + with self.subTest(intent=intent): + rng = random.Random(0) + variant = make_variant(text, intent, rng) + self.assertIsInstance(variant, str) + self.assertTrue(variant.strip(), f"empty variant for {intent!r}") + + def test_augment_file_preserves_labels(self) -> None: + """Rows written by augment() carry the same intent as their source file.""" + augmented = REPO / "data" / "augmented.jsonl" + if not augmented.exists(): + self.skipTest("augmented.jsonl not present; run python -m routelet.augment first") + valid_intents = {i.value for i in Intent} + with augmented.open() as f: + for n, line in enumerate(f, start=1): + if not line.strip(): + continue + with self.subTest(line=n): + row = json.loads(line) + self.assertIn( + row["intent"], + valid_intents, + msg=f"augmented.jsonl:{n} bad intent {row['intent']!r}", + ) + + +class TestHoldoutNonLeakage(unittest.TestCase): + """Variants generated for the sample rows do not appear in the holdout set.""" + + @classmethod + def setUpClass(cls) -> None: + cls.holdout_norms: set[str] = set() + with HOLDOUT.open() as f: + for line in f: + if line.strip(): + row = json.loads(line) + cls.holdout_norms.add(_norm(row["text"])) + + def test_variants_not_in_holdout(self) -> None: + for text, intent in SAMPLES: + for seed in range(10): + rng = random.Random(seed) + variant = make_variant(text, intent, rng) + with self.subTest(intent=intent, seed=seed): + self.assertNotIn( + _norm(variant), + self.holdout_norms, + msg=f"variant {variant!r} matched holdout entry", + ) + + +class TestConservativeChaining(unittest.TestCase): + """integration and find_action seeds must never produce "and then" variants. + + make_chain returns None when the intent is not in CHAINABLE, which is the + mechanism that prevents " and then " from appearing in integration/find_action + outputs. We test that: + 1. make_chain returns None for integration and find_action (unit-level). + 2. For the complementary check, make_chain returns a non-None string for + all intents in CHAINABLE (chat, memory, agent). + 3. When we call make_variant 50 times on integration/find_action rows the + results never contain " and then " (make_variant never splices chains in). + """ + + def test_make_chain_returns_none_for_non_chainable(self) -> None: + non_chainable = [Intent.INTEGRATION, Intent.FIND_ACTION] + for intent in non_chainable: + with self.subTest(intent=intent): + rng = random.Random(0) + result = make_chain("open spotify", _PARTNER, intent, rng) + self.assertIsNone( + result, + msg=f"make_chain should return None for {intent!r} but got {result!r}", + ) + + def test_chainable_intents_not_none(self) -> None: + for intent in CHAINABLE: + with self.subTest(intent=intent): + rng = random.Random(0) + result = make_chain("do something", _PARTNER, intent, rng) + self.assertIsNotNone( + result, + msg=f"make_chain returned None for chainable intent {intent!r}", + ) + + def test_make_variant_no_and_then_for_non_chainable(self) -> None: + """make_variant never embeds 'and then' in integration or find_action output.""" + non_chainable_samples = [ + ("open spotify", Intent.INTEGRATION), + ("send a slack message to alice", Intent.INTEGRATION), + ("find the settings icon", Intent.FIND_ACTION), + ("where is the volume button", Intent.FIND_ACTION), + ] + for text, intent in non_chainable_samples: + for seed in range(50): + rng = random.Random(seed) + variant = make_variant(text, intent, rng) + with self.subTest(intent=intent, seed=seed): + self.assertNotIn( + " and then ", + variant, + msg=f"variant for {intent!r} contains 'and then': {variant!r}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_data_validity.py b/tests/test_data_validity.py new file mode 100644 index 0000000..0fb0f9a --- /dev/null +++ b/tests/test_data_validity.py @@ -0,0 +1,97 @@ +"""Data integrity tests: JSONL schema, intent values, dedup, and holdout disjointness. + +No heavy deps: scikit-learn is pulled in by data.py but torch/setfit are not. + +Run with: .venv/bin/python -m unittest tests.test_data_validity +""" + +import json +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from routelet.data import Intent, load, load_dir + +REPO = Path(__file__).parent.parent +DATA_DIR = REPO / "data" +HOLDOUT = REPO / "evals" / "holdout.jsonl" + +VALID_INTENTS = {i.value for i in Intent} +DATA_FILES = sorted(DATA_DIR.glob("*.jsonl")) + + +class TestJsonlSchema(unittest.TestCase): + """Every line in every file has exactly {"text","intent"} with a valid intent.""" + + def _check_file(self, path: Path) -> None: + with path.open() as f: + for n, raw in enumerate(f, start=1): + if not raw.strip(): + continue + with self.subTest(file=path.name, line=n): + row = json.loads(raw) + self.assertEqual( + set(row.keys()), + {"text", "intent"}, + msg=f"{path.name}:{n} unexpected keys: {set(row.keys())}", + ) + self.assertIn( + row["intent"], + VALID_INTENTS, + msg=f"{path.name}:{n} unknown intent {row['intent']!r}", + ) + + @unittest.skipUnless(DATA_FILES, "data/*.jsonl is gitignored; present only locally") + def test_data_files(self) -> None: + for path in DATA_FILES: + self._check_file(path) + + def test_holdout_file(self) -> None: + self._check_file(HOLDOUT) + + +class TestTrainingPoolDedup(unittest.TestCase): + """No two rows in the combined training pool share the same text.""" + + @unittest.skipUnless(DATA_FILES, "data/*.jsonl is gitignored; present only locally") + def test_no_duplicate_texts(self) -> None: + examples = load_dir(DATA_DIR) + seen: dict[str, str] = {} # text -> first filename + duplicates: list[str] = [] + for ex in examples: + if ex.text in seen: + duplicates.append(ex.text) + else: + seen[ex.text] = ex.intent.value + self.assertEqual( + duplicates, + [], + msg=f"Duplicate texts in training pool ({len(duplicates)}):\n" + + "\n".join(f" {t!r}" for t in duplicates[:20]), + ) + + +class TestHoldoutDisjointness(unittest.TestCase): + """Zero overlap between holdout texts and training texts (normalized).""" + + @staticmethod + def _norm(text: str) -> str: + return text.strip().lower() + + @unittest.skipUnless(DATA_FILES, "data/*.jsonl is gitignored; present only locally") + def test_no_train_eval_overlap(self) -> None: + train_norms = {self._norm(e.text) for e in load_dir(DATA_DIR)} + holdout_norms = {self._norm(e.text) for e in load(HOLDOUT)} + overlap = train_norms & holdout_norms + self.assertEqual( + overlap, + set(), + msg=f"Holdout texts found in training pool ({len(overlap)} texts):\n" + + "\n".join(f" {t!r}" for t in sorted(overlap)), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_preprocess.py b/tests/test_preprocess.py new file mode 100644 index 0000000..cf7021a --- /dev/null +++ b/tests/test_preprocess.py @@ -0,0 +1,33 @@ +"""Conformance tests for preprocess() against the shared fixture. + +The same vectors live in the aegis repo; the Rust implementation must also +satisfy them. Do not edit the fixture; fix preprocess() if a vector fails. + +Run with: .venv/bin/python -m unittest tests.test_preprocess +""" + +import json +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from routelet.preprocess import preprocess + +FIXTURE = Path(__file__).parent / "preprocess_vectors.json" + + +class TestPreprocess(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.vectors = json.loads(FIXTURE.read_text())["vectors"] + + def test_all_vectors(self) -> None: + for v in self.vectors: + with self.subTest(input=v["in"]): + self.assertEqual(preprocess(v["in"]), v["out"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/uv.lock b/uv.lock index 7f7eb92..57e1444 100644 --- a/uv.lock +++ b/uv.lock @@ -265,7 +265,7 @@ name = "cuda-bindings" version = "13.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder" }, + { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/52/c8/b2589d68acf7e3d63e2be330b84bc25712e97ed799affbca7edd7eae25d6/cuda_bindings-13.2.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e865447abfb83d6a98ad5130ed3c70b1fc295ae3eeee39fd07b4ddb0671b6788", size = 5722404, upload-time = "2026-03-11T00:12:44.041Z" }, @@ -296,34 +296,34 @@ wheels = [ [package.optional-dependencies] cudart = [ - { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" }, ] cufft = [ - { name = "nvidia-cufft", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cufft", marker = "sys_platform == 'linux'" }, ] cufile = [ { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" }, ] curand = [ - { name = "nvidia-curand", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-curand", marker = "sys_platform == 'linux'" }, ] cusolver = [ - { name = "nvidia-cusolver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, ] [[package]] @@ -891,7 +891,7 @@ name = "nvidia-cublas" version = "13.1.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc" }, + { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, @@ -930,7 +930,7 @@ name = "nvidia-cudnn-cu13" version = "9.20.0.48" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, + { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, @@ -942,7 +942,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -972,9 +972,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, - { name = "nvidia-cusparse" }, - { name = "nvidia-nvjitlink" }, + { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cusparse", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -986,7 +986,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -1514,13 +1514,15 @@ name = "routelet" version = "0.0.1" source = { editable = "." } dependencies = [ - { name = "fastapi" }, - { name = "pydantic" }, { name = "scikit-learn" }, - { name = "uvicorn" }, ] [package.optional-dependencies] +serve = [ + { name = "fastapi" }, + { name = "pydantic" }, + { name = "uvicorn" }, +] train = [ { name = "datasets" }, { name = "torch" }, @@ -1530,14 +1532,14 @@ train = [ [package.metadata] requires-dist = [ { name = "datasets", marker = "extra == 'train'" }, - { name = "fastapi" }, - { name = "pydantic" }, + { name = "fastapi", marker = "extra == 'serve'" }, + { name = "pydantic", marker = "extra == 'serve'" }, { name = "scikit-learn" }, { name = "torch", marker = "extra == 'train'" }, { name = "transformers", marker = "extra == 'train'" }, - { name = "uvicorn" }, + { name = "uvicorn", marker = "extra == 'serve'" }, ] -provides-extras = ["train"] +provides-extras = ["train", "serve"] [[package]] name = "safetensors"