---
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)