diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..517ab4a
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,63 @@
+name: CI
+
+on:
+ push:
+ pull_request:
+
+permissions:
+ contents: read
+
+jobs:
+ test:
+ name: Test (Python ${{ matrix.python-version }})
+ runs-on: ubuntu-latest
+ strategy:
+ fail-fast: false
+ matrix:
+ python-version: ["3.10", "3.11", "3.12"]
+
+ steps:
+ - uses: actions/checkout@v4
+ - uses: actions/setup-python@v5
+ with:
+ python-version: ${{ matrix.python-version }}
+ cache: pip
+ - name: Install package and test dependencies
+ run: python -m pip install --upgrade pip && python -m pip install -e ".[test]"
+ - name: Run unit tests
+ run: python -m pytest
+
+ package:
+ name: Validate distributions
+ runs-on: ubuntu-latest
+
+ steps:
+ - uses: actions/checkout@v4
+ - uses: actions/setup-python@v5
+ with:
+ python-version: "3.12"
+ cache: pip
+ - name: Install build tools
+ run: python -m pip install --upgrade pip build twine
+ - name: Build source and wheel distributions
+ run: python -m build
+ - name: Inspect wheel contents
+ shell: python
+ run: |
+ import glob
+ import zipfile
+
+ wheel, = glob.glob("dist/*.whl")
+ with zipfile.ZipFile(wheel) as archive:
+ names = set(archive.namelist())
+ required = {
+ "clear/__init__.py",
+ "clear/hub.py",
+ "clear/concepts.py",
+ "clear/bpe_simple_vocab_16e6.txt.gz",
+ }
+ missing = required.difference(names)
+ if missing:
+ raise SystemExit(f"wheel is missing required files: {sorted(missing)}")
+ - name: Validate distribution metadata
+ run: python -m twine check dist/*
diff --git a/CITATION.cff b/CITATION.cff
new file mode 100644
index 0000000..0c08815
--- /dev/null
+++ b/CITATION.cff
@@ -0,0 +1,73 @@
+cff-version: 1.2.0
+message: "If you use CLEAR, please cite the Nature Biomedical Engineering article."
+title: "CLEAR: Concept-Level Embeddings for Auditable Radiology"
+type: software
+authors:
+ - family-names: Han
+ given-names: Tianyu
+ orcid: "https://orcid.org/0000-0002-8636-6462"
+ - family-names: Wu
+ given-names: Riga
+ - family-names: Tian
+ given-names: Yu
+ - family-names: Khader
+ given-names: Firas
+ - family-names: Adams
+ given-names: Lisa C.
+ orcid: "https://orcid.org/0000-0001-5836-4542"
+ - family-names: Bressem
+ given-names: Keno K.
+ orcid: "https://orcid.org/0000-0001-9249-8624"
+ - family-names: Davatzikos
+ given-names: Christos
+ orcid: "https://orcid.org/0000-0002-1025-8561"
+ - family-names: Kather
+ given-names: Jakob Nikolas
+ - family-names: Shen
+ given-names: Li
+ orcid: "https://orcid.org/0000-0002-5443-0503"
+ - family-names: Mankoff
+ given-names: David A.
+ - family-names: Barbosa
+ given-names: Eduardo Mortani
+ name-suffix: Jr.
+ - family-names: Truhn
+ given-names: Daniel
+ orcid: "https://orcid.org/0000-0002-9605-0728"
+repository-code: "https://github.com/peterhan91/CLEAR"
+url: "https://github.com/peterhan91/CLEAR"
+license: Apache-2.0
+version: 0.1.0
+preferred-citation:
+ type: article
+ authors:
+ - family-names: Han
+ given-names: Tianyu
+ - family-names: Wu
+ given-names: Riga
+ - family-names: Tian
+ given-names: Yu
+ - family-names: Khader
+ given-names: Firas
+ - family-names: Adams
+ given-names: Lisa C.
+ - family-names: Bressem
+ given-names: Keno K.
+ - family-names: Davatzikos
+ given-names: Christos
+ - family-names: Kather
+ given-names: Jakob Nikolas
+ - family-names: Shen
+ given-names: Li
+ - family-names: Mankoff
+ given-names: David A.
+ - family-names: Barbosa
+ given-names: Eduardo Mortani
+ name-suffix: Jr.
+ - family-names: Truhn
+ given-names: Daniel
+ title: "CLEAR: an auditable foundation model for radiology grounded in clinical concepts"
+ journal: Nature Biomedical Engineering
+ year: 2026
+ doi: "10.1038/s41551-026-01741-4"
+ url: "https://doi.org/10.1038/s41551-026-01741-4"
diff --git a/README.md b/README.md
index 10aeda7..942afbe 100644
--- a/README.md
+++ b/README.md
@@ -1,112 +1,204 @@
# CLEAR
-Concept-Level Embeddings for Auditable Radiology (CLEAR) packages the tooling behind the paper
-*CLEAR: An Auditable Foundation Model for Radiology Grounded in Clinical Concepts*. The code base
-follows a standard Python package layout so it is straightforward to install, extend, and import
-from other projects.
+**Concept-Level Embeddings for Auditable Radiology**
+
+> A medical AI model can be right for the wrong reason. CLEAR makes the reason
+> inspectable by grounding chest X-ray representations in radiological concepts.
+
+**Published in _Nature Biomedical Engineering_ on 22 July 2026.**
+
+[Open-access paper](https://www.nature.com/articles/s41551-026-01741-4) |
+[DOI](https://doi.org/10.1038/s41551-026-01741-4) |
+[Model and concept bank](https://huggingface.co/peterhan91/CLEAR) |
+[Code](https://github.com/peterhan91/CLEAR) |
+[Citation](#citation)
+
+CLEAR maps a chest X-ray to similarities with **368,294 report-derived
+radiological observations**, then projects those concept scores through a large
+language model (LLM) embedding space. This makes a prediction decomposable into
+weighted contributions from individual observations and supports concept-level
+zero-shot classification, model auditing, and concept bottleneck models.
+
+The published model was pretrained on **873,342 image-report pairs from 239,391
+patients** and externally evaluated on four physician-annotated datasets from
+the United States, Europe, and Asia. In the paper's enlarged-cardiomediastinum
+case study, concept auditing exposed reliance on atelectasis-related concepts;
+retaining clinically relevant mediastinal concepts and refitting the lightweight
+classifier improved AUROC from **0.727 to 0.784** without retraining the image
+encoder.
+
+```mermaid
+flowchart LR
+ A[Chest X-ray] --> B[CLEAR image encoder]
+ B --> C[Similarity to 368,294
radiological observations]
+ C --> D[LLM semantic projection]
+ D --> E[Concept-based image embedding]
+ C --> F[Concept-level auditing]
+ E --> G[Zero-shot and supervised tasks]
+ E --> H[Concept bottleneck models]
+```
-## Features
+## What is in this repository
-- Thin wrapper around CLIP checkpoints with reliable download and tokenisation utilities.
-- Concept-bank utilities for loading the radiological observation space (CLIP and LLM embeddings).
-- A streaming zero-shot pipeline that projects images into the LLM concept space without requiring
- the full concept matrix in accelerator memory.
-- Metric helpers for threshold search, F1, Matthews correlation coefficient analysis, and
- bootstrapped confidence intervals.
-- Example scripts for benchmarking, model auditing, and concept bottleneck experiments.
+- The CLEAR DINOv2 ViT-B/14 image encoder and paired radiology text encoder.
+- A Hugging Face-backed loader with the paper's 448 x 448 preprocessing.
+- Concept-bank and streaming projection utilities for the full observation
+ space without placing every concept embedding on accelerator memory at once.
+- Research scripts for benchmarking, zero-shot evaluation, model auditing,
+ linear probing, and concept bottleneck experiments.
+- Metric utilities for threshold search, F1, Matthews correlation coefficient,
+ and bootstrap confidence intervals.
## Installation
-1. Ensure you are running Python 3.9 or newer.
-2. (Optional) Create and activate a virtual environment.
-3. Install the package in editable mode:
+CLEAR requires Python 3.10 or newer. Clone the repository and install it in an
+environment with a PyTorch build appropriate for your hardware:
- ```bash
- pip install -e .[visualization]
- ```
+```bash
+git clone https://github.com/peterhan91/CLEAR.git
+cd CLEAR
+pip install -e '.[visualization]'
+```
- The optional `visualization` extra pulls in matplotlib for plotting utilities.
+The `visualization` extra installs Matplotlib for the plotting utilities. Model
+weights are downloaded from Hugging Face on first use and cached locally. The
+loader also obtains the DINOv2 architecture code through `torch.hub`; the CLEAR
+checkpoint supplies all model weights.
-## Quick Start
+## Quick start: vision-language backbone
-The example below loads the CLEAR vision-language backbone and computes
-similarity scores between a chest X-ray and radiological text prompts.
-Model weights are available at
-[huggingface.co/peterhan91/CLEAR](https://huggingface.co/peterhan91/CLEAR).
+`load_pretrained()` downloads the released checkpoint and returns both the model
+and the preprocessing transform used for inference:
```python
import torch
from PIL import Image
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, InterpolationMode
-
-from clear.model import CLIP
-from clear import tokenize
-from examples.train import load_clip
-
-# 1. Load the CLEAR checkpoint (DINOv2 ViT-B/14 + text encoder).
-model = load_clip(
- model_path="checkpoints/dinov2-multi-v1.0_vitb/best_model.pt",
- use_dinov2=True,
- dinov2_model_name="dinov2_vitb14",
-)
-model.eval()
-
-# 2. Preprocess a chest X-ray.
-preprocess = Compose([
- Resize(448, interpolation=InterpolationMode.BICUBIC),
- CenterCrop(448),
- ToTensor(),
- Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-])
-device = next(model.parameters()).device
-image = preprocess(Image.open("path/to/cxr.jpg").convert("RGB")).unsqueeze(0).to(device)
-
-# 3. Encode image and text.
-with torch.no_grad():
- image_features = model.encode_image(image)
- text_tokens = tokenize(["pleural effusion", "no pleural effusion"]).to(device)
- text_features = model.encode_text(text_tokens)
-
- # 4. Cosine similarity → softmax probability.
- image_features = image_features / image_features.norm(dim=-1, keepdim=True)
- text_features = text_features / text_features.norm(dim=-1, keepdim=True)
- logits = (image_features @ text_features.T).softmax(dim=-1)
-
-print(logits) # [[prob_positive, prob_negative]]
+
+from clear import load_pretrained, tokenize
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = load_pretrained(device=device)
+
+image = preprocess(
+ Image.open("path/to/cxr.jpg").convert("RGB")
+).unsqueeze(0).to(device)
+text = tokenize(["pleural effusion", "no pleural effusion"]).to(device)
+
+with torch.inference_mode():
+ image_features = torch.nn.functional.normalize(model.encode_image(image), dim=-1)
+ text_features = torch.nn.functional.normalize(model.encode_text(text), dim=-1)
+ pair_scores = (image_features @ text_features.T).softmax(dim=-1)
+
+print(pair_scores) # [[positive_score, negative_score]]
```
-The full CLEAR pipeline projects these concept similarity scores through
-LLM-derived semantic embeddings (SFR-Embedding-Mistral) for auditable
-zero-shot classification. See `examples/` for benchmarking, concept
-bottleneck models, and model auditing scripts.
+This is a **backbone demonstration**: it compares one image with two text prompts
+in the shared 768-dimensional image-text space. It does not reproduce the full
+CLEAR concept projection, auditing experiments, or intervention reported in the
+paper.
-## Repository Layout
+The same path is available as a command-line example. Adding `--top-concepts`
+also retrieves the highest-scoring observations from the complete concept bank;
+this option downloads the approximately 1.1 GB text-encoder embedding matrix.
+```bash
+python examples/quickstart.py path/to/cxr.jpg --top-concepts 10
```
-src/clear/ Core library modules (CLIP loader, concept bank, zero-shot pipeline, metrics).
-examples/ Research scripts and notebooks rewritten to import the packaged modules.
-scripts/ Standalone data preparation and legacy utilities.
-pyproject.toml Packaging metadata for building or installing CLEAR.
-README.md This document.
-LICENSE Apache-2.0 licence text.
+
+## Full concept pipeline
+
+The full published framework follows four stages:
+
+1. Encode a chest X-ray with the released image encoder.
+2. Compute similarities with all 368,294 CLEAR text-encoded observations.
+3. Multiply the concept-similarity vector by the aligned LLM concept embeddings
+ to obtain a concept-based image embedding.
+4. Use that representation for positive-negative zero-shot prompts, supervised
+ probes, auditing analyses, or concept bottleneck models.
+
+The released assets are hosted together on
+[Hugging Face](https://huggingface.co/peterhan91/CLEAR/tree/main):
+
+| File | Contents | Role |
+| --- | --- | --- |
+| `best_model.pt` | CLEAR DINOv2 ViT-B/14 image encoder and paired text encoder | Produces 768-dimensional image and text features |
+| `mimic_concepts.csv` | Ordered vocabulary of 368,294 radiological observations | Defines the rows shared by both concept-embedding files |
+| `concept_embeddings_368294.pt` | 368,294 x 768 CLEAR text-encoder embeddings | Computes image-observation similarity scores |
+| `concepts_embeddings_sfr_mistral.pickle` | 368,294 x 4,096 SFR-Embedding-Mistral embeddings | Projects similarity scores into the paper's LLM semantic space |
+
+`ConceptBank.from_files()` supports the released tensor formats. Because Python
+pickle can execute code while loading, the SFR asset requires the explicit
+`allow_unsafe_pickle=True` opt-in; use that option only with the official,
+trusted artifact.
+
+See [`examples/concepts`](examples/concepts) for concept projection, zero-shot,
+linear-probing, and concept-bottleneck research scripts, and
+[`examples/model_auditing`](examples/model_auditing) for the auditing analyses.
+These scripts expose the study workflow but are not a single-command reproduction
+of every paper result: they require the released assets, access to the relevant
+evaluation datasets, and experiment-specific paths and configuration.
+
+## Reproducibility notes
+
+- Use the released preprocessing returned by `load_pretrained()`; the published
+ model was evaluated with 448 x 448 inputs.
+- For a fixed model snapshot, pass a specific Hugging Face commit hash through
+ the loader's `revision` argument rather than relying on the repository default.
+- Keep the concept vocabulary and both embedding matrices in exactly the same row
+ order. Misaligned rows invalidate the concept projection.
+- The full concept artifacts are large. Plan for sufficient disk and host memory;
+ the library streams concept chunks to the selected accelerator.
+- Evaluation datasets have their own access requirements and licences. This
+ repository does not redistribute the chest X-rays used in the study.
+- Exact cohort construction, statistical procedures, and source data are
+ documented in the
+ [published article](https://www.nature.com/articles/s41551-026-01741-4).
+
+## Intended use and limitations
+
+CLEAR is released for research on chest X-ray representation learning,
+concept-level analysis, zero-shot evaluation, model auditing, and interpretable
+model development. It is **not a medical device and is not validated for clinical
+decision-making**. Predictions and concept attributions require independent
+validation for each population, acquisition setting, and task. They should not
+be used to diagnose, treat, or triage patients.
+
+The Nature article states that the model weights and full concept bank are
+available for academic research. Review the terms accompanying the
+[Hugging Face assets](https://huggingface.co/peterhan91/CLEAR) and the licences of
+all upstream models and datasets before reuse.
+
+## Repository layout
+
+```text
+src/clear/ Core library: checkpoint loader, concept bank, pipeline, and metrics
+examples/ Research scripts for benchmarks, concepts, and model auditing
+scripts/ Data-preparation and legacy utilities
+pyproject.toml Package and dependency metadata
+CITATION.cff Machine-readable citation metadata
+LICENSE Apache-2.0 licence text
```
## Citation
-If you find this work useful, please cite our paper:
+If you use CLEAR, please cite the published article:
```bibtex
@article{han2026clear,
- title={CLEAR: An Auditable Foundation Model for Radiology Grounded in Clinical Concepts},
- author={Han, Tianyu and Wu, Riga and Tian, Yu and Khader, Firas and Adams, Lisa and Bressem, Keno and Davatzikos, Christos and Kather, Jakob Nikolas and Shen, Li and Mankoff, David and others},
- journal={medRxiv},
- pages={2026--01},
- year={2026},
- publisher={Cold Spring Harbor Laboratory Press}
+ author = {Han, Tianyu and Wu, Riga and Tian, Yu and Khader, Firas and Adams, Lisa C. and Bressem, Keno K. and Davatzikos, Christos and Kather, Jakob Nikolas and Shen, Li and Mankoff, David A. and Barbosa, Jr., Eduardo Mortani and Truhn, Daniel},
+ title = {{CLEAR}: an auditable foundation model for radiology grounded in clinical concepts},
+ journal = {Nature Biomedical Engineering},
+ year = {2026},
+ doi = {10.1038/s41551-026-01741-4},
+ url = {https://doi.org/10.1038/s41551-026-01741-4}
}
```
+GitHub also reads the machine-readable [`CITATION.cff`](CITATION.cff) file and
+offers a **Cite this repository** action.
+
## License
-All code in this repository is released under the Apache-2.0 License. Portions derived from the OpenAI
-CLIP repository retain the original copyright notice.
+The code in this repository is released under the Apache-2.0 License. Portions
+derived from OpenAI CLIP retain their original copyright notice. Model weights,
+concept assets, upstream models, and datasets may be governed by separate terms.
diff --git a/examples/__init__.py b/examples/__init__.py
new file mode 100644
index 0000000..cd2261c
--- /dev/null
+++ b/examples/__init__.py
@@ -0,0 +1 @@
+"""Runnable examples for the CLEAR research codebase."""
diff --git a/examples/quickstart.py b/examples/quickstart.py
new file mode 100644
index 0000000..0407875
--- /dev/null
+++ b/examples/quickstart.py
@@ -0,0 +1,166 @@
+"""Run a single-image CLEAR backbone and concept-retrieval demonstration.
+
+The default path downloads only the released model checkpoint and compares an
+image with positive/negative text prompts. Passing ``--top-concepts`` also
+downloads the 768-dimensional concept matrix and reports the highest-scoring
+radiological observations. The latter artifact is approximately 1.1 GB.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+from pathlib import Path
+from typing import Sequence
+import warnings
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from clear import load_pretrained, tokenize
+
+
+DEFAULT_REPO_ID = "peterhan91/CLEAR"
+
+
+def _load_concepts(path: str | Path) -> list[str]:
+ with Path(path).open("r", encoding="utf-8", newline="") as handle:
+ reader = csv.DictReader(handle)
+ if not reader.fieldnames:
+ raise ValueError(f"No header row found in {path}")
+ field = "concept" if "concept" in reader.fieldnames else reader.fieldnames[0]
+ return [str(row[field]) for row in reader]
+
+
+def _load_tensor(path: str | Path) -> torch.Tensor:
+ try:
+ payload = torch.load(path, map_location="cpu", weights_only=True)
+ except TypeError: # Compatibility with PyTorch versions before weights_only.
+ warnings.warn(
+ "This PyTorch version cannot use weights-only loading; continue only "
+ f"if you trust {path}.",
+ RuntimeWarning,
+ stacklevel=2,
+ )
+ payload = torch.load(path, map_location="cpu")
+ if not torch.is_tensor(payload) or payload.ndim != 2:
+ raise ValueError(f"Expected one two-dimensional tensor in {path}")
+ return payload.float()
+
+
+@torch.inference_mode()
+def _rank_concepts(
+ image_features: torch.Tensor,
+ concept_embeddings: torch.Tensor,
+ concepts: Sequence[str],
+ *,
+ top_k: int,
+ chunk_size: int,
+) -> list[tuple[str, float]]:
+ if concept_embeddings.shape[0] != len(concepts):
+ raise ValueError(
+ "Concept rows and embedding rows are misaligned: "
+ f"{len(concepts)} concepts versus {concept_embeddings.shape[0]} embeddings"
+ )
+
+ device = image_features.device
+ best_scores = torch.empty(0, device=device)
+ best_indices = torch.empty(0, dtype=torch.long, device=device)
+
+ for start in range(0, concept_embeddings.shape[0], chunk_size):
+ stop = min(start + chunk_size, concept_embeddings.shape[0])
+ chunk = F.normalize(concept_embeddings[start:stop].to(device), dim=-1)
+ scores = (image_features @ chunk.T).squeeze(0)
+ local_k = min(top_k, scores.numel())
+ values, indices = scores.topk(local_k)
+ indices = indices + start
+
+ best_scores = torch.cat((best_scores, values))
+ best_indices = torch.cat((best_indices, indices))
+ keep = min(top_k, best_scores.numel())
+ best_scores, order = best_scores.topk(keep)
+ best_indices = best_indices[order]
+
+ return [
+ (concepts[index], float(score))
+ for index, score in zip(best_indices.cpu().tolist(), best_scores.cpu().tolist())
+ ]
+
+
+def _parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("image", type=Path, help="Path to a chest X-ray image")
+ parser.add_argument(
+ "--prompts",
+ nargs=2,
+ default=("pleural effusion", "no pleural effusion"),
+ metavar=("POSITIVE", "NEGATIVE"),
+ help="Positive and negative text prompts for the backbone demonstration",
+ )
+ parser.add_argument(
+ "--top-concepts",
+ type=int,
+ default=0,
+ metavar="N",
+ help="Also retrieve the N highest-scoring CLEAR observations",
+ )
+ parser.add_argument("--chunk-size", type=int, default=8192)
+ parser.add_argument("--repo-id", default=DEFAULT_REPO_ID)
+ parser.add_argument(
+ "--device",
+ default="cuda" if torch.cuda.is_available() else "cpu",
+ help="PyTorch device (default: cuda when available, otherwise cpu)",
+ )
+ args = parser.parse_args()
+ if args.top_concepts < 0:
+ parser.error("--top-concepts must be non-negative")
+ if args.chunk_size <= 0:
+ parser.error("--chunk-size must be positive")
+ return args
+
+
+def main() -> None:
+ args = _parse_args()
+ model, preprocess = load_pretrained(repo_id=args.repo_id, device=args.device)
+ with Image.open(args.image) as source_image:
+ image = preprocess(source_image).unsqueeze(0).to(args.device)
+
+ with torch.inference_mode():
+ image_features = F.normalize(model.encode_image(image), dim=-1)
+ prompt_tokens = tokenize(list(args.prompts)).to(args.device)
+ prompt_features = F.normalize(model.encode_text(prompt_tokens), dim=-1)
+ pair_scores = (image_features @ prompt_features.T).softmax(dim=-1).squeeze(0)
+
+ print("Backbone paired-prompt scores (not calibrated clinical probabilities):")
+ for prompt, score in zip(args.prompts, pair_scores.cpu().tolist()):
+ print(f" {prompt}: {score:.4f}")
+
+ if not args.top_concepts:
+ return
+
+ concepts_path = hf_hub_download(
+ repo_id=args.repo_id,
+ filename="mimic_concepts.csv",
+ )
+ embeddings_path = hf_hub_download(
+ repo_id=args.repo_id,
+ filename="concept_embeddings_368294.pt",
+ )
+ concepts = _load_concepts(concepts_path)
+ concept_embeddings = _load_tensor(embeddings_path)
+ ranked = _rank_concepts(
+ image_features,
+ concept_embeddings,
+ concepts,
+ top_k=args.top_concepts,
+ chunk_size=args.chunk_size,
+ )
+
+ print(f"\nTop {args.top_concepts} CLEAR radiological observations:")
+ for rank, (concept, score) in enumerate(ranked, start=1):
+ print(f" {rank:>2}. {concept} ({score:.4f})")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/pyproject.toml b/pyproject.toml
index 32c5744..1e350ec 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,21 +1,22 @@
[build-system]
-requires = ["setuptools>=65", "wheel"]
+requires = ["setuptools>=77", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "clear-med"
version = "0.1.0"
-description = "CLIP-based Evaluation and Analysis for Radiology"
+description = "Concept-Level Embeddings for Auditable Radiology"
readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
authors = [
{ name = "CLEAR Contributors" }
]
-licenses = [{ text = "MIT" }]
+license = "Apache-2.0"
+license-files = ["LICENSE"]
keywords = ["medical-imaging", "clip", "radiology", "zero-shot"]
dependencies = [
- "torch>=1.12",
- "torchvision>=0.13",
+ "torch>=2.0",
+ "torchvision>=0.15",
"numpy>=1.21",
"pandas>=1.3",
"pillow>=9.0",
@@ -24,11 +25,21 @@ dependencies = [
"scipy>=1.7",
"h5py>=3.6",
"ftfy>=6.1",
- "regex>=2021.8.3"
+ "regex>=2021.8.3",
+ "huggingface-hub>=0.20"
]
[project.optional-dependencies]
visualization = ["matplotlib>=3.5"]
+test = [
+ "build>=1.0",
+ "pytest>=7.0"
+]
+
+[project.urls]
+Paper = "https://doi.org/10.1038/s41551-026-01741-4"
+Repository = "https://github.com/peterhan91/CLEAR"
+Model = "https://huggingface.co/peterhan91/CLEAR"
[tool.setuptools]
package-dir = {"" = "src"}
diff --git a/src/clear/__init__.py b/src/clear/__init__.py
index 5425448..85d2a39 100644
--- a/src/clear/__init__.py
+++ b/src/clear/__init__.py
@@ -2,6 +2,7 @@
from .clip import available_models, load, tokenize
from .concepts import ConceptBank
+from .hub import build_cxr_preprocess, load_pretrained
from .pipeline import PromptPair, PromptSet, ZeroShotPipeline, build_prompt_embeddings
__all__ = [
@@ -9,6 +10,8 @@
"load",
"tokenize",
"ConceptBank",
+ "build_cxr_preprocess",
+ "load_pretrained",
"PromptPair",
"PromptSet",
"ZeroShotPipeline",
diff --git a/src/clear/concepts.py b/src/clear/concepts.py
index f0e8295..223be0e 100644
--- a/src/clear/concepts.py
+++ b/src/clear/concepts.py
@@ -2,16 +2,31 @@
from __future__ import annotations
from dataclasses import dataclass, field
+from numbers import Integral
from pathlib import Path
-from typing import Iterable, List, Optional, Sequence
+from typing import Any, Iterable, List, Mapping, Optional, Sequence
-import json
import csv
+import json
+import pickle
+import warnings
+
import numpy as np
__all__ = ["ConceptBank"]
+class _NumPyCompatUnpickler(pickle.Unpickler):
+ """Read NumPy 2 pickles on NumPy 1.x after the caller explicitly opts in."""
+
+ def find_class(self, module: str, name: str) -> Any:
+ # NumPy 2 renamed ``numpy.core`` to ``numpy._core``. The published SFR
+ # concept file uses the new path, while CLEAR still supports NumPy 1.x.
+ if module == "numpy._core" or module.startswith("numpy._core."):
+ module = "numpy.core" + module[len("numpy._core") :]
+ return super().find_class(module, name)
+
+
def _load_concept_list(path: Path) -> List[str]:
if not path.exists():
raise FileNotFoundError(path)
@@ -53,22 +68,154 @@ def _load_concept_list(path: Path) -> List[str]:
return [line.strip() for line in handle if line.strip()]
-def _load_array(path: Optional[Path]) -> Optional[np.ndarray]:
+def _safe_torch_load(path: Path) -> Any:
+ """Load tensor-only PyTorch files without enabling arbitrary unpickling."""
+
+ try:
+ import torch
+ except ImportError as exc: # pragma: no cover - torch is a core dependency
+ raise ImportError("Loading .pt embeddings requires PyTorch") from exc
+
+ try:
+ return torch.load(path, map_location="cpu", weights_only=True)
+ except TypeError as exc:
+ # ``weights_only`` was added after the oldest PyTorch version originally
+ # supported by CLEAR. Retain compatibility, but make the security
+ # difference explicit instead of silently falling back to pickle.
+ if "weights_only" not in str(exc):
+ raise
+ warnings.warn(
+ "This PyTorch version does not support safe weights-only loading; "
+ f"only continue if you trust {path}.",
+ RuntimeWarning,
+ stacklevel=2,
+ )
+ return torch.load(path, map_location="cpu")
+
+
+def _stack_indexed_embeddings(payload: Mapping[Any, Any], path: Path) -> np.ndarray:
+ """Convert the released ``{concept_idx: vector}`` pickle into a matrix."""
+
+ if not payload:
+ raise ValueError(f"Embedding dictionary in {path} is empty")
+ if not all(isinstance(key, Integral) and not isinstance(key, bool) for key in payload):
+ raise ValueError(
+ f"Embedding dictionary in {path} must use integer concept indices as keys"
+ )
+
+ keys = sorted(int(key) for key in payload)
+ missing = next((index for index, key in enumerate(keys) if index != key), None)
+ if missing is not None:
+ raise ValueError(
+ f"Embedding dictionary in {path} must have contiguous keys starting at 0; "
+ f"first missing index is {missing}"
+ )
+
+ vectors: List[np.ndarray] = []
+ vector_dim: Optional[int] = None
+ for index in keys:
+ vector = np.asarray(payload[index], dtype=np.float32)
+ if vector.ndim != 1:
+ raise ValueError(
+ f"Embedding at index {index} in {path} must be one-dimensional; "
+ f"got shape {vector.shape}"
+ )
+ if vector_dim is None:
+ vector_dim = int(vector.shape[0])
+ if vector_dim == 0:
+ raise ValueError(f"Embedding vectors in {path} cannot be empty")
+ elif vector.shape[0] != vector_dim:
+ raise ValueError(
+ f"Embedding at index {index} in {path} has dimension {vector.shape[0]}; "
+ f"expected {vector_dim}"
+ )
+ vectors.append(vector)
+
+ return np.stack(vectors, axis=0)
+
+
+def _coerce_embedding_payload(payload: Any, path: Path) -> np.ndarray:
+ """Extract one numeric embedding matrix from a supported file payload."""
+
+ try:
+ import torch
+ except ImportError: # pragma: no cover - torch is a core dependency
+ torch = None
+
+ if torch is not None and torch.is_tensor(payload):
+ array = payload.detach().cpu().numpy()
+ elif isinstance(payload, np.ndarray):
+ array = payload
+ elif isinstance(payload, Mapping):
+ for field in ("embeddings", "concept_embeddings"):
+ if field in payload:
+ return _coerce_embedding_payload(payload[field], path)
+ return _stack_indexed_embeddings(payload, path)
+ else:
+ raise ValueError(
+ f"Unsupported embedding payload in {path}: expected a tensor, array, "
+ "or an integer-indexed embedding dictionary"
+ )
+
+ try:
+ array = np.asarray(array, dtype=np.float32)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(f"Embeddings in {path} must be numeric") from exc
+ if array.ndim != 2:
+ raise ValueError(
+ f"Embeddings in {path} must be a two-dimensional matrix; got shape {array.shape}"
+ )
+ if 0 in array.shape:
+ raise ValueError(f"Embedding matrix in {path} cannot have an empty dimension")
+ return array
+
+
+def _load_array(
+ path: Optional[Path],
+ *,
+ allow_unsafe_pickle: bool = False,
+) -> Optional[np.ndarray]:
if path is None:
return None
if not path.exists():
raise FileNotFoundError(path)
- data = np.load(path)
+ suffix = path.suffix.lower()
+ if suffix in {".pt", ".pth"}:
+ return _coerce_embedding_payload(_safe_torch_load(path), path)
+
+ if suffix in {".pickle", ".pkl"}:
+ if not allow_unsafe_pickle:
+ raise ValueError(
+ f"Refusing to unpickle {path} because pickle files can execute arbitrary "
+ "code. Pass allow_unsafe_pickle=True only for a file you trust."
+ )
+ warnings.warn(
+ f"Unpickling {path}; only trusted pickle files should be loaded.",
+ RuntimeWarning,
+ stacklevel=2,
+ )
+ with path.open("rb") as handle:
+ payload = _NumPyCompatUnpickler(handle).load()
+ return _coerce_embedding_payload(payload, path)
+
+ # ``allow_pickle=False`` is intentional. It preserves safe loading for
+ # NumPy arrays and also supports extensionless legacy .npy/.npz paths.
+ data = np.load(path, allow_pickle=False)
if isinstance(data, np.lib.npyio.NpzFile):
- if "embeddings" in data:
- array = data["embeddings"]
- else:
- # fall back to the first entry
- array = data[data.files[0]]
+ try:
+ if not data.files:
+ raise ValueError(f"NumPy archive {path} contains no arrays")
+ if "embeddings" in data:
+ array = data["embeddings"]
+ else:
+ # Preserve the original behavior for unnamed archives.
+ array = data[data.files[0]]
+ finally:
+ data.close()
else:
array = data
- return np.asarray(array, dtype=np.float32)
+ return _coerce_embedding_payload(array, path)
@dataclass(slots=True)
@@ -99,14 +246,22 @@ def __post_init__(self) -> None:
n_concepts = len(self.concepts)
if self.clip_embeddings is not None:
+ self.clip_embeddings = _coerce_embedding_payload(
+ self.clip_embeddings, Path("")
+ )
if self.clip_embeddings.shape[0] != n_concepts:
raise ValueError(
- "clip_embeddings rows must match number of concepts"
+ "clip_embeddings rows must match number of concepts: "
+ f"got {self.clip_embeddings.shape[0]} rows for {n_concepts} concepts"
)
if self.llm_embeddings is not None:
+ self.llm_embeddings = _coerce_embedding_payload(
+ self.llm_embeddings, Path("")
+ )
if self.llm_embeddings.shape[0] != n_concepts:
raise ValueError(
- "llm_embeddings rows must match number of concepts"
+ "llm_embeddings rows must match number of concepts: "
+ f"got {self.llm_embeddings.shape[0]} rows for {n_concepts} concepts"
)
@classmethod
@@ -117,18 +272,36 @@ def from_files(
clip_embeddings_path: str | Path | None = None,
llm_embeddings_path: str | Path | None = None,
metadata: Optional[dict] = None,
+ allow_unsafe_pickle: bool = False,
) -> "ConceptBank":
"""Load a concept bank from disk.
``concepts_path`` may be a JSON Lines file containing objects with a
``concept`` field, a CSV/TSV file, or a simple newline-delimited text
- file. Embedding paths may point to ``.npy`` or ``.npz`` files containing
- an array named ``embeddings`` (or the first array when unnamed).
+ file. Embedding paths may point to ``.npy``/``.npz`` arrays, a
+ tensor-only PyTorch ``.pt`` file, or the released integer-indexed
+ ``.pickle`` dictionary. Because pickle can execute arbitrary code,
+ loading a pickle requires the explicit ``allow_unsafe_pickle=True``
+ opt-in and should only be used with a trusted file.
"""
concepts = _load_concept_list(Path(concepts_path))
- clip_embeddings = _load_array(Path(clip_embeddings_path)) if clip_embeddings_path else None
- llm_embeddings = _load_array(Path(llm_embeddings_path)) if llm_embeddings_path else None
+ clip_embeddings = (
+ _load_array(
+ Path(clip_embeddings_path),
+ allow_unsafe_pickle=allow_unsafe_pickle,
+ )
+ if clip_embeddings_path
+ else None
+ )
+ llm_embeddings = (
+ _load_array(
+ Path(llm_embeddings_path),
+ allow_unsafe_pickle=allow_unsafe_pickle,
+ )
+ if llm_embeddings_path
+ else None
+ )
return cls(
concepts=concepts,
clip_embeddings=clip_embeddings,
@@ -185,4 +358,3 @@ def iter_chunks(self, chunk_size: int) -> Iterable[tuple[np.ndarray, Optional[np
clip_chunk = clip[start:stop]
llm_chunk = llm[start:stop] if llm is not None else None
yield clip_chunk, llm_chunk
-
diff --git a/src/clear/hub.py b/src/clear/hub.py
new file mode 100644
index 0000000..b579dad
--- /dev/null
+++ b/src/clear/hub.py
@@ -0,0 +1,305 @@
+"""Official loading and chest-X-ray preprocessing for released CLEAR weights."""
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any, Callable, Optional, Tuple
+import warnings
+
+import torch
+import torch.nn as nn
+from torchvision.transforms import (
+ Compose,
+ InterpolationMode,
+ Lambda,
+ Normalize,
+ Resize,
+ ToTensor,
+)
+
+from .model import CLIP
+
+__all__ = ["build_cxr_preprocess", "load_pretrained"]
+
+
+DEFAULT_REPO_ID = "peterhan91/CLEAR"
+DEFAULT_CHECKPOINT_FILENAME = "best_model.pt"
+CXR_IMAGE_SIZE = 448
+CXR_PIXEL_MEAN = 101.48761
+CXR_PIXEL_STD = 83.43944
+
+
+class _DinoV2Visual(nn.Module):
+ """Expose DINOv2 under the names used by the published state dictionary."""
+
+ def __init__(self, backbone: nn.Module, output_dim: int = 768) -> None:
+ super().__init__()
+ backbone_dim = getattr(backbone, "embed_dim", None)
+ if backbone_dim is None:
+ backbone_dim = getattr(backbone, "num_features", None)
+ if backbone_dim is None:
+ raise ValueError("Could not determine the DINOv2 backbone embedding dimension")
+
+ self.backbone = backbone
+ self.projection = nn.Linear(int(backbone_dim), output_dim)
+ self.input_resolution = CXR_IMAGE_SIZE
+ self.output_dim = output_dim
+
+ def forward(self, images: torch.Tensor) -> torch.Tensor:
+ return self.projection(self.backbone(images))
+
+ @property
+ def conv1(self) -> nn.Module:
+ # ``CLIP.dtype`` uses ``visual.conv1.weight.dtype``. The compatibility
+ # property keeps that API intact after replacing the original vision
+ # tower with DINOv2.
+ return self.projection
+
+
+def _convert_rgb(image: Any) -> Any:
+ if not hasattr(image, "convert"):
+ raise TypeError("CLEAR preprocessing expects a PIL image")
+ return image.convert("RGB")
+
+
+def build_cxr_preprocess(image_size: int = CXR_IMAGE_SIZE) -> Compose:
+ """Build the preprocessing used for CLEAR chest-X-ray experiments.
+
+ The original experiments normalized 0--255 image tensors with a mean of
+ ``101.48761`` and standard deviation of ``83.43944``. ``ToTensor`` scales
+ PIL pixels to 0--1, so the values are divided by 255 here. Images are
+ resized to an exact square rather than aspect-ratio resized and cropped.
+ """
+
+ if not isinstance(image_size, int) or isinstance(image_size, bool) or image_size <= 0:
+ raise ValueError("image_size must be a positive integer")
+ mean = (CXR_PIXEL_MEAN / 255.0,) * 3
+ std = (CXR_PIXEL_STD / 255.0,) * 3
+ return Compose(
+ [
+ Lambda(_convert_rgb),
+ Resize(
+ (image_size, image_size),
+ interpolation=InterpolationMode.BICUBIC,
+ antialias=True,
+ ),
+ ToTensor(),
+ Normalize(mean=mean, std=std),
+ ]
+ )
+
+
+def _safe_torch_load(path: Path) -> Any:
+ try:
+ return torch.load(path, map_location="cpu", weights_only=True)
+ except TypeError as exc:
+ # Compatibility with PyTorch releases predating ``weights_only``. This
+ # warning is intentionally prominent because the fallback uses pickle.
+ if "weights_only" not in str(exc):
+ raise
+ warnings.warn(
+ "This PyTorch version does not support safe weights-only loading; "
+ f"only continue if you trust checkpoint {path}.",
+ RuntimeWarning,
+ stacklevel=2,
+ )
+ return torch.load(path, map_location="cpu")
+
+
+def _extract_state_dict(checkpoint: Any) -> dict[str, torch.Tensor]:
+ if not isinstance(checkpoint, Mapping):
+ raise ValueError("CLEAR checkpoint must contain a PyTorch state dictionary")
+
+ payload: Mapping[Any, Any] = checkpoint
+ for field in ("state_dict", "model_state_dict"):
+ candidate = payload.get(field)
+ if isinstance(candidate, Mapping):
+ payload = candidate
+ break
+
+ if not payload or not all(isinstance(key, str) for key in payload):
+ raise ValueError("CLEAR state-dictionary keys must be non-empty strings")
+ if not all(torch.is_tensor(value) for value in payload.values()):
+ raise ValueError("CLEAR state dictionary may contain tensors only")
+
+ state_dict = dict(payload)
+ if all(key.startswith("module.") for key in state_dict):
+ state_dict = {key.removeprefix("module."): value for key, value in state_dict.items()}
+
+ required = {
+ "visual.backbone.register_tokens",
+ "visual.backbone.patch_embed.proj.weight",
+ "visual.projection.weight",
+ "positional_embedding",
+ "text_projection",
+ "token_embedding.weight",
+ "ln_final.weight",
+ }
+ missing = sorted(required.difference(state_dict))
+ if missing:
+ raise ValueError(
+ "Checkpoint is not the released CLEAR DINOv2-register architecture; "
+ f"missing keys: {', '.join(missing)}"
+ )
+ return state_dict
+
+
+def _text_config(state_dict: Mapping[str, torch.Tensor]) -> dict[str, int]:
+ positional_embedding = state_dict["positional_embedding"]
+ text_projection = state_dict["text_projection"]
+ token_embedding = state_dict["token_embedding.weight"]
+ ln_final = state_dict["ln_final.weight"]
+
+ if positional_embedding.ndim != 2 or text_projection.ndim != 2:
+ raise ValueError("CLEAR text positional and projection weights must be matrices")
+ if token_embedding.ndim != 2 or ln_final.ndim != 1:
+ raise ValueError("CLEAR token embedding or final normalization has an invalid shape")
+
+ context_length, transformer_width = map(int, positional_embedding.shape)
+ vocab_size, token_width = map(int, token_embedding.shape)
+ projection_width, embed_dim = map(int, text_projection.shape)
+ if token_width != transformer_width or projection_width != transformer_width:
+ raise ValueError("CLEAR checkpoint contains inconsistent text-encoder dimensions")
+ if int(ln_final.shape[0]) != transformer_width or transformer_width % 64:
+ raise ValueError("CLEAR checkpoint contains an unsupported text-transformer width")
+
+ layer_indices = {
+ int(key.split(".")[2])
+ for key in state_dict
+ if key.startswith("transformer.resblocks.") and key.count(".") >= 3
+ }
+ if not layer_indices or layer_indices != set(range(max(layer_indices) + 1)):
+ raise ValueError("CLEAR checkpoint text-transformer layers are not contiguous")
+
+ return {
+ "embed_dim": embed_dim,
+ "context_length": context_length,
+ "vocab_size": vocab_size,
+ "transformer_width": transformer_width,
+ "transformer_heads": transformer_width // 64,
+ "transformer_layers": len(layer_indices),
+ }
+
+
+def _cached_dinov2_repo() -> Optional[Path]:
+ candidate = Path(torch.hub.get_dir()) / "facebookresearch_dinov2_main"
+ return candidate if (candidate / "hubconf.py").is_file() else None
+
+
+def _build_dinov2_backbone(*, local_files_only: bool) -> nn.Module:
+ cached_repo = _cached_dinov2_repo()
+ if cached_repo is not None:
+ return torch.hub.load(
+ str(cached_repo),
+ "dinov2_vitb14_reg",
+ source="local",
+ pretrained=False,
+ verbose=False,
+ )
+ if local_files_only:
+ raise FileNotFoundError(
+ "DINOv2 source code is not present in the torch hub cache and "
+ "local_files_only=True prevents downloading it"
+ )
+ return torch.hub.load(
+ "facebookresearch/dinov2",
+ "dinov2_vitb14_reg",
+ pretrained=False,
+ trust_repo=True,
+ verbose=False,
+ )
+
+
+def _build_model(
+ state_dict: Mapping[str, torch.Tensor],
+ *,
+ local_files_only: bool,
+) -> CLIP:
+ config = _text_config(state_dict)
+ if config["embed_dim"] != 768:
+ raise ValueError(
+ "The released DINOv2 ViT-B/14 checkpoint must use 768-dimensional features; "
+ f"got {config['embed_dim']}"
+ )
+
+ # CLIP owns the text encoder. A deliberately small temporary vision tower
+ # avoids allocating a second ViT-B while CLIP is initialized; it is replaced
+ # immediately by the published DINOv2-register architecture.
+ model = CLIP(
+ embed_dim=config["embed_dim"],
+ image_resolution=14,
+ vision_layers=1,
+ vision_width=64,
+ vision_patch_size=14,
+ context_length=config["context_length"],
+ vocab_size=config["vocab_size"],
+ transformer_width=config["transformer_width"],
+ transformer_heads=config["transformer_heads"],
+ transformer_layers=config["transformer_layers"],
+ )
+ model.visual = _DinoV2Visual(
+ _build_dinov2_backbone(local_files_only=local_files_only),
+ output_dim=config["embed_dim"],
+ )
+ return model
+
+
+def load_pretrained(
+ checkpoint_path: str | Path | None = None,
+ *,
+ repo_id: str = DEFAULT_REPO_ID,
+ filename: str = DEFAULT_CHECKPOINT_FILENAME,
+ revision: str | None = None,
+ cache_dir: str | Path | None = None,
+ device: str | torch.device | None = None,
+ local_files_only: bool = False,
+ strict: bool = True,
+) -> Tuple[CLIP, Callable[[Any], torch.Tensor]]:
+ """Load the published CLEAR backbone and its official CXR preprocessing.
+
+ When ``checkpoint_path`` is omitted, ``best_model.pt`` is downloaded from
+ ``peterhan91/CLEAR`` on Hugging Face. The checkpoint fully supplies the
+ DINOv2 ViT-B/14-register weights; no separate pretrained DINOv2 weights are
+ downloaded. ``local_files_only=True`` requires both the checkpoint and the
+ DINOv2 torch-hub source to already be cached.
+ """
+
+ if checkpoint_path is None:
+ try:
+ from huggingface_hub import hf_hub_download
+ except ImportError as exc:
+ raise ImportError(
+ "Loading CLEAR from Hugging Face requires huggingface-hub; "
+ "install the project dependencies or pass checkpoint_path"
+ ) from exc
+ downloaded = hf_hub_download(
+ repo_id=repo_id,
+ filename=filename,
+ revision=revision,
+ cache_dir=str(cache_dir) if cache_dir is not None else None,
+ local_files_only=local_files_only,
+ )
+ resolved_checkpoint = Path(downloaded)
+ else:
+ resolved_checkpoint = Path(checkpoint_path).expanduser()
+
+ if not resolved_checkpoint.is_file():
+ raise FileNotFoundError(resolved_checkpoint)
+
+ state_dict = _extract_state_dict(_safe_torch_load(resolved_checkpoint))
+ model = _build_model(state_dict, local_files_only=local_files_only)
+ incompatible = model.load_state_dict(state_dict, strict=strict)
+ if not strict and (incompatible.missing_keys or incompatible.unexpected_keys):
+ warnings.warn(
+ "CLEAR checkpoint was loaded non-strictly; missing keys: "
+ f"{incompatible.missing_keys}; unexpected keys: {incompatible.unexpected_keys}",
+ RuntimeWarning,
+ stacklevel=2,
+ )
+
+ target_device = torch.device(
+ device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
+ )
+ model.to(target_device)
+ model.eval()
+ return model, build_cxr_preprocess()
diff --git a/tests/test_concepts.py b/tests/test_concepts.py
new file mode 100644
index 0000000..1c4965e
--- /dev/null
+++ b/tests/test_concepts.py
@@ -0,0 +1,76 @@
+import pickle
+
+import numpy as np
+import pytest
+import torch
+
+from clear import ConceptBank
+
+
+def _write_concepts(path) -> None:
+ path.write_text("concept,concept_idx\npleural effusion,0\nno pleural effusion,1\n")
+
+
+@pytest.mark.parametrize("suffix", [".npy", ".npz", ".pt"])
+@pytest.mark.filterwarnings("ignore:TypedStorage is deprecated:UserWarning")
+def test_concept_bank_loads_safe_embedding_formats(tmp_path, suffix: str) -> None:
+ concepts_path = tmp_path / "mimic_concepts.csv"
+ embeddings_path = tmp_path / f"concept_embeddings{suffix}"
+ expected = np.asarray([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+ _write_concepts(concepts_path)
+
+ if suffix == ".npy":
+ np.save(embeddings_path, expected)
+ elif suffix == ".npz":
+ np.savez(embeddings_path, embeddings=expected)
+ else:
+ torch.save(torch.from_numpy(expected), embeddings_path)
+
+ bank = ConceptBank.from_files(
+ concepts_path,
+ clip_embeddings_path=embeddings_path,
+ )
+
+ assert bank.concepts == ("pleural effusion", "no pleural effusion")
+ np.testing.assert_array_equal(bank.ensure_clip_embeddings(), expected)
+
+
+def test_pickle_loading_requires_explicit_opt_in(tmp_path) -> None:
+ concepts_path = tmp_path / "mimic_concepts.csv"
+ embeddings_path = tmp_path / "concepts_embeddings_sfr_mistral.pickle"
+ _write_concepts(concepts_path)
+ with embeddings_path.open("wb") as handle:
+ pickle.dump({0: [1.0, 2.0], 1: [3.0, 4.0]}, handle)
+
+ with pytest.raises(ValueError, match="pickle"):
+ ConceptBank.from_files(
+ concepts_path,
+ llm_embeddings_path=embeddings_path,
+ )
+
+ with pytest.warns(RuntimeWarning, match="trusted pickle"):
+ bank = ConceptBank.from_files(
+ concepts_path,
+ llm_embeddings_path=embeddings_path,
+ allow_unsafe_pickle=True,
+ )
+ np.testing.assert_array_equal(
+ bank.ensure_llm_embeddings(),
+ np.asarray([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
+ )
+
+
+def test_pickle_mapping_indices_must_be_contiguous(tmp_path) -> None:
+ concepts_path = tmp_path / "mimic_concepts.csv"
+ embeddings_path = tmp_path / "concepts_embeddings_sfr_mistral.pickle"
+ _write_concepts(concepts_path)
+ with embeddings_path.open("wb") as handle:
+ pickle.dump({0: [1.0, 2.0], 2: [3.0, 4.0]}, handle)
+
+ with pytest.warns(RuntimeWarning, match="trusted pickle"):
+ with pytest.raises(ValueError, match="contiguous"):
+ ConceptBank.from_files(
+ concepts_path,
+ llm_embeddings_path=embeddings_path,
+ allow_unsafe_pickle=True,
+ )
diff --git a/tests/test_hub.py b/tests/test_hub.py
new file mode 100644
index 0000000..2741b0d
--- /dev/null
+++ b/tests/test_hub.py
@@ -0,0 +1,141 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+import clear.hub as hub
+
+
+def _synthetic_state_dict() -> dict[str, torch.Tensor]:
+ return {
+ "visual.backbone.register_tokens": torch.zeros(1, 4, 768),
+ "visual.backbone.patch_embed.proj.weight": torch.zeros(768, 3, 14, 14),
+ "visual.projection.weight": torch.zeros(768, 768),
+ "positional_embedding": torch.zeros(8, 64),
+ "text_projection": torch.zeros(64, 768),
+ "token_embedding.weight": torch.zeros(16, 64),
+ "ln_final.weight": torch.zeros(64),
+ "transformer.resblocks.0.attn.in_proj_weight": torch.zeros(192, 64),
+ }
+
+
+class _FakeModel(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.anchor = nn.Parameter(torch.zeros(()))
+ self.loaded_state_dict = None
+ self.loaded_strict = None
+
+ def load_state_dict(self, state_dict, strict=True):
+ self.loaded_state_dict = state_dict
+ self.loaded_strict = strict
+ return SimpleNamespace(missing_keys=[], unexpected_keys=[])
+
+
+def _mock_model_loading(monkeypatch, checkpoint):
+ fake_model = _FakeModel()
+ captured = {}
+
+ monkeypatch.setattr(hub, "_safe_torch_load", lambda path: checkpoint)
+
+ def fake_build_model(state_dict, *, local_files_only):
+ captured["state_dict"] = state_dict
+ captured["local_files_only"] = local_files_only
+ return fake_model
+
+ monkeypatch.setattr(hub, "_build_model", fake_build_model)
+ return fake_model, captured
+
+
+def test_extract_state_dict_unwraps_and_removes_data_parallel_prefix() -> None:
+ state_dict = _synthetic_state_dict()
+ checkpoint = {
+ "state_dict": {f"module.{key}": value for key, value in state_dict.items()}
+ }
+
+ extracted = hub._extract_state_dict(checkpoint)
+
+ assert extracted.keys() == state_dict.keys()
+ assert extracted["text_projection"] is state_dict["text_projection"]
+
+
+def test_load_pretrained_uses_local_checkpoint_without_hub_download(
+ tmp_path, monkeypatch
+) -> None:
+ checkpoint_path = tmp_path / "best_model.pt"
+ checkpoint_path.touch()
+ checkpoint = {
+ "model_state_dict": {
+ f"module.{key}": value for key, value in _synthetic_state_dict().items()
+ }
+ }
+ model, captured = _mock_model_loading(monkeypatch, checkpoint)
+ monkeypatch.setattr(
+ "huggingface_hub.hf_hub_download",
+ lambda **kwargs: pytest.fail("local checkpoint must not access Hugging Face"),
+ )
+
+ loaded_model, preprocess = hub.load_pretrained(
+ checkpoint_path,
+ device="cpu",
+ local_files_only=True,
+ strict=False,
+ )
+
+ assert loaded_model is model
+ assert model.loaded_strict is False
+ assert model.training is False
+ assert model.anchor.device.type == "cpu"
+ assert captured["local_files_only"] is True
+ assert "module.text_projection" not in captured["state_dict"]
+ assert callable(preprocess)
+
+
+def test_load_pretrained_forwards_hugging_face_download_options(
+ tmp_path, monkeypatch
+) -> None:
+ checkpoint_path = tmp_path / "downloaded.pt"
+ checkpoint_path.touch()
+ _mock_model_loading(monkeypatch, _synthetic_state_dict())
+ captured = {}
+
+ def fake_hf_hub_download(**kwargs):
+ captured.update(kwargs)
+ return str(checkpoint_path)
+
+ monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_hf_hub_download)
+
+ hub.load_pretrained(
+ repo_id="example/CLEAR",
+ filename="weights.pt",
+ revision="publication",
+ cache_dir=tmp_path / "cache",
+ device="cpu",
+ local_files_only=True,
+ )
+
+ assert captured == {
+ "repo_id": "example/CLEAR",
+ "filename": "weights.pt",
+ "revision": "publication",
+ "cache_dir": str(tmp_path / "cache"),
+ "local_files_only": True,
+ }
+
+
+def test_cxr_preprocess_converts_grayscale_and_uses_published_normalization() -> None:
+ preprocess = hub.build_cxr_preprocess(image_size=32)
+
+ output = preprocess(Image.new("L", (12, 20), color=101))
+
+ assert output.shape == (3, 32, 32)
+ expected = (101.0 - hub.CXR_PIXEL_MEAN) / hub.CXR_PIXEL_STD
+ torch.testing.assert_close(output, torch.full_like(output, expected), atol=1e-6, rtol=0)
+
+
+@pytest.mark.parametrize("image_size", [0, -1, 32.5, True])
+def test_cxr_preprocess_rejects_invalid_sizes(image_size) -> None:
+ with pytest.raises(ValueError, match="positive integer"):
+ hub.build_cxr_preprocess(image_size=image_size)
diff --git a/tests/test_package.py b/tests/test_package.py
new file mode 100644
index 0000000..0b1917e
--- /dev/null
+++ b/tests/test_package.py
@@ -0,0 +1,14 @@
+from importlib import resources
+
+import clear
+
+
+def test_package_import_exposes_supported_entry_points() -> None:
+ assert callable(clear.load_pretrained)
+ assert callable(clear.build_cxr_preprocess)
+ assert clear.ConceptBank.__module__ == "clear.concepts"
+
+
+def test_tokenizer_vocabulary_is_packaged() -> None:
+ vocabulary = resources.files("clear").joinpath("bpe_simple_vocab_16e6.txt.gz")
+ assert vocabulary.is_file()
diff --git a/tests/test_quickstart.py b/tests/test_quickstart.py
new file mode 100644
index 0000000..4362304
--- /dev/null
+++ b/tests/test_quickstart.py
@@ -0,0 +1,39 @@
+import pytest
+import torch
+
+from examples.quickstart import _rank_concepts
+
+
+def test_rank_concepts_streams_chunks_and_returns_global_top_k() -> None:
+ image_features = torch.tensor([[1.0, 0.0]])
+ concept_embeddings = torch.tensor(
+ [
+ [1.0, 0.0],
+ [0.0, 1.0],
+ [0.8, 0.2],
+ [-1.0, 0.0],
+ ]
+ )
+ concepts = ["first", "second", "third", "fourth"]
+
+ ranked = _rank_concepts(
+ image_features,
+ concept_embeddings,
+ concepts,
+ top_k=2,
+ chunk_size=2,
+ )
+
+ assert [concept for concept, _ in ranked] == ["first", "third"]
+ assert ranked[0][1] == pytest.approx(1.0)
+
+
+def test_rank_concepts_rejects_misaligned_artifacts() -> None:
+ with pytest.raises(ValueError, match="misaligned"):
+ _rank_concepts(
+ torch.tensor([[1.0, 0.0]]),
+ torch.tensor([[1.0, 0.0]]),
+ ["first", "second"],
+ top_k=1,
+ chunk_size=1,
+ )