| """BRAG - Bayesian RAG. |
| |
| Closed-form Bayesian orchestration for retrieval-augmented QA. The orchestrator itself makes no |
| LLM calls: every decision is a few arithmetic operations on conjugate (Beta) posteriors. |
| |
| Actions (each is an LLM call made by the *executor*, never by the policy): |
| R1 : read top-k retrieved passages -> answer or NONE, plus a follow-up search query |
| R2 : retrieve with the follow-up query, read old + new passages -> answer or NONE (needs R1) |
| CB : m parallel closed-book samples -> ONE vote (their mode), reliability by agreement level |
| |
| Relaxations |
| R1' self-samples of one LLM are one source (vote = mode, reliability kappa_g by agreement g) |
| R2' concentrated distractors: a wrong vote lands on a given rival w.p. eta (K_eff = 1/eta) |
| R3' reliability hyperpriors are population-level, calibrated offline on labelled data, frozen |
| R4' separate reliabilities per action type |
| R5' control by information rate: a* = argmax_a D_a / c_a, stop when P_err * D_a* / c_a* <= lambda |
| Posterior (per-query reliabilities integrated out in closed form): |
| log w_i = sum_t [lnB(a_t+n_ti, b_t+N_t-n_ti) - lnB(a_t, b_t+N_t)] - n_i ln eta ; log w_residual = -ln eta |
| """ |
| from __future__ import annotations |
| import json, math, re, string, threading, time |
| from collections import Counter |
| from concurrent.futures import ThreadPoolExecutor |
| from dataclasses import dataclass, field |
| from typing import Callable |
|
|
| import numpy as np |
| from scipy.special import betaln |
| from rank_bm25 import BM25Okapi |
|
|
| |
| _PUNCT = set(string.punctuation) |
| _STOP = set('a an the of in on at to for by with from and or is was were are be been who what which when ' |
| 'where whose how did does do that this as its it his her their'.split()) |
|
|
| def norm(s: str) -> str: |
| s = (s or '').lower() |
| s = ''.join(ch for ch in s if ch not in _PUNCT) |
| s = re.sub(r'\b(a|an|the)\b', ' ', s) |
| return ' '.join(s.split()) |
|
|
| def is_none(a) -> bool: |
| return a is None or norm(str(a)) in ('none', '', 'unknown', 'not stated', 'not mentioned') |
|
|
| def tokenize(s: str) -> list[str]: |
| return [w for w in re.findall(r'\w+', s.lower()) if w not in _STOP] |
|
|
| def est_tokens(s: str) -> int: |
| return max(1, len(s) // 4) |
|
|
| class Clusters: |
| """Answer strings -> clusters (equal after normalisation, or token containment).""" |
| def __init__(self): |
| self.reps: list[str] = [] |
| def id(self, a: str) -> int: |
| n = norm(a); nt = set(n.split()) |
| for i, r in enumerate(self.reps): |
| rt = set(r.split()) |
| if n == r or (nt and rt and (nt <= rt or rt <= nt)): |
| return i |
| self.reps.append(n) |
| return len(self.reps) - 1 |
|
|
| |
| class Corpus: |
| def __init__(self, passages: list[str]): |
| if not passages: |
| raise ValueError('corpus is empty') |
| self.passages = passages |
| self.bm25 = BM25Okapi([tokenize(p) or ['_'] for p in passages]) |
|
|
| @classmethod |
| def from_json(cls, path: str) -> 'Corpus': |
| rows = json.load(open(path)) |
| return cls([f"{r['title']}: {r['text']}" for r in rows]) |
|
|
| def search(self, query: str, k: int, exclude: set[int] = frozenset()) -> list[int]: |
| sc = self.bm25.get_scores(tokenize(query) or ['_']) |
| order = np.argsort(-sc) |
| return [int(j) for j in order[:k + len(exclude)] if int(j) not in exclude][:k] |
|
|
| |
| LLM = Callable[[list[dict], float, int], tuple[str, int, int]] |
|
|
| SYS = 'You are a precise question-answering component. Follow the output format exactly.' |
|
|
| def p_cb(q): |
| return [{'role': 'system', 'content': SYS}, |
| {'role': 'user', 'content': f'Answer from your own knowledge. Give only a short answer (1-6 words), ' |
| f'always a best guess.\nQuestion: {q}\nAnswer:'}] |
|
|
| def _ctx(ps): |
| return '\n'.join(f'[{i + 1}] {p}' for i, p in enumerate(ps)) |
|
|
| def p_r1(q, ps): |
| return [{'role': 'system', 'content': SYS}, |
| {'role': 'user', 'content': |
| 'Use ONLY the passages. Return JSON {"answer": "...", "followup": "..."} where answer is a short ' |
| 'answer (1-6 words) if the passages TOGETHER answer the question, else "NONE"; followup is the single ' |
| 'most useful next search query for the missing information (name the bridge entity if known). ' |
| f'Always give a followup.\n\nPassages:\n{_ctx(ps)}\n\nQuestion: {q}\nJSON:'}] |
|
|
| def p_r2(q, ps): |
| return [{'role': 'system', 'content': SYS}, |
| {'role': 'user', 'content': |
| 'Use ONLY the passages; combine facts across passages for multi-hop questions. Give only a short ' |
| 'answer (1-6 words), or NONE if the passages do not contain the answer.' |
| f'\n\nPassages:\n{_ctx(ps)}\n\nQuestion: {q}\nAnswer:'}] |
|
|
| def clean_answer(t: str) -> str: |
| t = (t or '').strip().strip('"').strip() |
| t = re.sub(r'^(final\s+)?answer\s*:\s*', '', t, flags=re.I) |
| return t.split('\n')[0].strip().strip('."') |
|
|
| def parse_r1(t: str) -> tuple[str, str]: |
| m = re.search(r'\{.*\}', t or '', flags=re.S) |
| if m: |
| try: |
| d = json.loads(m.group(0)) |
| return clean_answer(str(d.get('answer', 'NONE'))), str(d.get('followup', '')).strip() |
| except json.JSONDecodeError: |
| pass |
| a = re.search(r'"?answer"?\s*[:=]\s*"([^"]*)"', t or '', flags=re.I) |
| f = re.search(r'"?followup"?\s*[:=]\s*"([^"]*)"', t or '', flags=re.I) |
| return (clean_answer(a.group(1)) if a else 'NONE'), (f.group(1).strip() if f else '') |
|
|
| def hf_llm(model: str, token: str | None, provider: str = 'auto', timeout: float = 60.0) -> LLM: |
| """model: a Hugging Face model id (served by HF Inference Providers) or the URL of an |
| OpenAI-compatible endpoint (e.g. a dedicated HF Inference Endpoint, TGI or vLLM server).""" |
| from huggingface_hub import InferenceClient |
| if model.startswith(('http://', 'https://')): |
| client = InferenceClient(base_url=model, api_key=token, timeout=timeout) |
| else: |
| client = InferenceClient(model=model, token=token, provider=provider, timeout=timeout) |
| def call(messages, temperature, max_tokens): |
| r = client.chat_completion(messages, max_tokens=max_tokens, temperature=max(temperature, 0.01)) |
| text = r.choices[0].message.content or '' |
| u = getattr(r, 'usage', None) |
| tin = getattr(u, 'prompt_tokens', None) or sum(est_tokens(m['content']) for m in messages) |
| tout = getattr(u, 'completion_tokens', None) or est_tokens(text) |
| return text, int(tin), int(tout) |
| return call |
|
|
| |
| DEFAULT_CAL = { |
| 'rel': {'R1': [29, 6], 'R2': [54, 11]}, 'ans': {'R1': [34, 58], 'R2': [64, 28]}, |
| 'kap': {'1': [2, 30], '2': [15, 24], '3': [14, 11]}, 'pg': {'1': 0.333, '2': 0.409, '3': 0.258}, |
| 'eta': 0.591} |
|
|
| def _kl(rho, eta, q=1.0): |
| rho = min(max(rho, 1e-6), 1 - 1e-6) |
| return q * (rho - (1 - rho) * eta) * math.log(rho / ((1 - rho) * eta)) |
|
|
| def gold_match(pred, gold: list[str]) -> int: |
| if is_none(pred): |
| return 0 |
| p = norm(pred) |
| return int(any(norm(g) and (norm(g) in p or (len(p) > 3 and p in norm(g))) for g in gold)) |
|
|
| |
| @dataclass |
| class Result: |
| answer: str | None |
| trace: list = field(default_factory=list) |
| llm_calls: int = 0 |
| tokens_in: int = 0 |
| tokens_out: int = 0 |
| rounds: int = 0 |
| llm_seconds: float = 0.0 |
| policy_microseconds: float = 0.0 |
| stopped_because: str = '' |
| def to_dict(self): |
| return dict(answer=self.answer, stopped_because=self.stopped_because, llm_calls=self.llm_calls, |
| tokens_in=self.tokens_in, tokens_out=self.tokens_out, rounds=self.rounds, |
| llm_wall_seconds=round(self.llm_seconds, 3), |
| policy_microseconds_total=round(self.policy_microseconds, 1), trace=self.trace) |
|
|
| class BRAG: |
| def __init__(self, llm: LLM, calibration: dict | None = None, lam: float = 0.05, |
| k: int = 5, cb_samples: int = 3, cb_temperature: float = 0.7): |
| c = json.loads(json.dumps(calibration or DEFAULT_CAL)) |
| c['kap'] = {int(g): v for g, v in c['kap'].items()}; c['pg'] = {int(g): v for g, v in c['pg'].items()} |
| self.P, self.llm, self.lam, self.k = c, llm, lam, k |
| self.m, self.temp = cb_samples, cb_temperature |
|
|
| |
| def posterior(self, votes): |
| eta = self.P['eta']; cl = sorted({i for _, i in votes}); lw = {i: 0.0 for i in cl} |
| by = {} |
| for t, i in votes: |
| by.setdefault(t, Counter())[i] += 1 |
| for t, cnt in by.items(): |
| a, b = self.P['kap'][t] if isinstance(t, int) else self.P['rel'][t]; N = sum(cnt.values()) |
| for i in cl: |
| n = cnt.get(i, 0) |
| lw[i] += betaln(a + n, b + N - n) - betaln(a, b + N) - n * math.log(eta) |
| v = np.array([lw[i] for i in cl] + [-math.log(eta)]); v -= v.max(); p = np.exp(v); p /= p.sum() |
| return cl, p |
|
|
| def rates(self, done): |
| eta = self.P['eta']; out = {} |
| for t in ('R1', 'R2'): |
| if t in done or (t == 'R2' and 'R1' not in done): |
| continue |
| a, b = self.P['ans'][t]; ra, rb = self.P['rel'][t] |
| out[t] = _kl(ra / (ra + rb), eta, a / (a + b)) / 1.0 |
| if 'CB' not in done: |
| d = sum(self.P['pg'][g] * _kl(self.P['kap'][g][0] / sum(self.P['kap'][g]), eta) for g in (1, 2, 3)) |
| out['CB'] = d / self.m |
| return out |
|
|
| _lock = threading.Lock() |
|
|
| def _call(self, res, messages, temperature, max_tokens): |
| t0 = time.perf_counter(); text, ti, to = self.llm(messages, temperature, max_tokens) |
| with self._lock: |
| res.llm_calls += 1 |
| res.tokens_in += ti; res.tokens_out += to |
| return text |
|
|
| def run(self, question: str, corpus: Corpus) -> Result: |
| res = Result(answer=None); C = Clusters(); votes = []; raw = {}; done = set(); st = {} |
| while True: |
| t0 = time.perf_counter() |
| cl, p = self.posterior(votes) |
| pbest = float(p[:-1].max()) if cl else 0.0; perr = 1 - pbest |
| rates = self.rates(done); act = None; why = '' |
| if not rates: |
| why = 'no actions left' |
| else: |
| a_star = max(rates, key=rates.get) |
| if perr * rates[a_star] > self.lam: |
| act = a_star |
| else: |
| why = f'P_err*D/c = {perr * rates[a_star]:.4f} <= lambda = {self.lam}' |
| dt = 1e6 * (time.perf_counter() - t0); res.policy_microseconds += dt |
| step = dict(posterior={(C.reps[i] if i < len(C.reps) else str(i)): round(float(pp), 4) |
| for i, pp in zip(cl, p[:-1])}, |
| residual=round(float(p[-1]), 4), p_error=round(perr, 4), |
| info_rate_per_call={k_: round(v, 4) for k_, v in rates.items()}, |
| action=act or 'STOP', policy_us=round(dt, 1)) |
| if act is None: |
| res.stopped_because = why; res.trace.append(step); break |
| done.add(act); res.rounds += 1; tw = time.perf_counter() |
| if act == 'CB': |
| with ThreadPoolExecutor(self.m) as ex: |
| outs = list(ex.map(lambda _: self._call(res, p_cb(question), self.temp, 24), range(self.m))) |
| samples = [clean_answer(o) for o in outs]; good = [s for s in samples if not is_none(s)] |
| step['observation'] = {'samples': samples} |
| if good: |
| ids = [C.id(s) for s in good]; mo, g = Counter(ids).most_common(1)[0] |
| votes.append((g, mo)); raw.setdefault(mo, good[ids.index(mo)]) |
| step['observation']['vote'] = good[ids.index(mo)]; step['observation']['agreement'] = g |
| elif act == 'R1': |
| idx = corpus.search(question, self.k); st['idx1'] = idx |
| ps = [corpus.passages[j][:600] for j in idx] |
| ans, fu = parse_r1(self._call(res, p_r1(question, ps), 0.0, 96)); st['followup'] = fu or question |
| step['observation'] = {'answer': ans, 'followup': fu, 'passages': [x[:80] for x in ps]} |
| if not is_none(ans): |
| cid = C.id(ans); votes.append(('R1', cid)); raw.setdefault(cid, ans) |
| else: |
| idx2 = corpus.search(st['followup'], self.k, exclude=set(st['idx1'])) |
| ps = [corpus.passages[j][:600] for j in st['idx1'] + idx2] |
| ans = clean_answer(self._call(res, p_r2(question, ps), 0.0, 32)) |
| step['observation'] = {'answer': ans, 'new_passages': [corpus.passages[j][:80] for j in idx2]} |
| if not is_none(ans): |
| cid = C.id(ans); votes.append(('R2', cid)); raw.setdefault(cid, ans) |
| step['wall_seconds'] = round(time.perf_counter() - tw, 3); res.llm_seconds += step['wall_seconds'] |
| res.trace.append(step) |
| cl, p = self.posterior(votes) |
| if cl: |
| res.answer = raw[cl[int(np.argmax(p[:-1]))]] |
| else: |
| tw = time.perf_counter() |
| res.answer = clean_answer(self._call(res, p_cb(question), 0.0, 24)); res.rounds += 1 |
| res.llm_seconds += time.perf_counter() - tw |
| res.stopped_because += ' | forced closed-book answer' |
| return res |
|
|
| |
| def calibrate(self, labelled: list[dict], corpus: Corpus) -> dict: |
| """labelled: [{question, gold:[aliases]}]. Runs every action once per question and returns |
| calibration pseudo-counts (+1 smoothing) for this LLM, retriever and corpus.""" |
| rel = {'R1': [1, 1], 'R2': [1, 1]}; ans = {'R1': [1, 1], 'R2': [1, 1]} |
| kap = {g: [1, 1] for g in (1, 2, 3)}; pg = Counter({1: 1, 2: 1, 3: 1}); pairs = coll = 0 |
| dummy = Result(answer=None) |
| for ex_ in labelled: |
| q, gold = ex_['question'], ex_['gold'] if isinstance(ex_['gold'], list) else [ex_['gold']] |
| idx = corpus.search(q, self.k); ps = [corpus.passages[j][:600] for j in idx] |
| a1, fu = parse_r1(self._call(dummy, p_r1(q, ps), 0.0, 96)) |
| idx2 = corpus.search(fu or q, self.k, exclude=set(idx)) |
| a2 = clean_answer(self._call(dummy, p_r2(q, ps + [corpus.passages[j][:600] for j in idx2]), 0.0, 32)) |
| for t, a in (('R1', a1), ('R2', a2)): |
| if is_none(a): ans[t][1] += 1 |
| else: ans[t][0] += 1; c = gold_match(a, gold); rel[t][0] += c; rel[t][1] += 1 - c |
| with ThreadPoolExecutor(self.m) as ex: |
| outs = [clean_answer(o) for o in ex.map(lambda _: self._call(dummy, p_cb(q), self.temp, 24), range(self.m))] |
| good = [s for s in outs if not is_none(s)] |
| if good: |
| C = Clusters(); ids = [C.id(s) for s in good]; mo, g = Counter(ids).most_common(1)[0] |
| ok = gold_match(good[ids.index(mo)], gold); kap[g][0] += ok; kap[g][1] += 1 - ok; pg[g] += 1 |
| wrong = [x for x in (a1, a2, outs[0] if outs else None) if not is_none(x) and not gold_match(x, gold)] |
| C = Clusters(); ids = [C.id(x) for x in wrong] |
| for i in range(len(ids)): |
| for j in range(i + 1, len(ids)): pairs += 1; coll += ids[i] == ids[j] |
| tot = sum(pg.values()) |
| return dict(rel=rel, ans=ans, kap={str(g): v for g, v in kap.items()}, |
| pg={str(g): pg[g] / tot for g in pg}, eta=(1 + coll) / (2 + pairs), |
| n_questions=len(labelled), llm_calls=dummy.llm_calls) |
|
|
