Skip to content

Commit 759860d

Browse files
Add bounded cache-only local summarizer backends (#3598)
Add a pinned offline MLX summarizer, explicit extractive backend, guarded runtime admission, shared task contract and privacy-safe error handling. Include focused regression tests and setup documentation. Closes #3234.
1 parent ac7a3db commit 759860d

9 files changed

Lines changed: 683 additions & 30 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
99

1010
### Added
1111

12+
- Added cache-only, registry-pinned MLX summarization with capability, memory,
13+
prompt-digest and source-token leakage guards. The deterministic CPU baseline
14+
now requires explicit `model="extractive"`; missing runtime or artifacts fail
15+
closed instead of silently changing backends (#3234).
1216
- Added offline benchmark suite listing and description, explicit task metadata,
1317
and aggregate report comparison with missing-evidence and regression verdicts
1418
(#286).

‎docs/clinical/summarization.md‎

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22

33
OpenMed's clinical summarization stage is generative-last. The public
44
`openmed.clinical.summarize()` entry point de-identifies a note first, then
5-
passes only the de-identified text to a summarizer backend. A deterministic
6-
extractive fallback is used when no backend is supplied.
5+
passes only the de-identified text to a summarizer backend. The default `mlx`
6+
backend resolves the reviewed, revision-pinned Maple alias through the registry.
7+
It loads cached artifacts only, with outbound Python sockets blocked. It never
8+
downloads a model or silently falls back to extraction.
79

810
```python
911
from openmed.clinical import summarize
1012

11-
result = summarize(note)
13+
result = summarize(note, model="extractive") # explicit, offline CPU baseline
1214
assert result.leakage_check.passed
1315
print(result.summary)
1416
```
@@ -22,7 +24,36 @@ Pipeline code that already performed de-identification may call
2224
`summarize_deidentified()` with its `DeidentificationResult`. Passing a plain
2325
string to that guarded stage raises an ordering error.
2426

25-
The default path is local and deterministic. A trained SLM backend is a
26-
separate task and must be supplied as a local or on-device callable; this
27-
stage never selects a cloud model or sends raw PHI to one. Summaries are
28-
assistive outputs and require qualified clinical review.
27+
Use `model="mlx"` (or `"maple"`) on Apple Silicon after separately installing
28+
`openmed[mlx]` and downloading the pinned artifact. Raw-note input also requires
29+
the cached `OpenMed/OpenMed-PII-SuperClinical-Small-44M-v1-mlx` artifact and its
30+
source model's tokenizer. Download these separately before offline inference;
31+
no conversion, dependency installation or download runs inside `summarize()`.
32+
With the default cache locations, provision once using:
33+
34+
```bash
35+
hf download deepgrove/maple-preview-2bit-mlx --revision 361db5da5e74ff6fcdd852d478e1f266ce11013a
36+
hf download OpenMed/OpenMed-PII-SuperClinical-Small-44M-v1-mlx --cache-dir ~/.cache/openmed
37+
hf download OpenMed/OpenMed-PII-SuperClinical-Small-44M-v1 --cache-dir ~/.cache/openmed --include '*.json' --include '*.txt' --include '*.model'
38+
```
39+
40+
Missing dependencies raise
41+
`MissingOptionalDependencyError`; missing weights, context overflow, insufficient
42+
estimated memory, and malformed responses fail closed with content-free errors.
43+
The adapter runs the shared `summarize` Maple task, template digest, offline
44+
capability probe and conservative memory preflight before model construction.
45+
The default memory estimate is capped at 16 GiB; applications can pass an
46+
`MLXSummarizerBackend(memory_budget_bytes=...)` with an explicit device budget.
47+
This estimate is not a measured peak-memory guarantee.
48+
The bounded generation allowance is 2,048 tokens, including internal reasoning;
49+
it is reserved before model loading. Final summary text is limited to 8 KiB.
50+
51+
`result.metadata` contains only a fixed backend ID and template digest. Do not
52+
log `result.summary` or `result.to_dict()`: the explicit output contains clinical
53+
text. Custom callables are trusted local application code, not sandboxed plugins;
54+
they run with OpenMed's Python socket guard but must not spawn remote tools.
55+
56+
Migration: callers relying on the former implicit sentence-picker must now use
57+
`model="extractive"`. Maple is a general model, not a released clinical SLM or
58+
clinically validated summarizer. Summaries require qualified clinical review;
59+
the source-token leakage check cannot prove that every identifier was detected.

‎openmed/clinical/summarize.py‎

Lines changed: 88 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,12 @@
22
33
Summarization is deliberately the last generative stage in the clinical
44
pipeline. :func:`summarize` accepts a raw note, de-identifies it locally, and
5-
then sends only the de-identified text to a caller-supplied backend or the
6-
deterministic extractive fallback. :func:`summarize_deidentified` exposes the
5+
then sends only the de-identified text to a local backend or explicitly selected
6+
deterministic extraction. :func:`summarize_deidentified` exposes the
77
guarded stage for callers that already own a :class:`DeidentificationResult`.
88
9-
The default backend is intentionally small and deterministic. A trained
10-
on-device SLM, including an MLX implementation, can be supplied through the
11-
``model`` argument without changing the privacy boundary.
9+
The default is a cache-only MLX backend. Missing artifacts or runtime fail
10+
closed; use ``model="extractive"`` for the deterministic CPU baseline.
1211
"""
1312

1413
from __future__ import annotations
@@ -17,7 +16,7 @@
1716
import inspect
1817
import re
1918
from collections.abc import Iterator, Mapping
20-
from dataclasses import dataclass
19+
from dataclasses import dataclass, field
2120
from typing import Any, Protocol
2221

2322
from openmed.core.pii import DeidentificationResult, deidentify
@@ -122,10 +121,11 @@ class SummarizationResult:
122121
``result.leakage_check`` or unpack ``summary, leakage_check``.
123122
"""
124123

125-
summary: str
124+
summary: str = field(repr=False)
126125
leakage_check: LeakageCheck
127126
mode: str = DEFAULT_SUMMARIZATION_MODE
128127
backend: str = "deterministic-extractive"
128+
template_digest: str | None = None
129129

130130
def __post_init__(self) -> None:
131131
if not isinstance(self.summary, str):
@@ -136,13 +136,23 @@ def __post_init__(self) -> None:
136136
raise ValueError("mode must be a non-empty string")
137137
if not isinstance(self.backend, str) or not self.backend:
138138
raise ValueError("backend must be a non-empty string")
139+
if self.template_digest is not None and (
140+
not isinstance(self.template_digest, str)
141+
or re.fullmatch(r"sha256:[0-9a-f]{64}", self.template_digest) is None
142+
):
143+
raise ValueError("invalid template digest")
139144

140145
def __iter__(self) -> Iterator[str | LeakageCheck]:
141146
"""Yield the summary and leakage check for tuple-style consumers."""
142147

143148
yield self.summary
144149
yield self.leakage_check
145150

151+
@property
152+
def metadata(self) -> dict[str, Any]:
153+
"""Return value-free backend and template provenance for audit logging."""
154+
return {"backend_id": self.backend, "template_digest": self.template_digest}
155+
146156
def to_dict(self) -> dict[str, Any]:
147157
"""Return a JSON-compatible result without source PHI metadata."""
148158

@@ -151,6 +161,7 @@ def to_dict(self) -> dict[str, Any]:
151161
"leakage_check": self.leakage_check.to_dict(),
152162
"mode": self.mode,
153163
"backend": self.backend,
164+
"metadata": self.metadata,
154165
}
155166

156167

@@ -171,8 +182,8 @@ def summarize(
171182
:func:`openmed.core.pii.deidentify`.
172183
mode: Summarization mode forwarded to compatible backends. ``"bhc"``
173184
denotes a brief hospital-course style summary.
174-
model: Optional local/on-device backend. It may be callable with the
175-
de-identified text or expose ``summarize(text, *, mode=...)``.
185+
model: Registry alias (default ``mlx``), explicit ``extractive``, or
186+
a caller-owned local callable accepting de-identified text.
176187
177188
Returns:
178189
A summary and a passing :class:`LeakageCheck`.
@@ -191,8 +202,31 @@ def summarize(
191202
return summarize_deidentified(text, mode=normalized_mode, model=model)
192203
if not isinstance(text, str):
193204
raise TypeError("text must be a string or DeidentificationResult")
194-
result = deidentify(text, method="mask")
195-
return summarize_deidentified(result, mode=normalized_mode, model=model)
205+
from openmed.clinical.summarize_backends import (
206+
LocalSummarizerError,
207+
MLXSummarizerBackend,
208+
_require_runtime,
209+
resolve_summarizer_backend,
210+
)
211+
from openmed.core.config import OpenMedConfig
212+
from openmed.core.offline import network_blocked_if_offline
213+
214+
backend = resolve_summarizer_backend(model)
215+
if isinstance(backend, MLXSummarizerBackend):
216+
_require_runtime()
217+
failed = False
218+
try:
219+
with network_blocked_if_offline(local_only=True):
220+
result = deidentify(
221+
text, method="mask", config=OpenMedConfig(local_only=True)
222+
)
223+
except Exception:
224+
failed = True
225+
if failed:
226+
raise LocalSummarizerError(
227+
"local de-identification failed; prepare cached PII artifacts"
228+
)
229+
return summarize_deidentified(result, mode=normalized_mode, model=backend)
196230

197231

198232
def summarize_deidentified(
@@ -230,7 +264,10 @@ def summarize_deidentified(
230264
"the de-identified text still contains a source token"
231265
)
232266

233-
summary = _invoke_backend(model, source.deidentified_text, normalized_mode)
267+
from openmed.clinical.summarize_backends import resolve_summarizer_backend
268+
269+
backend = resolve_summarizer_backend(model)
270+
summary = _invoke_backend(backend, source.deidentified_text, normalized_mode)
234271
leakage_check = _build_leakage_check(source, summary)
235272
if not leakage_check.passed:
236273
raise SummarizationLeakageError(leakage_check)
@@ -239,7 +276,12 @@ def summarize_deidentified(
239276
summary=summary,
240277
leakage_check=leakage_check,
241278
mode=normalized_mode,
242-
backend=_backend_name(model),
279+
backend=_backend_name(backend),
280+
template_digest=(
281+
backend.template_digest
282+
if _backend_name(backend) != "caller-supplied-local"
283+
else None
284+
),
243285
)
244286

245287

@@ -274,6 +316,29 @@ def _require_deidentification_result(value: object) -> DeidentificationResult:
274316

275317

276318
def _invoke_backend(model: object | None, text: str, mode: str) -> str:
319+
from openmed.clinical.summarize_backends import (
320+
LocalSummarizerError,
321+
MLXSummarizerBackend,
322+
)
323+
from openmed.core.capabilities import MissingOptionalDependencyError
324+
from openmed.core.offline import network_blocked_if_offline
325+
326+
failed = False
327+
try:
328+
with network_blocked_if_offline(local_only=True):
329+
return _call_backend(model, text, mode)
330+
except MissingOptionalDependencyError:
331+
if type(model) is MLXSummarizerBackend:
332+
raise
333+
failed = True
334+
except Exception:
335+
failed = True
336+
if failed:
337+
raise LocalSummarizerError("local summarizer execution failed")
338+
raise AssertionError("unreachable summarizer state")
339+
340+
341+
def _call_backend(model: object | None, text: str, mode: str) -> str:
277342
if model is None:
278343
return _extractive_summary(text)
279344

@@ -303,7 +368,7 @@ def _invoke_backend(model: object | None, text: str, mode: str) -> str:
303368
else:
304369
output = callback(text)
305370

306-
if not isinstance(output, str):
371+
if not isinstance(output, str) or len(output.encode("utf-8")) > 8192:
307372
raise TypeError("summarizer backend must return a string")
308373
return output.strip()
309374

@@ -319,11 +384,16 @@ def _extractive_summary(text: str) -> str:
319384

320385

321386
def _backend_name(model: object | None) -> str:
322-
if model is None:
387+
from openmed.clinical.summarize_backends import (
388+
ExtractiveSummarizerBackend,
389+
MLXSummarizerBackend,
390+
)
391+
392+
if type(model) is ExtractiveSummarizerBackend:
323393
return "deterministic-extractive"
324-
if hasattr(model, "summarize"):
325-
return f"{type(model).__name__}.summarize"
326-
return type(model).__name__
394+
if type(model) is MLXSummarizerBackend:
395+
return "local-mlx"
396+
return "caller-supplied-local"
327397

328398

329399
def _source_phi_surfaces(deidentified: Any) -> tuple[str, ...]:

0 commit comments

Comments
 (0)