Skip to content

Add Phase 2 learned router training and evaluation#1

Merged
ujjwalkumar26 merged 2 commits into
mainfrom
phase-2-learned-router
Apr 11, 2026
Merged

Add Phase 2 learned router training and evaluation#1
ujjwalkumar26 merged 2 commits into
mainfrom
phase-2-learned-router

Conversation

@ujjwalkumar26

Copy link
Copy Markdown
Owner

Summary

This PR delivers roadmap Part 2 by turning LearnedRouter into a real trainable router and adding the end-to-end pieces needed to use it in practice.

It adds:

  • relevance-driven router training
  • learned-router evaluation APIs
  • persistence via save/load
  • public exports for the learned router types
  • runnable examples for training and evaluation
  • unit coverage for the new behavior
  • README updates for training, inference, and evaluation workflows

What Changed

  • Reworked LearnedRouter to learn metric weights from labeled relevance data instead of training toward uniform weights.
  • Added fit_from_relevance(...) to derive routing targets from per-metric retrieval quality.
  • Added evaluate_from_relevance(...) and LearnedRouterEvaluation for measuring learned routing quality.
  • Kept the runtime API simple: train a router once, then pass it into refract.search(...).
  • Added save() / load() support for trained routers.
  • Exported the learned-router classes from the package root and routing package.
  • Added examples:
    • examples/train_learned_router.py
    • examples/evaluate_learned_router.py
  • Added tests in tests/unit/test_learned_router.py.
  • Updated the README to document Phase 2 usage and examples.
  • Removed the stale PyTorch requirement from the learned extra, since the new implementation uses existing sklearn-based dependencies.

Notes

  • The new evaluation/example flow is designed to help users validate learning on their own data instead of claiming a universal quality lift.
  • The sample evaluation on the tiny bundled dataset shows the learned router working end to end; on held-out sample queries it matches heuristic routing, which is expected on such a small corpus.

Validation

Ran locally:

  • ruff check src tests
  • ruff format --check src tests
  • pytest tests/ -x -q --tb=short
  • pytest tests/
  • python examples/train_learned_router.py
  • python examples/evaluate_learned_router.py
  • ruff check src tests examples/train_learned_router.py examples/evaluate_learned_router.py
  • ruff format --check src tests examples/train_learned_router.py examples/evaluate_learned_router.py

Reviewer Focus

  • API shape and ergonomics of LearnedRouter
  • training target derivation from per-metric relevance quality
  • docs/examples clarity for the new Phase 2 workflow
  • package metadata change removing the PyTorch learned extra requirement

@ujjwalkumar26 ujjwalkumar26 merged commit 45ad47c into main Apr 11, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant