How owwo99/ragrouter works
1. What this repo is
- RAGRouter is a training/evaluation codebase for learning a router that assigns each query to the best model in a fixed pool of candidate LMs.
- The router is trained on precomputed supervision for whether each model answer is correct, with separate labels for no RAG, local retrieval, and online retrieval settings.
2. Top-level architecture
- The repository is split into two main areas:
- data/: all preprocessed dataset assets and supervision files.
- src/: the actual training pipeline, dataset loading, model definition, and utilities.
- There is no separate server, API layer, or interactive application; the main runtime is a batch training script.
3. Data subsystem
- Each dataset lives under its own folder: medmcqa, nq, popqa, triviaqa, and webq.
- Every dataset directory contains:
- query_map.json: query IDs to raw query text.
- doc_map.json: document IDs to document text.
- train.csv / test.csv: binary correctness labels for model outputs without RAG.
- train_local.csv / test_local.csv: labels for outputs using local retrieval.
- train_online.csv / test_online.csv: labels for outputs using online retrieval where provided.
- data/model_map.json defines the mapping from model IDs to the candidate model names in the pool.
4. Data boundary and preprocessing
- The data folder already contains preprocessed artifacts, so the repo does not include raw corpus ingestion or retrieval indexing logic.
- This means architecture-wise the retrieval and labeling work happens before training; the trainer consumes static CSV/JSON assets rather than calling live retrieval services.
5. Entry point and orchestration
- src/main.py is the primary entry point.
- It accepts --data_type to choose the dataset and --task_type to choose the supervision mode (local or online).
- It also exposes the main training hyperparameters: embedding size, number of models, batch size, epochs, learning rate, classification weight, and contrastive temperature.
6. Dataset layer
- src/dataset.py is the adapter between the on-disk assets and the model.
- Its role is to join together query text, document text, and model correctness labels into training examples.
- It likely handles split selection, sample construction, and batching-compatible tensors or records for the router model.
7. Model layer
- src/model.py contains the routing model itself.
- Based on the README, the model learns two key representations:
- a knowledge representation for the query/document context,
- a RAG capability vector for modeling which candidate model is best suited.
- The training objective combines:
- a classification loss for routing correctness,
- a contrastive loss controlled by tau,
- weighted by lambda1.
8. Utility layer
- src/utils.py holds shared helpers used across the pipeline.
- In a repo like this, that usually includes things such as metric computation, seeding, logging, file loading, and experiment setup.
- It is not a separate subsystem from the router; it supports the training/evaluation flow.
9. Core data flow
- Query ID and document ID come from the dataset files.
- dataset.py resolves those IDs into text using query_map.json and doc_map.json.
- The router model consumes the resulting sample and produces scores over the model pool.
- Those scores are compared against the correctness labels from the selected CSV split to train and evaluate routing quality.
10. Retrieval-mode branching
- The key architectural switch is task_type:
- local uses the local-retrieval labels/files.
- online uses the online-retrieval labels/files where available.
- This is not a runtime retrieval pipeline; it is a supervision-mode selection that changes which labels the router learns from.
11. Model pool boundary
- The candidate model pool is fixed and parameterized by --num_models and model_map.json.
- The router does not generate answers itself; it decides which downstream model should be trusted for a query under the selected retrieval regime.
12. What an architecture diagram should show
- A clean diagram would have:
- Dataset assets on the left (query_map, doc_map, CSV labels, model map),
- Data loader / dataset adapter in the middle,
- Router model with classification and contrastive objectives,
- Training/evaluation loop in main.py,
- and an output of model-selection decisions / evaluation metrics on the right.
- The main boundary to emphasize is that all retrieval and labeling inputs are offline precomputed artifacts, not live external systems.
Main entry points: dataset.py Dataset · main.py Main · model.py Router model · model_map.json Model map
Components 13
Data
Offline precomputed query, document, and label assets consumed by training.
- Model map · Maps model IDs to the fixed set of candidate LMs.
data/model_map.json - MedMCQA · Preprocessed query/doc maps and split labels for one task.
data/medmcqa - NQ · Preprocessed query/doc maps and split labels for one task.
data/nq - PopQA · Preprocessed query/doc maps and split labels for one task.
data/popqa - TriviaQA · Preprocessed query/doc maps and split labels for one task.
data/triviaqa - WebQ · Preprocessed query/doc maps and split labels for one task.
data/webq
Source
Batch training pipeline, dataset adapter, model, and shared utilities.
- Main · Primary batch training script and CLI for dataset and supervision selection.
src/main.py - Dataset · Builds training examples by resolving IDs to text and labels.
src/dataset.py - Router model · Learns query/document representations and scores candidate models.
src/model.py - Routing loss · Supervises correct model selection over the candidate pool.
src/model.py - Contrastive loss · Shapes the knowledge and capability representations with temperature tau.
src/model.py - Utils · Common support for metrics, seeding, logging, and experiment setup.
src/utils.py - Evaluation · Reports routing quality and selected model decisions.
Connections 15
- Model map → Dataset: candidate mapping · Provides the fixed model pool used by the router.
- MedMCQA → Dataset: offline data · Task-specific precomputed query, document, and label assets.
- NQ → Dataset: offline data · Task-specific precomputed query, document, and label assets.
- PopQA → Dataset: offline data · Task-specific precomputed query, document, and label assets.
- TriviaQA → Dataset: offline data · Task-specific precomputed query, document, and label assets.
- WebQ → Dataset: offline data · Task-specific precomputed query, document, and label assets.
- Main → Dataset: load split · Selects dataset and supervision mode from the CLI.
- Main → Router model: train · Runs the batch optimization loop for the routing model.
- Main → Utils: support · Uses shared helpers for reproducibility and experiment plumbing.
- Dataset → Router model: examples · Feeds resolved training records into the model.
- Router model → Routing loss: optimizes · Updates routing decisions against correctness labels.
- Router model → Contrastive loss: optimizes · Updates representation learning with temperature-controlled contrast.
- Routing loss → Evaluation: scores · Produces routing accuracy and related metrics.
- Contrastive loss → Evaluation: shapes · Improves the learned embedding space used by the router.
- Main → Evaluation: report · Aggregates and prints final training and evaluation results.