README.md

metadata

title: BRAG - Bayesian RAG
emoji: 🎯
colorFrom: indigo
colorTo: blue
sdk: gradio
sdk_version: 6.31.0
python_version: '3.11'
app_file: app.py
pinned: false
short_description: Closed-form Bayesian orchestration for RAG

BRAG — Bayesian RAG

BRAG orchestrates retrieval-augmented question answering as a partially observable decision problem solved in closed form. The orchestrator makes no LLM calls. Each decision is a few arithmetic operations on conjugate (Beta) posteriors and takes microseconds. Only the actions call the LLM.

Actions

• R1: read the top-k BM25 passages and return an answer or NONE, plus a follow-up search query.
• R2: retrieve with the follow-up query, then read the old and new passages (the second hop).
• CB: take m parallel closed-book samples, which count as one vote (their mode). The vote's reliability depends on how many of the samples agree.

Posterior (per-query reliabilities integrated out): log w_i = Σ_t [lnB(a_t+n_ti, b_t+N_t−n_ti) − lnB(a_t, b_t+N_t)] − n_i ln η, with residual −ln η.

Control

• Choose a* = argmax_a D_a / c_a, where D_a = q_a(ρ_a − (1−ρ_a)η) ln(ρ_a / ((1−ρ_a)η)) is the KL information rate of action a.
• Stop when P_err · D_a* / c_a* ≤ λ.

Hyperpriors

• The reliabilities ρ, the answerability rates q, the agreement-level reliabilities κ_g and the distractor concentration η are population-level.
• They are calibrated on labelled data and frozen; they are not self-trained.
• The shipped calibration (data/calibration.json) was fit with Claude Haiku as the LLM. When you use another model, calibrate first with the /calibrate endpoint.

Benchmark (300 test questions, Claude Haiku, BM25 over HippoRAG corpora)

Method	Acc	F1	LLM calls	Tokens
Closed-book	0.270	0.275	1.00	69
Vanilla RAG (1 hop)	0.387	0.377	1.66	591
IRCoT / self-ask style (always 2 hops)	0.553	0.548	2.38	1561
Adaptive / CRAG-style (hop if needed)	0.550	0.529	2.05	1221
BRAG, λ = 0.05	0.567	0.545	3.26	1622

• Against the best baseline, the accuracy gain is +1.3 to +1.7 points. Paired-bootstrap 95% confidence intervals include zero, so the gain is not statistically significant.
• The orchestrator's overhead is about 13 µs per decision.
• The deployed code reproduces the benchmark exactly: replayed on the cached LLM outputs it gives the same accuracy and call count.

Programmatic use

from gradio_client import Client
c = Client("Angshul/BRAG")
answer, report = c.predict(
    question="Friedrich Ratzel coined the phrase Lebensraum; which political movement made it notorious?",
    corpus="HotpotQA",          # HotpotQA | 2WikiMultihopQA | MuSiQue
    documents="",               # or your own passages: one per line, or a JSON list
    lam=0.05,                   # cost of one LLM call in units of answer error
    model="Qwen/Qwen2.5-7B-Instruct",   # HF model id, or an OpenAI-compatible endpoint URL
    hf_token="",                # your own token, or empty to use the Space secret
    provider="auto",
    calibration="",             # JSON returned by /calibrate for this model
    api_name="/brag")

report contains, for every step:

• the posterior over candidate answers;
• the information rate of each action;
• the action taken and what it observed;
• the time the policy took.

It also contains the total LLM calls, tokens and wall-clock time.

To calibrate for a new model or corpus:

cal = c.predict(labelled_json='[{"question": "...", "gold": ["..."]}, ...]',   # 30+ questions
                corpus="HotpotQA", documents="", model="<model>", hf_token="", provider="auto",
                api_name="/calibrate")

REST:

POST /gradio_api/call/brag   {"data": [question, corpus, documents, lam, model, hf_token, provider, calibration]}
GET  /gradio_api/call/brag/<event_id>

Scaling

• The policy is not the bottleneck. It is stateless per request and needs no GPU. Throughput is bounded by the LLM provider's rate limits and by the Space's single replica.
• Concurrency. The Space queues requests and runs BRAG_CONCURRENCY of them at once (default 16). Each request's closed-book samples run in parallel.
• Callers can bring their own quota. Pass hf_token so each caller is billed for their own inference.
• Higher throughput:
  • Point model at a dedicated, autoscaling Inference Endpoint, or at a vLLM or TGI server.
  • Duplicate the Space.
  • Import brag.py directly. It is a single dependency-light module: BRAG(llm).run(question, Corpus(passages)).

Setup

1. Add a Space secret HF_TOKEN, with permission to call Inference Providers.
2. Optionally set the variables BRAG_MODEL and BRAG_CONCURRENCY.

Data

The corpora are from the HippoRAG reproduction set (OSU-NLP-Group/HippoRAG):

• HotpotQA (CC BY-SA 4.0)
• 2WikiMultihopQA (Apache-2.0)
• MuSiQue (CC BY 4.0)