Grounded Pointer QA: checkpoint, standalone inference, model card

README.md

@@ -0,0 +1,97 @@
+---
+language: en
+license: mit
+library_name: pytorch
+base_model: FacebookAI/roberta-base
+pipeline_tag: question-answering
+tags:
+  - extractive-qa
+  - grounded-qa
+  - hallucination-free
+  - selective-prediction
+  - abstention
+  - rag
+datasets:
+  - rajpurkar/squad_v2
+metrics:
+  - exact_match
+model-index:
+  - name: grounded-pointer-qa
+    results:
+      - task:
+          type: question-answering
+        dataset:
+          name: SQuAD v2 (held-out test half, retrieval setting)
+          type: rajpurkar/squad_v2
+        metrics:
+          - type: exact_match
+            value: 74.6
+            name: EM (EM-optimal gate)
+          - type: precision
+            value: 91.7
+            name: Answered precision @ 90%-precision gate (9% coverage)
+---
+
+# Grounded Pointer QA
+
+An extractive question-answering model that **cannot hallucinate by
+construction**: its output layer can only point at spans inside retrieved
+passages of *your* documents — it has no vocabulary to generate from. A
+trained abstention head refuses when the loaded knowledge doesn't contain the
+answer, decoding is deterministic (argmax; same question + same documents =
+same answer forever), and knowledge is **hot-swappable**: point it at a new
+folder of text files and it answers from those, no retraining.
+
+Built on `roberta-base` (125M params) with pointer + abstention heads,
+finetuned on SQuAD v2 **against real TF-IDF retrieval** — the model never saw
+gold passages during training, only what the retriever actually returned, so
+it learned to abstain on retrieval misses too.
+
+## Operating modes
+
+The checkpoint ships with a calibrated confidence gate
+(`P(answerable) × P(span)`), selected on a calibration split and verified on
+a held-out test split:
+
+| Mode | Gate | Coverage | Answered precision | EM |
+|---|---|---|---|---|
+| "Right or silent" (shipped default) | 0.965 | 9% | **91.7%** | 57.1 |
+| EM-optimal | 0.295 | 45% | 72.3% | 74.6 |
+
+Pass a lower `gate` to `ask()` for more coverage at lower precision.
+
+## Usage
+
+```python
+# files needed: proqa.pt, modeling_proqa.py (both in this repo)
+# pip install torch transformers scikit-learn numpy  (+ pypdf for PDFs)
+from modeling_proqa import GroundedQA
+
+qa = GroundedQA("proqa.pt")
+qa.load_folder("path/to/your/notes")     # .txt / .md / .pdf
+
+qa.ask("when does the vendor contract expire?")
+# {'answer': '30 November 2026', 'source': '...expires on 30 November 2026...',
+#  'confidence': 0.98}
+
+qa.ask("what is the capital of France?")   # not in your docs
+# {'answer': None, 'source': None, 'confidence': 0.0}
+```
+
+## Honest limitations
+
+- Single-span extraction only: no summarization, no aggregation across
+  passages, no multi-turn conversation.
+- Trained on Wikipedia-style prose; **tables read poorly** — convert rows to
+  sentences ("ATA chapter 32 is Landing Gear.") for big accuracy gains.
+- TF-IDF retrieval is lexical: paraphrases sharing no words with your
+  documents may cause (safe) abstentions.
+- Single training run, single seed; in-domain calibration.
+
+Full writeup, ablations (from-scratch 22.4 EM → 74.6 EM ladder), and negative
+results: see the project repository's `paper/`.
+
+## Training
+
+One NVIDIA RTX 5060 Ti (16 GB): ~2.5 h finetune (batch 8 × 4 passages × 384
+tokens, bf16, lr 2e-5, 2 epochs) + calibration pass.

modeling_proqa.py

@@ -0,0 +1,168 @@
+"""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})")

proqa.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27f876d2b24a30f86f4afb3c4328a1fec22c3374c993f0abbfc4d36d298fec90
+size 501023103