| """BRAG - Bayesian RAG: Gradio Space (live mode) with a programmatic API.""" |
| import json, os, hashlib, threading |
| from functools import lru_cache |
| import gradio as gr |
| from brag import BRAG, Corpus, hf_llm |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| CORPORA = {'HotpotQA': 'hotpotqa_corpus.json', '2WikiMultihopQA': '2wikimultihopqa_corpus.json', |
| 'MuSiQue': 'musique_corpus.json'} |
| DEFAULT_MODEL = os.getenv('BRAG_MODEL', 'Qwen/Qwen2.5-7B-Instruct') |
| CAL_PATH = os.path.join(HERE, 'data', 'calibration.json') |
| DEFAULT_CAL = json.load(open(CAL_PATH)) if os.path.exists(CAL_PATH) else None |
| _lock = threading.Lock() |
|
|
| @lru_cache(maxsize=len(CORPORA)) |
| def builtin_corpus(name): |
| return Corpus.from_json(os.path.join(HERE, 'data', CORPORA[name])) |
|
|
| @lru_cache(maxsize=32) |
| def custom_corpus(digest, text): |
| text = text.strip() |
| if text.startswith('['): |
| rows = json.loads(text) |
| ps = [r if isinstance(r, str) else f"{r.get('title', '')}: {r.get('text', '')}" for r in rows] |
| else: |
| ps = [l.strip() for l in text.splitlines() if l.strip()] |
| return Corpus(ps) |
|
|
| def get_corpus(corpus, documents): |
| if documents and documents.strip(): |
| return custom_corpus(hashlib.sha1(documents.encode()).hexdigest(), documents) |
| if corpus not in CORPORA: |
| raise gr.Error(f'corpus must be one of {list(CORPORA)} or documents must be given') |
| with _lock: |
| return builtin_corpus(corpus) |
|
|
| def get_llm(model, hf_token, provider): |
| if os.getenv('BRAG_MOCK') == '1': |
| from mock_llm import mock_llm |
| return mock_llm |
| token = (hf_token or '').strip() or os.getenv('HF_TOKEN') |
| if not token: |
| raise gr.Error('No Hugging Face token: pass hf_token, or set the HF_TOKEN secret on the Space.') |
| return hf_llm((model or DEFAULT_MODEL).strip(), token, (provider or 'auto').strip()) |
|
|
| def parse_cal(calibration): |
| if calibration and str(calibration).strip(): |
| try: |
| return json.loads(calibration) if isinstance(calibration, str) else calibration |
| except json.JSONDecodeError as e: |
| raise gr.Error(f'calibration is not valid JSON: {e}') |
| return DEFAULT_CAL |
|
|
| def brag(question: str, corpus: str = 'HotpotQA', documents: str = '', lam: float = 0.05, |
| model: str = DEFAULT_MODEL, hf_token: str = '', provider: str = 'auto', calibration: str = ''): |
| """Answer a question with BRAG (Bayesian RAG). |
| |
| Args: |
| question: the question to answer. |
| corpus: built-in corpus to search: HotpotQA, 2WikiMultihopQA or MuSiQue (ignored if documents given). |
| documents: optional custom corpus: one passage per line, or a JSON list of strings / {title, text}. |
| lam: cost of one LLM call in units of answer error (higher = cheaper, earlier stopping). |
| model: Hugging Face model id, or URL of an OpenAI-compatible endpoint (dedicated Inference Endpoint, TGI, vLLM). |
| hf_token: optional Hugging Face token (defaults to the Space's HF_TOKEN secret). |
| provider: Hugging Face inference provider ('auto' picks one that serves the model). |
| calibration: optional JSON calibration (output of the calibrate endpoint) for this model. |
| Returns: |
| answer, and a JSON report with the per-step posterior, information rates, actions and costs. |
| """ |
| if not question or not question.strip(): |
| raise gr.Error('question is empty') |
| engine = BRAG(get_llm(model, hf_token, provider), parse_cal(calibration), lam=float(lam)) |
| try: |
| res = engine.run(question.strip(), get_corpus(corpus, documents)) |
| except gr.Error: |
| raise |
| except Exception as e: |
| raise gr.Error(f'LLM/provider error: {type(e).__name__}: {e}') |
| rep = res.to_dict(); rep['model'] = model; rep['lambda'] = float(lam) |
| return res.answer or '', rep |
|
|
| def calibrate(labelled_json: str, corpus: str = 'HotpotQA', documents: str = '', |
| model: str = DEFAULT_MODEL, hf_token: str = '', provider: str = 'auto'): |
| """Calibrate BRAG's hyperpriors for a model/corpus from labelled questions. |
| |
| Args: |
| labelled_json: JSON list of {"question": ..., "gold": [aliases]} (30+ recommended, max 200). |
| corpus / documents / model / hf_token / provider: as in the brag endpoint. |
| Returns: |
| calibration JSON to pass as `calibration` to the brag endpoint. |
| """ |
| try: |
| rows = json.loads(labelled_json) |
| assert isinstance(rows, list) and rows and all('question' in r and 'gold' in r for r in rows) |
| except Exception: |
| raise gr.Error('labelled_json must be a JSON list of {"question": ..., "gold": [...]}') |
| if len(rows) > 200: |
| raise gr.Error('at most 200 labelled questions per call') |
| engine = BRAG(get_llm(model, hf_token, provider), DEFAULT_CAL) |
| try: |
| return engine.calibrate(rows, get_corpus(corpus, documents)) |
| except Exception as e: |
| raise gr.Error(f'LLM/provider error: {type(e).__name__}: {e}') |
|
|
| API_DOC = """ |
| ### Programmatic use |
| ```python |
| from gradio_client import Client |
| c = Client("Angshul/BRAG") # or Client("Angshul/BRAG", hf_token="hf_...") |
| answer, report = c.predict( |
| question="Friedrich Ratzel coined the phrase Lebensraum; which political movement made it notorious?", |
| corpus="HotpotQA", documents="", lam=0.05, |
| model="Qwen/Qwen2.5-7B-Instruct", hf_token="", provider="auto", calibration="", |
| api_name="/brag") |
| ``` |
| Own documents: pass `documents` (one passage per line, or a JSON list). Different model: run |
| `api_name="/calibrate"` once on 30+ labelled questions and pass the returned JSON as `calibration`. |
| REST: `POST /gradio_api/call/brag` with `{"data": [question, corpus, documents, lam, model, hf_token, provider, calibration]}`. |
| """ |
|
|
| with gr.Blocks(title='BRAG - Bayesian RAG') as demo: |
| gr.Markdown('# BRAG — Bayesian RAG\nClosed-form Bayesian orchestration of retrieval-augmented QA. ' |
| 'The orchestrator makes **no LLM calls**: each decision is arithmetic on conjugate posteriors ' |
| '(microseconds). It chooses between reading retrieved passages, a second retrieval hop and ' |
| 'closed-book sampling, and stops when the expected error reduction per call falls below λ.') |
| with gr.Tab('Ask'): |
| q = gr.Textbox(label='Question', lines=2) |
| with gr.Row(): |
| corpus = gr.Dropdown(list(CORPORA), value='HotpotQA', label='Corpus') |
| lam = gr.Slider(0.005, 0.3, value=0.05, step=0.005, label='λ (cost per LLM call)') |
| docs = gr.Textbox(label='Own documents (optional): one passage per line, or JSON list', lines=3) |
| with gr.Accordion('Model', open=False): |
| model = gr.Textbox(value=DEFAULT_MODEL, label='HF model id or OpenAI-compatible endpoint URL') |
| provider = gr.Textbox(value='auto', label='Inference provider') |
| tok = gr.Textbox(label='HF token (optional; default: Space secret)', type='password') |
| cal = gr.Textbox(label='Calibration JSON (optional)', lines=2) |
| go = gr.Button('Answer', variant='primary') |
| ans = gr.Textbox(label='Answer'); rep = gr.JSON(label='Trace: posterior, information rates, actions, costs') |
| go.click(brag, [q, corpus, docs, lam, model, tok, provider, cal], [ans, rep], api_name='brag') |
| with gr.Tab('Calibrate'): |
| lab = gr.Textbox(label='Labelled questions JSON: [{"question": ..., "gold": [...]}, ...]', lines=6) |
| with gr.Row(): |
| corpus2 = gr.Dropdown(list(CORPORA), value='HotpotQA', label='Corpus') |
| model2 = gr.Textbox(value=DEFAULT_MODEL, label='HF model id') |
| docs2 = gr.Textbox(label='Own documents (optional)', lines=2) |
| tok2 = gr.Textbox(label='HF token (optional)', type='password'); prov2 = gr.Textbox(value='auto', label='Provider') |
| go2 = gr.Button('Calibrate'); out2 = gr.JSON(label='Calibration (pass to the brag endpoint)') |
| go2.click(calibrate, [lab, corpus2, docs2, model2, tok2, prov2], out2, api_name='calibrate') |
| gr.Markdown(API_DOC) |
|
|
| demo.queue(default_concurrency_limit=int(os.getenv('BRAG_CONCURRENCY', '16')), max_size=256) |
|
|
| if __name__ == '__main__': |
| demo.launch() |
|
|
