rag-wright 0.1.0__py3-none-any.whl

rag_wright/spans/legalbert_classifier.py

@@ -0,0 +1,83 @@
+"""T56 (FR-R, ADR-0025): LegalBERT fine-tuned function classifier (the upgrade over the linear head).
+
+A fine-tuned `nlpaueb/legal-bert-base-uncased` sequence classifier: an operative span's text -> one of the 41
+CUAD clause types or NONE. Self-contained (its own tokenizer + encoder), sub-second local inference (MPS/CPU),
+no LLM call. Chosen to make the classifier a controlled, high-accuracy variable (it captures fine legal
+distinctions the frozen-embedding linear head could not, e.g. Cap-on-Liability vs Uncapped-Liability), so a
+later downstream recall shortfall is attributable to the property/rerank stages, not the classifier.
+
+The tokenizer + model are injected, so the classify plumbing (batching, argmax, id->label) is unit-tested with
+fakes and no model download; training/loading use the real `transformers` classes (see
+`scripts/train_legalbert_function.py`).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+
+
+class LegalBertFunctionClassifier:
+    """A fine-tuned LegalBERT sequence classifier over operative spans (predicts a CUAD type or NONE)."""
+
+    def __init__(self, model: Any, tokenizer: Any, *, device: str = "cpu", max_length: int = 256) -> None:
+        self._model = model.to(device).eval()
+        self._tokenizer = tokenizer
+        self._device = device
+        self._max_length = max_length
+
+    @classmethod
+    def load(cls, path: Path, *, device: str = "cpu") -> "LegalBertFunctionClassifier":
+        from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(str(path))
+        model = AutoModelForSequenceClassification.from_pretrained(str(path))
+        return cls(model, tokenizer, device=device)
+
+    @torch.no_grad()
+    def classify(self, texts: list[str], *, batch_size: int = 32) -> list[str]:
+        """One label per span (a CUAD type or NONE). Empty input -> empty output. Batched forward on the device."""
+        if not texts:
+            return []
+        id2label = self._model.config.id2label
+        out: list[str] = []
+        for start in range(0, len(texts), batch_size):
+            batch = texts[start : start + batch_size]
+            enc = self._tokenizer(
+                batch, truncation=True, max_length=self._max_length, padding=True, return_tensors="pt"
+            ).to(self._device)
+            logits = self._model(**enc).logits
+            out.extend(str(id2label[int(i)]) for i in logits.argmax(dim=-1).tolist())
+        return out
+
+    @torch.no_grad()
+    def classify_topk(self, texts: list[str], *, k: int = 2, batch_size: int = 32) -> list[list[str]]:
+        """The top-`k` labels per span (highest logit first). Feeds the hybrid classifier's routing decision
+        (route to the LLM when the top-2 are confusable siblings). Empty input -> empty output."""
+        if not texts:
+            return []
+        id2label = self._model.config.id2label
+        out: list[list[str]] = []
+        for start in range(0, len(texts), batch_size):
+            batch = texts[start : start + batch_size]
+            enc = self._tokenizer(
+                batch, truncation=True, max_length=self._max_length, padding=True, return_tensors="pt"
+            ).to(self._device)
+            logits = self._model(**enc).logits
+            topk = logits.topk(min(k, logits.shape[-1]), dim=-1).indices.tolist()
+            out.extend([str(id2label[int(i)]) for i in row] for row in topk)
+        return out
+
+
+def register_clause_function_classification(registry) -> None:
+    """CAP-REG-2: register `clause_function_classification` (model; fine-tuned LegalBERT classifier)."""
+    from rag_wright.contracts.function import FunctionClassification
+
+    registry.register(
+        "clause_function_classification",
+        contract=FunctionClassification,
+        kind="model",
+        display_name="Clause function classification (LegalBERT)",
+    )

rag_wright/spans/model_capabilities.py

@@ -0,0 +1,107 @@
+"""EP-RT-7 (ADR-0117): the SINGLE binding of model-capability name -> implementation.
+
+This is the ONE place the engine's model capabilities (`clause_function_classification`, `clause_property_classification`)
+are bound to their fleets and dispatched. It lives at the `spans/` layer -- below BOTH the engine API invoker
+(`api/invoke.py`) and the ingestion pipeline (`subgraphs/`) -- so every consumer routes through it by name. No
+capability's fleet is constructed anywhere else: production goes THROUGH the capability layer, never around it
+(see memory route-production-through-capability-layer). Subgraph capabilities stay bound in `api/invoke.py` because
+they are workspace-bound (need the handle's store/models); model capabilities are store-independent, so they bind
+here where both the API and the pipeline can reach them.
+
+Usage/cost is captured by the caller's ambient `api.measure_usage()` scope (EP-API-5); this wraps each dispatch in a
+trace span only."""
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable
+
+from rag_wright.capabilities.invoke import capability_impl
+from rag_wright.models.tracing import traced_step
+
+# --- cached fleets: loaded ONCE per process (heavy) ---
+
+_DIM_REGISTRY: Any = None
+
+
+def _dim_registry() -> Any:
+    """The 29-dim property-classifier fleet (shared Laya group agents + SetFit + abstain heads)."""
+    global _DIM_REGISTRY
+    if _DIM_REGISTRY is None:
+        from rag_wright.spans.dim_classifier import load_dim_registry
+
+        _DIM_REGISTRY = load_dim_registry()
+    return _DIM_REGISTRY
+
+
+_SETFIT_CLAUSE: Any = None
+
+
+def _setfit_clause_classifier() -> Any:
+    """The trained SetFit/LegalBERT clause-function soft-tagger."""
+    global _SETFIT_CLAUSE
+    if _SETFIT_CLAUSE is None:
+        from rag_wright.spans.clause_function_classifier import production_setfit_clause_classifier
+
+        _SETFIT_CLAUSE = production_setfit_clause_classifier()
+    return _SETFIT_CLAUSE
+
+
+# --- the model-capability adapters: name -> (inputs -> output). Store-independent. ---
+
+# The model-capability invoke factories: `(resources, inputs) -> result`, uniform with the subgraph factories
+# (EP-CORE-2). Model capabilities are store-independent, so `resources` is ignored. These are the targets of the
+# manifests' `impl_ref` ("rag_wright.spans.model_capabilities:clause_function_classification", etc.).
+
+def clause_function_classification(resources: Any, inputs: dict) -> Any:  # noqa: ARG001 - store-independent
+    return _setfit_clause_classifier().classify_spans(inputs["chunk_text"], inputs["span_texts"])
+
+
+def clause_property_classification(resources: Any, inputs: dict) -> Any:  # noqa: ARG001 - store-independent
+    from rag_wright.spans.property_extractor import HybridPropertyExtractor
+
+    ext = HybridPropertyExtractor(_dim_registry(), runnable=object())  # classifier lane only; runnable unused
+    return ext.classify_properties(inputs["text"], functions=tuple(inputs.get("functions", ())))
+
+
+def dispatch_model(name: str, inputs: dict) -> Any:
+    """Dispatch a model-capability by name, inside a trace span -- resolving the implementation via the manifest
+    `impl_ref` (the SAME adapter-free path the engine API invoker uses, EP-CORE-2). This is the ONE production path
+    the ingestion pipeline also routes through. Model factories are store-independent, so resources is `None`."""
+    factory = capability_impl(name)  # resolves impl_ref; raises KeyError/NotImplementedError on an unknown/unwired name
+    with traced_step(f"invoke:{name}"):
+        return factory(None, inputs)
+
+
+async def adispatch_model(name: str, inputs: dict, *, sem: asyncio.Semaphore | None = None) -> Any:
+    """Async dispatch: run the (CPU-bound, in-process) model capability off the event loop; `sem` bounds in-flight
+    work across concurrent callers (the ingestion segment stage passes one shared semaphore)."""
+    if sem is None:
+        return await asyncio.to_thread(dispatch_model, name, inputs)
+    async with sem:
+        return await asyncio.to_thread(dispatch_model, name, inputs)
+
+
+# --- thin shims the ingestion pipeline wires its classify stages to (so production routes through the capability) ---
+
+class CapabilityFunctionClassifier:
+    """The ingestion segment stage's clause-function classifier, routed through the `clause_function_classification`
+    capability. Same `(classify_spans / aclassify_spans(sem=))` shape the segment stage expects."""
+
+    def classify_spans(self, chunk_text: str, span_texts: list[str]) -> Any:
+        return dispatch_model("clause_function_classification",
+                              {"chunk_text": chunk_text, "span_texts": span_texts})
+
+    async def aclassify_spans(self, chunk_text: str, span_texts: list[str],
+                              *, sem: asyncio.Semaphore | None = None) -> Any:
+        return await adispatch_model("clause_function_classification",
+                                     {"chunk_text": chunk_text, "span_texts": span_texts}, sem=sem)
+
+
+def capability_property_classifier_fn() -> Callable[[str, tuple], list[dict]]:
+    """The classifier LANE for the ingestion extract stage, routed through the `clause_property_classification`
+    capability. Returns `(text, functions) -> [{dimension, value, confidence}]`; the pipeline composes the residual
+    LLM call + the semantic judge AROUND it (those are not part of this capability)."""
+    def _classify(text: str, functions: tuple = ()) -> list[dict]:
+        return dispatch_model("clause_property_classification", {"text": text, "functions": tuple(functions)})
+
+    return _classify

rag_wright/spans/new_function_labels.py

@@ -0,0 +1,111 @@
+"""T60 (FR-C.3, ADR-0026): source training labels for the 3 extended function classes.
+
+CUAD annotates only its 41 `ClauseCategory` types, so it has no spans for the ACORD families the
+function taxonomy adds (Indemnification, the indirect/consequential damages waiver, the warranty
+disclaimer). We bootstrap them from the SAME CUAD contracts (ACORD stays pure eval -- no leakage): the
+operative spans CUAD leaves unlabeled (NONE) are candidates, narrowed by a cheap high-recall KEYWORD
+pre-filter and then confirmed by an LLM (DeepSeek via the model-profile seam).
+
+The keyword pre-filter is unit-testable and free; the LLM confirm is the same structured-output seam
+pattern the OKF classifier uses (stub-injectable, so the plumbing is tested without a model call).
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional, Protocol, runtime_checkable
+
+from pydantic import BaseModel
+
+from rag_wright.contracts.function import ExtendedFunction
+from rag_wright.models.profiles import ModelRole, model_for
+from rag_wright.models.seam import build_structured
+
+NONE_LABEL = "NONE"  # matches the T56 classifier's off-taxonomy sentinel
+
+# High-recall keyword pre-filter (lowercased substring match). Precision is the LLM's job; this only
+# has to avoid dropping true positives, so it is deliberately broad.
+NEW_FUNCTION_KEYWORDS: dict[str, tuple[str, ...]] = {
+    ExtendedFunction.INDEMNIFICATION.value: (
+        "indemnif", "hold harmless", "harmless from", "defend, indemnify",
+    ),
+    ExtendedFunction.INDIRECT_DAMAGES_WAIVER.value: (
+        "consequential", "indirect damage", "incidental damage", "punitive damage",
+        "special damage", "in no event", "lost profits", "loss of profit",
+    ),
+    ExtendedFunction.WARRANTY_DISCLAIMER.value: (
+        "disclaim", "as is", "as-is", "merchantability", "fitness for a particular purpose",
+        "implied warrant", "no warrant",
+    ),
+}
+
+
+class NewFunctionTag(str, Enum):
+    """The LLM confirm's closed output: one of the 3 extended functions, or NONE (none of them)."""
+
+    INDEMNIFICATION = ExtendedFunction.INDEMNIFICATION.value
+    INDIRECT_DAMAGES_WAIVER = ExtendedFunction.INDIRECT_DAMAGES_WAIVER.value
+    WARRANTY_DISCLAIMER = ExtendedFunction.WARRANTY_DISCLAIMER.value
+    NONE = NONE_LABEL
+
+
+class NewFunctionLabel(BaseModel):
+    """Structured LLM output: which extended-function class the span is (or NONE)."""
+
+    label: NewFunctionTag
+
+
+def keyword_candidates(text: str) -> frozenset[str]:
+    """The extended-function labels a span MIGHT be, by keyword (the pre-filter). Empty => not a
+    candidate for any new class (skip the LLM call entirely)."""
+    low = text.lower()
+    hits = {
+        label
+        for label, keywords in NEW_FUNCTION_KEYWORDS.items()
+        if any(kw in low for kw in keywords)
+    }
+    return frozenset(hits)
+
+
+def label_prompt(text: str, candidates: frozenset[str]) -> str:
+    """The confirm prompt: given the pre-filter's candidate classes, decide the one true class or NONE."""
+    options = "\n- ".join(sorted(candidates))
+    return (
+        "You are labeling a contract clause span for a function classifier. A keyword filter flagged it "
+        "as possibly one of these clause types:\n- " + options + "\n\n"
+        "Decide which ONE it actually is, or NONE if it is none of them (the keyword was incidental). "
+        "Judge by the span's operative meaning, not a passing mention.\n\n"
+        "Span:\n" + text[:2000]
+    )
+
+
+@runtime_checkable
+class NewFunctionConfirmer(Protocol):
+    """Span text + candidate classes -> a confirmed label. The seam a test stubs (no model call)."""
+
+    def __call__(self, text: str, candidates: frozenset[str]) -> NewFunctionTag: ...
+
+
+class SeamNewFunctionConfirmer:
+    """The real confirmer: structured output through the model-profile seam on STRUCTURED_REASONING
+    (DeepSeek V4 Pro by default; quality-sensitive labeling). Retries a bare `None` (a transient miss
+    that `with_structured_output` returns instead of raising), and raises only when it persists."""
+
+    def __init__(self, model_id: Optional[str] = None, *, retries: int = 3) -> None:
+        self._runnable = build_structured(
+            model_id or model_for(ModelRole.STRUCTURED_REASONING), NewFunctionLabel
+        )
+        self._retries = retries
+
+    def __call__(self, text: str, candidates: frozenset[str]) -> NewFunctionTag:
+        prompt = label_prompt(text, candidates)
+        last_error: Exception | None = None
+        for _ in range(self._retries):
+            try:
+                v = self._runnable.invoke(prompt)
+            except Exception as e:  # noqa: BLE001 - transient provider/parse error; retry
+                last_error = e
+                continue
+            if v is not None:
+                return v.label
+        raise last_error or ValueError("no structured output after retries")

rag_wright/spans/page_map.py

@@ -0,0 +1,68 @@
+"""issue 0032 / CU-B5: map a span's canonical-text character range to its source page(s).
+
+Docling records each parsed item's `prov[0].page_no` (+ a best-effort bbox); `content_items` now carries it
+(`ContentItem.page`/`.bbox`). The chunker builds a canonical document text (`_SEP`-join of the chunk texts,
+each a `_SEP`-join of the stripped item texts) that span `doc_start`/`doc_end` offsets index into. This module
+reconstructs a `[char_range) -> page` map over that canonical text by locating each item's text in reading
+order, so a span's pages fall out of the char range it already has -- correctly handling a clause that crosses
+a page boundary (the pages come back as a LIST). Page-only is enough for the honest scanned-PDF fallback
+("page 7 -- exact location unavailable"); bbox is carried best-effort where the parser produced one.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+BBox = tuple[float, float, float, float]
+
+
+@dataclass(frozen=True)
+class PageRange:
+    """One content item's [start, end) character range in the canonical document text, its 1-based source
+    page, and a best-effort bbox (l, t, r, b on that page) when the parser produced one."""
+
+    start: int
+    end: int
+    page: int
+    bbox: Optional[BBox]
+
+
+def build_page_offset_map(items: list[Any], canonical_text: str) -> list[PageRange]:
+    """Reconstruct the page map for `canonical_text` by locating each content item's STRIPPED text in reading
+    order with a forward cursor (the canonical text is the `_SEP`-join of exactly those stripped item texts).
+
+    Robust to the chunker's below-floor merges and bare-heading folds -- they regroup which items land in
+    which chunk but never change the item-text SEQUENCE, so the cursor still finds each in order. The only
+    lossy case is an over-cap item hard-split across chunks (rare): its text is not contiguous, so it is left
+    unlocated (best-effort). Items with no page or empty text contribute nothing."""
+    out: list[PageRange] = []
+    cursor = 0
+    for it in items:
+        page = getattr(it, "page", None)
+        text = (getattr(it, "text", "") or "").strip()
+        if page is None or not text:
+            continue
+        idx = canonical_text.find(text, cursor)
+        if idx == -1:  # a merged/reordered edge: fall back to a global search before giving up
+            idx = canonical_text.find(text)
+            if idx == -1:
+                continue
+        out.append(PageRange(idx, idx + len(text), int(page), getattr(it, "bbox", None)))
+        cursor = idx + len(text)
+    return out
+
+
+def pages_for(
+    page_map: list[PageRange], doc_start: Optional[int], doc_end: Optional[int]
+) -> tuple[list[int], Optional[BBox]]:
+    """The distinct source pages a `[doc_start, doc_end)` span overlaps (sorted ascending -- a LIST, because a
+    clause can cross a page boundary), plus a best-effort bbox: the box of the SINGLE overlapping item when
+    exactly one overlaps and it carries a box, else `None` (a multi-item/multi-page span has no single box).
+    An empty map, missing offsets, or no overlap -> `([], None)`."""
+    if doc_start is None or doc_end is None:
+        return [], None
+    overlapping = [pr for pr in page_map if pr.start < doc_end and pr.end > doc_start]
+    pages = sorted({pr.page for pr in overlapping})
+    bbox = overlapping[0].bbox if len(overlapping) == 1 else None
+    return pages, bbox