"""Standalone inference for Grounded Pointer QA (proqa). Self-contained: everything needed to load the checkpoint and ask questions over your own documents. Requires: torch, transformers, scikit-learn, numpy. from modeling_proqa import GroundedQA qa = GroundedQA("proqa.pt") # or a hf_hub_download path qa.load_folder(r"C:\\my\\notes") # .txt / .md files print(qa.ask("who approved the budget?")) # quote + source, or None """ import os from dataclasses import dataclass import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from sklearn.feature_extraction.text import TfidfVectorizer from transformers import AutoConfig, AutoModel, AutoTokenizer @dataclass class ProConfig: backbone: str = "roberta-base" max_len: int = 384 k_passages: int = 4 class ProReaderQA(nn.Module): def __init__(self, cfg: ProConfig): super().__init__() self.cfg = cfg self.backbone = AutoModel.from_config(AutoConfig.from_pretrained(cfg.backbone)) h = self.backbone.config.hidden_size self.span_head = nn.Linear(h, 2) self.abstain_head = nn.Sequential(nn.Linear(h, h), nn.Tanh(), nn.Linear(h, 1)) def forward(self, input_ids, attention_mask, context_mask): b, k, L = input_ids.shape out = self.backbone(input_ids=input_ids.view(b * k, L), attention_mask=attention_mask.view(b * k, L) ).last_hidden_state start_logits, end_logits = self.span_head(out).unbind(dim=-1) neg_inf = torch.finfo(start_logits.dtype).min cm = context_mask.view(b * k, L) start_logits = start_logits.masked_fill(~cm, neg_inf).view(b, k * L) end_logits = end_logits.masked_fill(~cm, neg_inf).view(b, k * L) cls = out[:, 0].view(b, k, -1).mean(dim=1) return start_logits, end_logits, self.abstain_head(cls).squeeze(-1) def chunk_text(text, chunk_words=150, overlap=40): words = text.split() if not words: return [] chunks, step = [], max(chunk_words - overlap, 1) for i in range(0, len(words), step): chunks.append(" ".join(words[i:i + chunk_words])) if i + chunk_words >= len(words): break return chunks class GroundedQA: def __init__(self, checkpoint_path: str, device: str = None): self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") ckpt = torch.load(checkpoint_path, map_location=self.device) self.cfg = ProConfig(**ckpt["config"]) self.model = ProReaderQA(self.cfg).to(self.device).eval() self.model.load_state_dict(ckpt["state_dict"]) self.gate = ckpt.get("gate_threshold", 0.5) self.tok = AutoTokenizer.from_pretrained(self.cfg.backbone) self.passages, self._vec, self._mat = [], None, None # ---------- knowledge ---------- def load_passages(self, passages): self.passages = list(passages) self._vec = TfidfVectorizer(lowercase=True, ngram_range=(1, 2), sublinear_tf=True, min_df=1) self._mat = self._vec.fit_transform(self.passages) def load_folder(self, folder, chunk_words=150): passages = [] for root, _, files in os.walk(folder): for name in sorted(files): path = os.path.join(root, name) if name.lower().endswith((".txt", ".md")): with open(path, encoding="utf-8", errors="ignore") as f: passages.extend(chunk_text(f.read(), chunk_words)) elif name.lower().endswith(".pdf"): try: from pypdf import PdfReader text = "\n".join(p.extract_text() or "" for p in PdfReader(path).pages) passages.extend(chunk_text(text, chunk_words)) except ImportError: pass # pip install pypdf for PDF support if not passages: raise ValueError(f"no readable documents under {folder}") self.load_passages(passages) # ---------- ask ---------- @torch.no_grad() def ask(self, question: str, gate: float = None): """Returns dict(answer, source, confidence) or dict(answer=None, ...).""" assert self.passages, "load knowledge first (load_folder / load_passages)" gate = self.gate if gate is None else gate k, L = self.cfg.k_passages, self.cfg.max_len q = self._vec.transform([question]) sims = (q @ self._mat.T).toarray()[0] top = list(np.argsort(-sims)[:k]) while len(top) < k: top.append(top[-1]) passages = [self.passages[j] for j in top] q_ids = self.tok(question, add_special_tokens=False, truncation=True, max_length=64)["input_ids"] question = self.tok.decode(q_ids) enc = self.tok([question] * k, passages, truncation="only_second", max_length=L, padding="max_length", return_offsets_mapping=True, return_tensors="pt") cm = torch.zeros(k, L, dtype=torch.bool) for s in range(k): cm[s] = torch.tensor([sid == 1 for sid in enc.sequence_ids(s)]) if top[s] in top[:s]: # dedup tiny indexes cm[s] = False with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=self.device == "cuda"): s_log, e_log, a_log = self.model( enc["input_ids"].unsqueeze(0).to(self.device), enc["attention_mask"].bool().unsqueeze(0).to(self.device), cm.unsqueeze(0).to(self.device)) s_lp = F.log_softmax(s_log.float(), -1).view(1, k, L) e_lp = F.log_softmax(e_log.float(), -1).view(1, k, L) scores = s_lp.unsqueeze(3) + e_lp.unsqueeze(2) valid = torch.ones(L, L, dtype=torch.bool, device=scores.device).triu() valid &= ~torch.ones(L, L, dtype=torch.bool, device=scores.device).triu(80) flat = scores.masked_fill(~valid, float("-inf")).view(1, -1) best, idx = flat.max(-1) pi, rem = int(idx // (L * L)), int(idx % (L * L)) s, e = rem // L, rem % L conf = float(a_log.float().sigmoid() * best.exp()) if conf < gate: return {"answer": None, "source": None, "confidence": conf} o = enc["offset_mapping"][pi] return {"answer": passages[pi][int(o[s][0]): int(o[e][1])], "source": passages[pi], "confidence": conf} if __name__ == "__main__": import sys qa = GroundedQA(sys.argv[1] if len(sys.argv) > 1 else "proqa.pt") qa.load_folder(sys.argv[2] if len(sys.argv) > 2 else ".") print(f"loaded {len(qa.passages)} passages; gate={qa.gate:.2f}") while True: q = input("question> ").strip() if q in ("", "/quit"): break r = qa.ask(q) if r["answer"] is None: print(f" [abstains] (conf={r['confidence']:.2f})") else: print(f" \"{r['answer']}\" (conf={r['confidence']:.2f})")