File size: 5,034 Bytes
44dbefd 1a3c3b3 44dbefd 1a3c3b3 44dbefd 1a3c3b3 44dbefd 1a3c3b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | ---
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)
|