| --- |
| 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 |
|
|
| ```python |
| 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: |
|
|
| ```python |
| 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) |
|
|
