SproutRAG

SproutRAG is a retrieval-augmented generation stack built for structured, multi-granularity evidence. It combines hierarchical attention-based indexing with flexible retrieval, optional reranking, and generation pipelines. The result is a practical system for building and evaluating RAG workflows end-to-end.

Quick Start Overview

You will typically follow this flow:

1. Install the package and optional dependencies.
2. Create or modify a config file (YAML or JSON).
3. Index documents with sproutrag index.
4. Run retrieval with sproutrag retrieve.
5. Generate answers with sproutrag answer.
6. Train the model with sproutrag train (optional).
7. Evaluate retrieval and generation with sproutrag evaluate.

The CLI uses typed configs so you can keep everything reproducible and easy to modify.

Installation

Install in editable mode:

pip install -e .

Optional extras:

pip install -e .[yaml]
pip install -e .[eval]
pip install -e .[spacy]

Notes:

• YAML configs require PyYAML (.[yaml]).
• Evaluation extras enable ROUGE-L, METEOR, and BERTScore (.[eval]).
• spaCy improves sentence splitting (.[spacy]).

Config Files

SproutRAG runs from YAML or JSON config files. The CLI expects a config per command. You can start from the examples below and adjust as needed.

Index Config

runtime:
  device: null
  seed: 42
  num_workers: 0
  use_cuda_if_available: true

encoder:
  model_name_or_path: sentence-transformers/all-MiniLM-L6-v2
  max_length: 4096
  normalize_embeddings: true
  trust_remote_code: true

aggregator:
  init_strategy: uniform
  mask_diagonal: true

indexing:
  input_path: data/documents.jsonl
  output_dir: indexes
  index_name: demo
  max_sentences_per_chunk: 2
  batch_size: 1
  overwrite: true

The index input file is JSONL with one document per line:

{"doc_id": "doc-1", "text": "Your document text", "title": "Optional title", "metadata": {"source": "example"}}

Retrieve Config

runtime:
  device: null
  seed: 42
  num_workers: 0
  use_cuda_if_available: true

encoder:
  model_name_or_path: sentence-transformers/all-MiniLM-L6-v2
  max_length: 4096
  normalize_embeddings: true
  trust_remote_code: true

retrieval:
  index_dir: indexes
  index_name: demo
  query: "How does SproutRAG build the retrieval hierarchy?"
  query_path: null
  output_path: outputs/retrieval.json
  top_k: 10
  per_document_top_k: 5
  beam_width: 5
  threshold: 0.0
  collect_strategy: threshold
  include_root: false
  multi_document: true

reranker:
  enabled: false
  type: none
  model_name_or_path: null
  max_length: 512
  batch_size: 8
  metadata_score_key: null
  trust_remote_code: true

Answer Config

runtime:
  device: null
  seed: 42
  num_workers: 0
  use_cuda_if_available: true

encoder:
  model_name_or_path: sentence-transformers/all-MiniLM-L6-v2
  max_length: 4096
  normalize_embeddings: true
  trust_remote_code: true

retrieval:
  index_dir: indexes
  index_name: demo
  query: "How does SproutRAG retrieve multi-granularity evidence?"
  query_path: null
  output_path: outputs/answer.json
  top_k: 10
  per_document_top_k: 5
  beam_width: 5
  threshold: 0.0
  collect_strategy: threshold
  include_root: false
  multi_document: true

reranker:
  enabled: false
  type: none
  model_name_or_path: null
  max_length: 512
  batch_size: 8
  metadata_score_key: null
  trust_remote_code: true

context:
  include_scores: true
  include_metadata: false
  include_node_ids: true
  context_separator: "\n\n"
  max_context_chars: 12000

generator:
  type: echo
  model_name_or_path: null
  max_input_length: 4096
  max_new_tokens: 256
  temperature: 0.0
  do_sample: false
  include_system_prompt: true
  system_prompt: null
  trust_remote_code: true

Train Config

runtime:
  device: null
  seed: 42
  num_workers: 0
  use_cuda_if_available: true

encoder:
  model_name_or_path: sentence-transformers/all-MiniLM-L6-v2
  max_length: 4096
  normalize_embeddings: true
  trust_remote_code: true

aggregator:
  init_strategy: uniform
  mask_diagonal: true

training:
  train_path: data/train.jsonl
  output_dir: runs
  run_name: sproutrag_train
  num_epochs: 3
  batch_size: 32
  max_negatives: null
  require_equal_negatives: true
  learning_rate: 0.00002
  aggregator_learning_rate: 0.001
  weight_decay: 0.0
  warmup_ratio: 0.05
  gradient_clip_norm: 1.0
  checkpoint_every_steps: 100
  log_every_steps: 1
  max_sentences_per_chunk: 2
  normalize_passage_embeddings: true

loss:
  retrieval_temperature: 0.05
  attention_lambda: 0.1
  use_in_batch_negatives: false
  retrieval_reduction: mean
  attention_reduction: mean
  attention_example_reduction: mean
  allow_self_pairs: true
  empty_attention_policy: zero

Evaluate Retrieval Config

runtime:
  device: null
  seed: 42
  num_workers: 0
  use_cuda_if_available: true

encoder:
  model_name_or_path: sentence-transformers/all-MiniLM-L6-v2
  max_length: 4096
  normalize_embeddings: true
  trust_remote_code: true

retrieval:
  index_dir: indexes
  index_name: demo
  query: "dummy"
  query_path: null
  output_path: null
  top_k: 10
  per_document_top_k: 5
  beam_width: 5
  threshold: 0.0
  collect_strategy: threshold
  include_root: false
  multi_document: true

reranker:
  enabled: false
  type: none
  model_name_or_path: null
  max_length: 512
  batch_size: 8
  metadata_score_key: null
  trust_remote_code: true

evaluation:
  task: retrieval
  examples_path: data/retrieval_eval.jsonl
  index_dir: indexes
  index_name: demo
  output_path: outputs/retrieval_eval.json
  ks: [1, 3, 5]
  include_optional_generation_metrics: false
  include_bertscore: false
  bertscore_model_type: null

Evaluate Generation Config

runtime:
  device: null
  seed: 42
  num_workers: 0
  use_cuda_if_available: true

encoder:
  model_name_or_path: sentence-transformers/all-MiniLM-L6-v2
  max_length: 4096
  normalize_embeddings: true
  trust_remote_code: true

retrieval:
  index_dir: indexes
  index_name: demo
  query: "dummy"
  query_path: null
  output_path: null
  top_k: 10
  per_document_top_k: 5
  beam_width: 5
  threshold: 0.0
  collect_strategy: threshold
  include_root: false
  multi_document: true

reranker:
  enabled: false
  type: none
  model_name_or_path: null
  max_length: 512
  batch_size: 8
  metadata_score_key: null
  trust_remote_code: true

context:
  include_scores: true
  include_metadata: false
  include_node_ids: true
  context_separator: "\n\n"
  max_context_chars: 12000

generator:
  type: echo
  model_name_or_path: null
  max_input_length: 4096
  max_new_tokens: 256
  temperature: 0.0
  do_sample: false
  include_system_prompt: true
  system_prompt: null
  trust_remote_code: true

evaluation:
  task: generation
  examples_path: data/generation_eval.jsonl
  index_dir: indexes
  index_name: demo
  output_path: outputs/generation_eval.json
  ks: [1, 3, 5]
  include_optional_generation_metrics: false
  include_bertscore: false
  bertscore_model_type: null

Running the CLI

Index documents

sproutrag index --config configs/index.yaml

Retrieve results

sproutrag retrieve --config configs/retrieve.yaml

Answer queries

sproutrag answer --config configs/answer.yaml

Train

sproutrag train --config configs/train.yaml

Evaluate retrieval

sproutrag evaluate --config configs/evaluate_retrieval.yaml

Evaluate generation

sproutrag evaluate --config configs/evaluate_generation.yaml

Override config values from the CLI

sproutrag retrieve \
  --config configs/retrieve.yaml \
  --override retrieval.top_k=20 \
  --override retrieval.beam_width=8

Notes

• All outputs are JSON for easy integration with downstream tools.
• Keep configs in version control to make runs reproducible.
• For MS MARCO training, the training pipeline currently loads the first 30k examples from the v2.1 train split.

Citation

@misc{abaskohi2026sproutragattentionguidedtreesearch,
      title={SproutRAG: Attention-Guided Tree Search with Progressive Embeddings for Long-Document RAG}, 
      author={Amirhossein Abaskohi and Issam H. Laradji and Peter West and Giuseppe Carenini},
      year={2026},
      eprint={2606.18381},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2606.18381}, 
}