LTR training dataset pipeline

Builds a sharded, resumable LightGBM LambdaMART dataset for approximately 50k queries. The pipeline:

1. Retrieves the top 20k BM25 and ANN rankings.
2. Deterministically samples 65 documents from fixed rank buckets in each list.
3. Deduplicates the union while retaining both ranks, scores, and full-list percentiles.
4. Hydrates and grades documents in batches of 10 with GPT-5.4-nano.
5. Fetches pointwise Vespa features and writes grouped train/validation datasets.

Documents that disappear from Vespa between collection and hydration are counted in the annotation summary and omitted rather than failing the shard. If the model returns the wrong number of judgments, that batch is recursively split and retried.

Each completed Parquet shard is an atomic checkpoint. Re-running a stage skips completed shards and retries shards with a failure report.

Input

input/queries.csv must contain exactly one non-empty query column value per row. Repeated whitespace is normalized. The one-based CSV row number becomes query_id; blank and duplicate normalized queries cause validation to fail. A stable text hash assigns queries 90% to train and 10% to validation.

Configuration

The .env file or environment must provide:

VESPA_HOST="https://..."
VESPA_TOKEN="..."
CEREBRIUM_EMBEDDER_BASE_URL="https://..."
CEREBRIUM_API_KEY="..."
OPENAI_API_KEY="..."

# Optional
OPENAI_BASE_URL="https://eu.api.openai.com/v1"
OPENAI_INPUT_COST_PER_MILLION="0.05"
OPENAI_OUTPUT_COST_PER_MILLION="0.40"

The cost rates only affect shard summaries; billing is measured from API token usage.

Run

Run one production-code-path canary query:

uv run python src/ltr_training/pipeline.py all --dataset-version canary \
  --shard-size 1 --max-shards 1

Run the 50k pipeline stage by stage:

uv run python src/ltr_training/pipeline.py collect --dataset-version v1
uv run python src/ltr_training/pipeline.py annotate --dataset-version v1
uv run python src/ltr_training/pipeline.py features --dataset-version v1

Use --shard-id 12 to run or retry one shard. Important controls include --shard-size, --query-concurrency, --vespa-concurrency, --openai-concurrency, and --openai-requests-per-minute.

Use features --overwrite-features to rebuild feature shards from existing annotations. Legacy annotation rows without domain_count are backfilled by rerunning only the Vespa retrievals; no annotation requests are made.

The one-time candidate-feature migration is resumable and performs no network requests:

uv run python scripts/backfill_candidate_features.py --dry-run
uv run python scripts/backfill_candidate_features.py

Changing the sampling policy or prompt requires a new --dataset-version.

Training

Each training creates an isolated run directory. The run ID defaults to a UTC timestamp and can be set explicitly:

uv run python src/ltr_training/train.py --dataset-version v1
uv run python src/ltr_training/train.py --dataset-version v1 --run-id experiment-1

The run directory contains the model, metadata, complete validation metric history, and native TreeSHAP feature importance across the full validation set. SHAP values are aggregated per query before queries are averaged equally. Existing run IDs are never overwritten. A SHAP failure is recorded separately and does not invalidate a successfully trained model.

After SHAP, the model is evaluated on the first 100 unique queries from the Brave evaluation CSV, preserving file order and using every reference doc_id for those queries. The report records mean document-ID recall at 100, 400, 1,000, 2,000, 5,000, 10,000, and 20,000, plus per-query results. Evaluation failures are recorded separately and do not invalidate the trained model.

Model-driven sampling

After autoresearch has produced a champion, run the outer data-improvement loop:

uv run model-driven-sampling \
  --source-dataset v1 \
  --max-iterations 10 \
  --queries-per-iteration 100

Each iteration deterministically chooses unused training queries, scores the full union of BM25 top 20k and ANN top 20k, and selects three rank-disagreement documents plus three random controls from each evaluation rank band. Selected features are saved before annotation, so they are not fetched from Vespa twice.

The loop adds labels cumulatively, removes duplicate query/document pairs and query groups having only one distinct grade, trains with a frozen snapshot of the current champion recipe, and applies the same 100-query then 500-query recall funnel as autoresearch. Promotion also requires recall at every cutoff to remain at least as high as the current champion. Use --max-cutoff-regression to permit a small regression explicitly.

Progress is resumable under:

output/<source-dataset>/model_sampling/
  state.json
  best.json
  frozen_autoresearch_train.py
  iterations/001/
    manifest.json
    queries.parquet
    mined/
    annotations/
    dataset/
    training/
    evaluation/

Useful safety limits are --max-labels, --max-cost-usd, and --patience. Full autoresearch should be run once after selecting the best data iteration, not inside every model-driven sampling iteration.

Evaluation API

The API retrieves the full BM25 and ANN rankings, takes the configured top LTR_CANDIDATE_DEPTH candidates from each, computes all model features, and returns the top LTR_RESULT_LIMIT documents ordered by the LambdaMART score. Feature fetching is split into VESPA_FEATURE_BATCH_SIZE-sized requests to stay below Vespa's query-tree limit. Full-ranking sizes are retained when computing rank-percentile features.

Start the API:

uv run --env-file .env python src/ltr_training/api.py \
  --model-path output/v1/trainings/<run-id>/lambdamart.txt

Or use the Makefile:

make api API_ARGS="--model-path output/v1/trainings/<run-id>/lambdamart.txt"

Rank one query:

curl http://127.0.0.1:8000/rank \
  --header 'content-type: application/json' \
  --data '{"query":"example search"}'

Recall evaluators that only need the final ordering can use the compact endpoint to avoid serializing candidate scores and retrieval metadata:

curl http://127.0.0.1:8000/rank/doc-ids \
  --header 'content-type: application/json' \
  --data '{"query":"example search"}'

The model is loaded once when the application starts. Its embedded feature names must exactly match the canonical schema in features.py; startup fails if training and evaluation features have drifted.

Outputs

output/<dataset-version>/
  candidates/part-XXXXX.parquet
  annotations/part-XXXXX.parquet
  features/part-XXXXX.parquet
  train.parquet
  validation.parquet
  train_groups.json
  validation_groups.json
  trainings/<run-id>/
    lambdamart.txt
    metadata.json
    metrics.json
    shap_importance.json
    shap_feature_importance.png
    recall_evaluation.json

Each stage also writes compact per-shard summaries. Failed shards write part-XXXXX.failures.json and are not marked complete.