modeling_proqa.py

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