lorebase-inference 0.2.3__py3-none-any.whl

lorebase_inference/interfaces.py

@@ -0,0 +1,151 @@
+"""Vendor-neutral model interfaces and their result types (design 9.4).
+
+Callers depend on these protocols only; the provider behind them is chosen by configuration.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal, Protocol, runtime_checkable
+
+JsonSchema = Mapping[str, Any]
+"""A JSON Schema document describing the expected structure of a model answer."""
+
+
+@dataclass(frozen=True, slots=True)
+class Message:
+    """One chat message."""
+
+    role: Literal["system", "user", "assistant"]
+    content: str
+
+
+@dataclass(frozen=True, slots=True)
+class ChatOptions:
+    """Per-call options; unset values fall back to the configured defaults."""
+
+    task: str | None = None
+    """Task name used for per-task routing (for example `extraction`)."""
+    temperature: float | None = None
+    max_tokens: int | None = None
+    seed: int | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class Usage:
+    """Token accounting for one call."""
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+
+
+@dataclass(frozen=True, slots=True)
+class ChatResult:
+    """Parsed JSON answer with accounting."""
+
+    data: Any
+    """The answer parsed from JSON; conforms to the requested schema."""
+    usage: Usage
+    model: str
+    """Identifier of the model that produced the answer."""
+
+
+@dataclass(frozen=True, slots=True)
+class Label:
+    """An entity type a span extractor should look for."""
+
+    name: str
+    description: str = ""
+    examples: tuple[str, ...] = field(default_factory=tuple)
+    """Surface forms known for this label; used by dictionary-based extractors."""
+
+
+@dataclass(frozen=True, slots=True)
+class Span:
+    """A labelled text span with character offsets (`end` exclusive) into the input text."""
+
+    label: str
+    text: str
+    start: int
+    end: int
+    score: float = 1.0
+
+    def __post_init__(self) -> None:
+        if not 0 <= self.start <= self.end:
+            raise ValueError(f"invalid span offsets {self.start}..{self.end}")
+
+
+@dataclass(frozen=True, slots=True)
+class RelationLabel:
+    """A relation type a relation extractor should look for between spans."""
+
+    name: str
+    description: str = ""
+    head_labels: tuple[str, ...] = field(default_factory=tuple)
+    """Span labels allowed as head; empty allows any."""
+    tail_labels: tuple[str, ...] = field(default_factory=tuple)
+    """Span labels allowed as tail; empty allows any."""
+
+    def allows(self, head: Span, tail: Span) -> bool:
+        """True when the span labels satisfy the head and tail constraints."""
+        return (not self.head_labels or head.label in self.head_labels) and (
+            not self.tail_labels or tail.label in self.tail_labels
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class SpanRelation:
+    """A candidate relation `label` between two spans of the same text."""
+
+    label: str
+    head: Span
+    tail: Span
+    score: float = 1.0
+
+
+@runtime_checkable
+class ChatModel(Protocol):
+    """Chat completion constrained to a JSON schema."""
+
+    async def complete_json(
+        self,
+        messages: Sequence[Message],
+        schema: JsonSchema,
+        options: ChatOptions | None = None,
+    ) -> ChatResult:
+        """Answer `messages` with JSON that conforms to `schema`."""
+        ...
+
+
+@runtime_checkable
+class EmbeddingModel(Protocol):
+    """Text embeddings of a fixed dimension."""
+
+    name: str
+    """Model name; vectors are stored keyed by it."""
+    dims: int
+
+    async def embed(self, texts: Sequence[str]) -> list[list[float]]:
+        """One vector of `dims` floats per input text, in order."""
+        ...
+
+
+@runtime_checkable
+class SpanExtractor(Protocol):
+    """Tier-1 extraction of labelled spans from text."""
+
+    async def extract(self, text: str, labels: Sequence[Label]) -> list[Span]:
+        """Spans found in `text` for the given labels, ordered by start offset."""
+        ...
+
+
+@runtime_checkable
+class RelationExtractor(Protocol):
+    """Tier-1 extraction of candidate relations between already extracted spans."""
+
+    async def extract(
+        self, text: str, spans: Sequence[Span], labels: Sequence[RelationLabel]
+    ) -> list[SpanRelation]:
+        """Relations between `spans` (offsets into `text`), ordered by head then tail start."""
+        ...

lorebase_inference/local_embeddings.py

@@ -0,0 +1,58 @@
+"""In-process CPU embeddings: lazy model load, batching, dimension check.
+
+The model library is reached through a backend callable supplied by the registry, so this
+module imports nothing heavy. The backend is built on first use (loading or downloading the
+model takes seconds), in a worker thread so that the event loop is never blocked.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Callable, Iterable, Sequence
+
+from lorebase_inference.errors import DimensionMismatchError
+from lorebase_inference.settings import EmbeddingSettings
+
+EmbedBackend = Callable[[Sequence[str]], Iterable[Sequence[float]]]
+"""Embeds one batch of texts; yields one vector per text."""
+
+
+class LocalEmbeddingModel:
+    """`EmbeddingModel` running a local model through a backend."""
+
+    def __init__(
+        self, embeddings: EmbeddingSettings, backend_factory: Callable[[], EmbedBackend]
+    ) -> None:
+        self.name = embeddings.model
+        self.dims = embeddings.dims
+        self._batch_size = embeddings.batch_size
+        self._factory = backend_factory
+        self._backend: EmbedBackend | None = None
+        self._lock = threading.Lock()
+
+    async def embed(self, texts: Sequence[str]) -> list[list[float]]:
+        """One vector of `dims` floats per text, in order."""
+        if not texts:
+            return []
+        return await asyncio.to_thread(self._embed, list(texts))
+
+    def _embed(self, texts: list[str]) -> list[list[float]]:
+        with self._lock:
+            if self._backend is None:
+                self._backend = self._factory()
+            vectors: list[list[float]] = []
+            for start in range(0, len(texts), self._batch_size):
+                batch = texts[start : start + self._batch_size]
+                vectors.extend([float(v) for v in vector] for vector in self._backend(batch))
+        if len(vectors) != len(texts):
+            raise DimensionMismatchError(
+                f"model {self.name!r} returned {len(vectors)} vectors for {len(texts)} texts"
+            )
+        for vector in vectors:
+            if len(vector) != self.dims:
+                raise DimensionMismatchError(
+                    f"model {self.name!r} returned {len(vector)} dimensions, "
+                    f"embeddings.dims is {self.dims}"
+                )
+        return vectors

lorebase_inference/local_relations.py

@@ -0,0 +1,157 @@
+"""Tier-1 relation candidates between spans: a local model, or a deterministic fallback.
+
+`LocalRelationExtractor` runs a zero-shot relation model behind a backend callable supplied by
+the registry; this module tokenizes the text, maps spans to token ranges and maps predictions
+back to spans.
+
+`CooccurrenceRelationExtractor` needs no model and is the documented fallback when the
+relation model cannot be installed (it is an optional extra, and its dependencies are heavy).
+It pairs spans that share a sentence and emits every relation label whose head and tail
+constraints the pair satisfies. It knows nothing about meaning: the score is a proximity value
+`1 / (1 + words between the spans)`, not a probability, and the output is a superset of what a
+model would propose. Downstream validation, evidence checks and review decide. Sentences end at
+`.`, `!` or `?` followed by whitespace, and at blank lines; abbreviations therefore split
+sentences, which only loses candidates.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import re
+import threading
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+from lorebase_inference.interfaces import RelationLabel, Span, SpanRelation
+from lorebase_inference.local_spans import split_windows
+from lorebase_inference.settings import CooccurrenceSettings, ExtractorSettings
+
+RelationBackend = Callable[[list[str], list[list[Any]], list[str], float], list[Mapping[str, Any]]]
+"""`(tokens, entities, label names, threshold)` to relations.
+
+Entities are `[first token, last token (inclusive), span label, text]`. Each relation holds
+`head_pos` and `tail_pos` (`[first token, end token (exclusive)]`), `label` and `score`.
+"""
+
+_TOKEN = re.compile(r"\w+(?:['\u2019-]\w+)*|[^\w\s]")
+_SENTENCE_BREAK = re.compile(r"(?<=[.!?])\s+|\n\s*\n")
+
+
+def sentence_bounds(text: str) -> list[tuple[int, int]]:
+    """`(start, end)` offsets of the sentences of `text`, whitespace excluded."""
+    bounds: list[tuple[int, int]] = []
+    start = 0
+    for match in [*_SENTENCE_BREAK.finditer(text), None]:
+        end = match.start() if match else len(text)
+        if text[start:end].strip():
+            bounds.append((start, end))
+        start = match.end() if match else len(text)
+    return bounds
+
+
+class CooccurrenceRelationExtractor:
+    """`RelationExtractor` pairing spans of the same sentence, deterministically."""
+
+    def __init__(self, settings: CooccurrenceSettings | None = None) -> None:
+        self._limit = (settings or CooccurrenceSettings()).max_pairs_per_sentence
+
+    async def extract(
+        self, text: str, spans: Sequence[Span], labels: Sequence[RelationLabel]
+    ) -> list[SpanRelation]:
+        """Candidate relations, ordered by head start, tail start and label."""
+        relations: list[SpanRelation] = []
+        ordered = sorted(spans, key=lambda s: (s.start, s.end, s.label))
+        for start, end in sentence_bounds(text):
+            inside = [s for s in ordered if s.start >= start and s.end <= end]
+            pairs = [
+                (a, b) for i, a in enumerate(inside) for b in inside[i + 1 :] if b.start >= a.end
+            ]
+            pairs.sort(key=lambda p: (_gap_words(text, *p), p[0].start, p[1].start))
+            for head, tail in pairs[: self._limit]:
+                score = 1 / (1 + _gap_words(text, head, tail))
+                for label in labels:
+                    if label.allows(head, tail):
+                        relations.append(SpanRelation(label.name, head, tail, score))
+                    elif label.allows(tail, head):
+                        relations.append(SpanRelation(label.name, tail, head, score))
+        relations.sort(key=lambda r: (r.head.start, r.tail.start, r.label))
+        return relations
+
+
+def _gap_words(text: str, head: Span, tail: Span) -> int:
+    return len(_TOKEN.findall(text[head.end : tail.start]))
+
+
+class LocalRelationExtractor:
+    """`RelationExtractor` running a local zero-shot relation model through a backend."""
+
+    def __init__(
+        self, settings: ExtractorSettings, backend_factory: Callable[[], RelationBackend]
+    ) -> None:
+        self._settings = settings
+        self._factory = backend_factory
+        self._backend: RelationBackend | None = None
+        self._lock = threading.Lock()
+
+    async def extract(
+        self, text: str, spans: Sequence[Span], labels: Sequence[RelationLabel]
+    ) -> list[SpanRelation]:
+        """Candidate relations between `spans`, ordered by head start, tail start and label."""
+        if not labels or len(spans) < 2:
+            return []
+        return await asyncio.to_thread(self._extract, text, list(spans), list(labels))
+
+    def _extract(
+        self, text: str, spans: list[Span], labels: list[RelationLabel]
+    ) -> list[SpanRelation]:
+        spans = sorted(spans, key=lambda s: (s.start, s.end, s.label))
+        by_name = {label.name.casefold(): label for label in labels}
+        names = [label.name for label in labels]
+        found: dict[tuple[int, int, str], SpanRelation] = {}
+        with self._lock:
+            if self._backend is None:
+                self._backend = self._factory()
+            backend = self._backend
+            for offset, window in split_windows(text, self._settings.window_chars):
+                local = [s for s in spans if s.start >= offset and s.end <= offset + len(window)]
+                if len(local) < 2:
+                    continue
+                tokens = [(m.group(), m.start(), m.end()) for m in _TOKEN.finditer(window)]
+                placed = _place(tokens, local, offset)
+                if len(placed) < 2:
+                    continue
+                entities = [[first, last, span.label, span.text] for span, first, last in placed]
+                lookup = {(first, last): span for span, first, last in placed}
+                raw = backend([t[0] for t in tokens], entities, names, self._settings.threshold)
+                for item in raw:
+                    label = by_name.get(str(item.get("label", "")).casefold())
+                    head = _find(lookup, item.get("head_pos"))
+                    tail = _find(lookup, item.get("tail_pos"))
+                    if label is None or head is None or tail is None or head == tail:
+                        continue
+                    if not label.allows(head, tail):
+                        continue
+                    key = (head.start, tail.start, label.name)
+                    score = float(item.get("score", 1.0))
+                    if key not in found or found[key].score < score:
+                        found[key] = SpanRelation(label.name, head, tail, score)
+        return sorted(found.values(), key=lambda r: (r.head.start, r.tail.start, r.label))
+
+
+def _place(
+    tokens: list[tuple[str, int, int]], spans: list[Span], offset: int
+) -> list[tuple[Span, int, int]]:
+    """Spans of the window with the inclusive `(first, last)` token indices covering them."""
+    placed: list[tuple[Span, int, int]] = []
+    for span in spans:
+        start, end = span.start - offset, span.end - offset
+        covered = [i for i, (_t, s, e) in enumerate(tokens) if e > start and s < end]
+        if covered:
+            placed.append((span, covered[0], covered[-1]))
+    return placed
+
+
+def _find(lookup: dict[tuple[int, int], Span], position: Any) -> Span | None:
+    if isinstance(position, Sequence) and len(position) == 2:
+        return lookup.get((int(position[0]), int(position[1]) - 1))
+    return None

lorebase_inference/local_spans.py

@@ -0,0 +1,93 @@
+"""Tier-1 zero-shot span extraction with a local model, behind a backend callable.
+
+The registry supplies the backend (the only code that touches the model library); this module
+handles what is independent of it: windowing of long texts, mapping labels, translating
+offsets back to the input text, and running the model in a worker thread.
+
+The models read a few hundred tokens at a time. Texts longer than `window_chars` are cut at
+sentence ends or whitespace into consecutive windows, so an entity that straddles a cut is
+missed; callers are expected to pass chunks (design 9.2, step 3) and the cut is a safety net.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import re
+import threading
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from typing import Any
+
+from lorebase_inference.interfaces import Label, Span
+from lorebase_inference.settings import ExtractorSettings
+
+SpanBackend = Callable[[str, list[str], float], list[Mapping[str, Any]]]
+"""`(text, label names, threshold)` to entities with `start`, `end`, `label` and `score`."""
+
+_SENTENCE_END = re.compile(r"[.!?]\s|\n")
+
+
+def split_windows(text: str, size: int) -> Iterator[tuple[int, str]]:
+    """Consecutive `(offset, window)` pieces of at most `size` characters covering `text`.
+
+    A cut prefers the last sentence end in the second half of the window, then the last
+    whitespace, then falls back to a hard cut. Windows that are only whitespace are skipped.
+    """
+    position, length = 0, len(text)
+    while position < length:
+        end = min(length, position + size)
+        cut = end - position
+        if end < length:
+            piece = text[position:end]
+            ends = [m.end() for m in _SENTENCE_END.finditer(piece) if m.end() > size // 2]
+            if ends:
+                cut = ends[-1]
+            else:
+                space = max(piece.rfind(" "), piece.rfind("\n"), piece.rfind("\t"))
+                cut = space + 1 if space > 0 else cut
+        if text[position : position + cut].strip():
+            yield position, text[position : position + cut]
+        position += cut
+
+
+class LocalSpanExtractor:
+    """`SpanExtractor` running a local zero-shot model through a backend."""
+
+    def __init__(
+        self, settings: ExtractorSettings, backend_factory: Callable[[], SpanBackend]
+    ) -> None:
+        self._settings = settings
+        self._factory = backend_factory
+        self._backend: SpanBackend | None = None
+        self._lock = threading.Lock()
+
+    async def extract(self, text: str, labels: Sequence[Label]) -> list[Span]:
+        """Spans of the given labels in `text`, ordered by start offset."""
+        if not labels or not text.strip():
+            return []
+        return await asyncio.to_thread(self._extract, text, list(labels))
+
+    def _extract(self, text: str, labels: list[Label]) -> list[Span]:
+        names = list(dict.fromkeys(label.name for label in labels))
+        by_folded = {name.casefold(): name for name in names}
+        spans: list[Span] = []
+        with self._lock:
+            if self._backend is None:
+                self._backend = self._factory()
+            backend = self._backend
+            for offset, window in split_windows(text, self._settings.window_chars):
+                for entity in backend(window, names, self._settings.threshold):
+                    name = by_folded.get(str(entity.get("label", "")).casefold())
+                    start, end = int(entity["start"]) + offset, int(entity["end"]) + offset
+                    if name is None or not 0 <= start < end <= len(text):
+                        continue
+                    spans.append(
+                        Span(
+                            label=name,
+                            text=text[start:end],
+                            start=start,
+                            end=end,
+                            score=float(entity.get("score", 1.0)),
+                        )
+                    )
+        spans.sort(key=lambda s: (s.start, s.end, s.label))
+        return spans