# owwo99/ragrouter architecture

> Architecture diagram of the GitHub repository owwo99/ragrouter, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/owwo99/ragrouter
- Repository: https://github.com/owwo99/ragrouter

## Overview

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: [`src/dataset.py`](https://github.com/owwo99/ragrouter/blob/main/src/dataset.py) (Dataset), [`src/main.py`](https://github.com/owwo99/ragrouter/blob/main/src/main.py) (Main), [`src/model.py`](https://github.com/owwo99/ragrouter/blob/main/src/model.py) (Router model), [`data/model_map.json`](https://github.com/owwo99/ragrouter/blob/main/data/model_map.json) (Model map)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data"]
  node_model_map["Model map<br/>candidate pool map<br/>[model_map.json]"]
  node_medmcqa[("MedMCQA<br/>dataset bundle")]
  node_nq[("NQ<br/>dataset bundle")]
  node_popqa[("PopQA<br/>dataset bundle")]
  node_triviaqa[("TriviaQA<br/>dataset bundle")]
  node_webq[("WebQ<br/>dataset bundle")]
end

subgraph group_src["Source"]
  node_main["Main<br/>trainer entrypoint<br/>[main.py]"]
  node_dataset["Dataset<br/>data adapter<br/>[dataset.py]"]
  node_router_model["Router model<br/>ranking model<br/>[model.py]"]
  node_cls_loss(("Routing loss<br/>classification objective<br/>[model.py]"))
  node_contrastive(("Contrastive loss<br/>representation objective<br/>[model.py]"))
  node_utils["Utils<br/>shared helpers<br/>[utils.py]"]
  node_eval["Evaluation<br/>metrics output"]
end

node_model_map -->|"candidate mapping"| node_dataset
node_medmcqa -->|"offline data"| node_dataset
node_nq -->|"offline data"| node_dataset
node_popqa -->|"offline data"| node_dataset
node_triviaqa -->|"offline data"| node_dataset
node_webq -->|"offline data"| node_dataset
node_main -->|"load split"| node_dataset
node_main -->|"train"| node_router_model
node_main -->|"support"| node_utils
node_dataset -->|"examples"| node_router_model
node_router_model -->|"optimizes"| node_cls_loss
node_router_model -->|"optimizes"| node_contrastive
node_cls_loss -->|"scores"| node_eval
node_contrastive -->|"shapes"| node_eval
node_main -->|"report"| node_eval

click node_model_map "https://github.com/owwo99/ragrouter/blob/main/data/model_map.json"
click node_medmcqa "https://github.com/owwo99/ragrouter/tree/main/data/medmcqa"
click node_nq "https://github.com/owwo99/ragrouter/tree/main/data/nq"
click node_popqa "https://github.com/owwo99/ragrouter/tree/main/data/popqa"
click node_triviaqa "https://github.com/owwo99/ragrouter/tree/main/data/triviaqa"
click node_webq "https://github.com/owwo99/ragrouter/tree/main/data/webq"
click node_main "https://github.com/owwo99/ragrouter/blob/main/src/main.py"
click node_dataset "https://github.com/owwo99/ragrouter/blob/main/src/dataset.py"
click node_router_model "https://github.com/owwo99/ragrouter/blob/main/src/model.py"
click node_cls_loss "https://github.com/owwo99/ragrouter/blob/main/src/model.py"
click node_contrastive "https://github.com/owwo99/ragrouter/blob/main/src/model.py"
click node_utils "https://github.com/owwo99/ragrouter/blob/main/src/utils.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_model_map,node_medmcqa,node_nq,node_popqa,node_triviaqa,node_webq toneBlue
class node_main,node_dataset,node_router_model,node_cls_loss,node_contrastive,node_utils,node_eval toneAmber
```

## Components

### Data

Offline precomputed query, document, and label assets consumed by training.

- [Model map](https://github.com/owwo99/ragrouter/blob/main/data/model_map.json): candidate pool map · Maps model IDs to the fixed set of candidate LMs. · `data/model_map.json`
- [MedMCQA](https://github.com/owwo99/ragrouter/tree/main/data/medmcqa): dataset bundle · Preprocessed query/doc maps and split labels for one task. · `data/medmcqa`
- [NQ](https://github.com/owwo99/ragrouter/tree/main/data/nq): dataset bundle · Preprocessed query/doc maps and split labels for one task. · `data/nq`
- [PopQA](https://github.com/owwo99/ragrouter/tree/main/data/popqa): dataset bundle · Preprocessed query/doc maps and split labels for one task. · `data/popqa`
- [TriviaQA](https://github.com/owwo99/ragrouter/tree/main/data/triviaqa): dataset bundle · Preprocessed query/doc maps and split labels for one task. · `data/triviaqa`
- [WebQ](https://github.com/owwo99/ragrouter/tree/main/data/webq): dataset bundle · Preprocessed query/doc maps and split labels for one task. · `data/webq`

### Source

Batch training pipeline, dataset adapter, model, and shared utilities.

- [Main](https://github.com/owwo99/ragrouter/blob/main/src/main.py): trainer entrypoint · Primary batch training script and CLI for dataset and supervision selection. · `src/main.py`
- [Dataset](https://github.com/owwo99/ragrouter/blob/main/src/dataset.py): data adapter · Builds training examples by resolving IDs to text and labels. · `src/dataset.py`
- [Router model](https://github.com/owwo99/ragrouter/blob/main/src/model.py): ranking model · Learns query/document representations and scores candidate models. · `src/model.py`
- [Routing loss](https://github.com/owwo99/ragrouter/blob/main/src/model.py): classification objective · Supervises correct model selection over the candidate pool. · `src/model.py`
- [Contrastive loss](https://github.com/owwo99/ragrouter/blob/main/src/model.py): representation objective · Shapes the knowledge and capability representations with temperature tau. · `src/model.py`
- [Utils](https://github.com/owwo99/ragrouter/blob/main/src/utils.py): shared helpers · Common support for metrics, seeding, logging, and experiment setup. · `src/utils.py`
- **Evaluation**: metrics output · Reports routing quality and selected model decisions.

## Connections

- 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.