diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000..e0ea6cc
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,52 @@
+# Dependabot configuration for mmcontext
+# Docs: https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file
+#
+# Security updates (PRs that fix Dependabot *alerts*) are enabled separately in
+# Settings -> Code security and obey target-branch below. This file additionally
+# enables proactive *version* updates on a schedule, grouped to reduce PR noise.
+version: 2
+updates:
+  # Python dependencies declared in pyproject.toml
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+      day: "monday"
+    # PRs target the integration branch, never main, per the project branch strategy.
+    target-branch: "dev-claude"
+    open-pull-requests-limit: 5
+    labels:
+      - "dependencies"
+    commit-message:
+      prefix: "deps"
+    # One grouped PR for routine minor/patch bumps instead of one PR per package.
+    groups:
+      python-minor-patch:
+        update-types:
+          - "minor"
+          - "patch"
+    # Heavy, tightly version-pinned core libs: keep these as individual PRs so a
+    # major bump (e.g. torch, transformers) is reviewed in isolation.
+    ignore:
+      - dependency-name: "torch"
+        update-types: ["version-update:semver-major"]
+      - dependency-name: "sentence-transformers"
+        update-types: ["version-update:semver-major"]
+
+  # GitHub Actions used in .github/workflows/*
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+      day: "monday"
+    target-branch: "dev-claude"
+    open-pull-requests-limit: 5
+    labels:
+      - "dependencies"
+      - "github-actions"
+    commit-message:
+      prefix: "ci"
+    groups:
+      actions-all:
+        patterns:
+          - "*"
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 94562c6..971230b 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -14,9 +14,9 @@ jobs:
   package:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - name: Set up Python 3.12
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: "3.12"
       - name: Install build dependencies
diff --git a/.github/workflows/claude-implement.yaml b/.github/workflows/claude-implement.yaml
new file mode 100644
index 0000000..4257f1d
--- /dev/null
+++ b/.github/workflows/claude-implement.yaml
@@ -0,0 +1,61 @@
+name: Claude Code Implement
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: write
+  pull-requests: write
+  issues: write
+  actions: read
+  id-token: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.issue.number }}-${{ github.actor }}
+  cancel-in-progress: true
+
+jobs:
+  implement:
+    # Trigger on @claude mentions in issue comments — only by repo owner, ignore bots
+    if: >
+      github.actor == 'mengerj' &&
+      github.event.sender.type != 'Bot' &&
+      !github.event.issue.pull_request &&
+      contains(github.event.comment.body, '@claude')
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          ref: dev-claude
+          fetch-depth: 0
+      - uses: anthropics/claude-code-action@v1
+        id: claude
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+          show_full_output: true
+          base_branch: dev-claude
+          additional_permissions: |
+            actions: read
+          claude_args: |
+            --allowedTools "Bash,Read,Write,Edit,Glob,Grep,Task,WebSearch,WebFetch"
+            --max-turns 30
+            --output-format stream-json
+          # The CLAUDE.md in the repo root provides the plan-first protocol,
+          # branch naming conventions, and PR linking instructions.
+          # Claude reads it automatically on checkout.
+
+      - name: Archive Claude session log
+        if: always()
+        run: |
+          mkdir -p claude-logs
+          echo '${{ steps.claude.outputs.result }}' > claude-logs/session.json
+        shell: bash
+
+      - name: Upload session log
+        if: always()
+        uses: actions/upload-artifact@v7
+        with:
+          name: claude-implement-issue${{ github.event.issue.number }}-${{ github.run_number }}
+          path: claude-logs/
+          retention-days: 30
diff --git a/.github/workflows/claude-review.yaml b/.github/workflows/claude-review.yaml
new file mode 100644
index 0000000..92a1069
--- /dev/null
+++ b/.github/workflows/claude-review.yaml
@@ -0,0 +1,56 @@
+name: Claude Code Review
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: write
+  pull-requests: write
+  issues: write
+  actions: read
+  id-token: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.event.issue.number }}-${{ github.actor }}
+  cancel-in-progress: true
+
+jobs:
+  review:
+    # Respond to @claude in PR comments (owner only, ignore bots)
+    if: >
+      github.event.sender.type != 'Bot' &&
+      github.actor == 'mengerj' &&
+      github.event.issue.pull_request &&
+      contains(github.event.comment.body, '@claude')
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - uses: anthropics/claude-code-action@v1
+        id: claude
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+          show_full_output: true
+          additional_permissions: |
+            actions: read
+          claude_args: |
+            --allowedTools "Bash,Read,Write,Edit,Glob,Grep,Task"
+            --max-turns 30
+            --output-format stream-json
+
+      - name: Archive Claude session log
+        if: always()
+        run: |
+          mkdir -p claude-logs
+          echo '${{ steps.claude.outputs.result }}' > claude-logs/session.json
+        shell: bash
+
+      - name: Upload session log
+        if: always()
+        uses: actions/upload-artifact@v7
+        with:
+          name: claude-review-pr${{ github.event.pull_request.number || github.event.issue.number }}-${{ github.run_number }}
+          path: claude-logs/
+          retention-days: 30
diff --git a/.github/workflows/codeql.yaml b/.github/workflows/codeql.yaml
new file mode 100644
index 0000000..d24c6de
--- /dev/null
+++ b/.github/workflows/codeql.yaml
@@ -0,0 +1,43 @@
+name: CodeQL
+
+# Semantic code scanning for Python. Findings appear under the repo's
+# Security -> Code scanning tab and as PR annotations.
+# Docs: https://docs.github.com/code-security/code-scanning
+
+on:
+  push:
+    branches: [main, dev-claude]
+  pull_request:
+    branches: [main, dev-claude]
+  schedule:
+    # Weekly full scan (Mondays 06:00 UTC) to catch newly-disclosed query updates.
+    - cron: "0 6 * * 1"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  analyze:
+    name: Analyze (python)
+    runs-on: ubuntu-latest
+    permissions:
+      actions: read
+      contents: read
+      security-events: write
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+
+      - name: Initialize CodeQL
+        uses: github/codeql-action/init@v4
+        with:
+          languages: python
+          # security-and-quality adds maintainability/quality queries on top of
+          # the default security set. Drop to "security-extended" if too noisy.
+          queries: security-and-quality
+
+      - name: Perform CodeQL Analysis
+        uses: github/codeql-action/analyze@v4
+        with:
+          category: "/language:python"
diff --git a/.github/workflows/dependabot-auto-merge.yaml b/.github/workflows/dependabot-auto-merge.yaml
new file mode 100644
index 0000000..07e1872
--- /dev/null
+++ b/.github/workflows/dependabot-auto-merge.yaml
@@ -0,0 +1,36 @@
+name: Dependabot auto-merge
+
+# Auto-approves and enables auto-merge for low-risk Dependabot PRs (patch/minor,
+# including security updates). GitHub still waits for all required status checks
+# (tests, CodeQL) to pass before the merge actually happens.
+#
+# Prerequisites (one-time, in repo Settings):
+#   - "Allow auto-merge" enabled (Settings -> General -> Pull Requests)
+#   - Branch protection on dev-claude requiring the Test + CodeQL checks
+# Docs: https://docs.github.com/code-security/dependabot/working-with-dependabot/automating-dependabot-with-github-actions
+
+on: pull_request
+
+permissions:
+  contents: write
+  pull-requests: write
+
+jobs:
+  auto-merge:
+    runs-on: ubuntu-latest
+    if: github.actor == 'dependabot[bot]'
+    steps:
+      - name: Fetch Dependabot metadata
+        id: meta
+        uses: dependabot/fetch-metadata@v3
+        with:
+          github-token: "${{ secrets.GITHUB_TOKEN }}"
+
+      - name: Enable auto-merge for patch/minor updates
+        if: steps.meta.outputs.update-type == 'version-update:semver-patch' || steps.meta.outputs.update-type == 'version-update:semver-minor'
+        run: |
+          gh pr review --approve "$PR_URL"
+          gh pr merge --auto --squash "$PR_URL"
+        env:
+          PR_URL: ${{ github.event.pull_request.html_url }}
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml
index d0a1a6d..0473e17 100644
--- a/.github/workflows/release.yaml
+++ b/.github/workflows/release.yaml
@@ -15,11 +15,11 @@ jobs:
     permissions:
       id-token: write # IMPORTANT: this permission is mandatory for trusted publishing
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           filter: blob:none
           fetch-depth: 0
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.x"
           cache: "pip"
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 8e23d94..2d025ae 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -2,9 +2,9 @@ name: Test
 
 on:
   push:
-    branches: [main]
+    branches: [main, dev-claude]
   pull_request:
-    branches: [main]
+    branches: [main, dev-claude]
   schedule:
     - cron: "0 5 1,15 * *"
 
@@ -41,7 +41,7 @@ jobs:
       PYTHON: ${{ matrix.python }}
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
 
       - name: Free up disk space
         run: |
@@ -52,7 +52,7 @@ jobs:
           sudo rm -rf /usr/local/lib/android
 
       - name: Set up Python ${{ matrix.python }}
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python }}
 
@@ -77,4 +77,4 @@ jobs:
         run: |
           coverage report
       - name: Upload coverage
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v7
diff --git a/.gitignore b/.gitignore
index da806f3..343bb0b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -112,3 +112,4 @@ modules/.calmate/cache/
 cache/
 
 mmcontext_out/
+tutorials/outputs/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 81057ec..3785c34 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -11,11 +11,11 @@ repos:
     hooks:
       - id: prettier
   - repo: https://github.com/tox-dev/pyproject-fmt
-    rev: "v2.5.0"
+    rev: "v2.24.1"
     hooks:
       - id: pyproject-fmt
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.3
+    rev: v0.15.17
     hooks:
       - id: ruff
         types_or: [python, pyi, jupyter]
@@ -23,7 +23,7 @@ repos:
       - id: ruff-format
         types_or: [python, pyi, jupyter]
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: detect-private-key
       - id: check-ast
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 0000000..7dbfc68
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,93 @@
+# CLAUDE.md — mmcontext
+
+## Project Overview
+
+mmcontext is a Python library for multimodal contrastive learning that aligns text and omics (single-cell gene expression) embeddings using the sentence-transformers (>=5.4) framework. It enables joint embedding spaces where text descriptions and biological data can be compared directly.
+
+**Repository**: `github.com/mengerj/mmcontext`
+**Maintainer**: Jonatan Menger
+**Python**: >=3.11, <3.14
+**Key dependency**: `sentence-transformers>=5.4` (multimodal API)
+
+## Package Structure
+
+```
+src/mmcontext/
+├── modules/           # sentence-transformers InputModules
+│   ├── mmcontext_module.py    # Core: text tokenization + omics vector pass-through
+│   ├── adapter_module.py      # Projects omics vectors into shared embedding space
+│   └── omics_attention_module.py  # Optional self-attention for omics tokens
+├── embed/             # Embedding pipeline, dataset utilities, model utils
+├── eval/              # Evaluation: kNN, ARI, scIB metrics, label similarity
+├── io/                # VectorStore, data preparation
+├── pl/                # Plotting utilities
+├── _legacy/           # Archived modules (do not modify unless removing)
+├── callback.py        # Training callbacks
+├── hub_utils.py       # HuggingFace Hub integration
+├── simulator.py       # Synthetic data generation for testing
+└── utils.py           # Shared utilities
+```
+
+## Code Style
+
+- **Formatter/linter**: ruff (line-length=120)
+- **Docstrings**: numpy-style, required for public classes and functions (D100/D104/D105/D107 ignored)
+- **Imports**: sorted by isort (via ruff)
+- **Type hints**: use them for all public function signatures
+- **Pre-commit**: prettier, pyproject-fmt, ruff, detect-private-key, check-ast
+
+## Testing Conventions
+
+- Tests live in `tests/` at repo root
+- Use pytest fixtures; shared fixtures go in `conftest.py`
+- Test files: `test_<module_name>.py`
+- The `simulator.py` module generates synthetic data for tests — prefer it over loading real datasets
+- Tests must pass on Python 3.11, 3.12, 3.13
+
+## PR Conventions
+
+- PRs target `dev-claude` (not `main`) unless it's a release
+- PR description must reference the issue: `Fixes #<number>` or `Closes #<number>`
+- All CI checks must pass before merge
+- Keep PRs focused — one logical change per PR
+
+## Issue Implementation Protocol
+
+When implementing a feature from a GitHub issue (via @claude or otherwise):
+
+1. **If the issue is ambiguous**: Post clarifying questions as a comment. Do NOT start implementation until the questions are answered.
+
+2. **Plan first**: Before writing any code, post an implementation plan as a comment on the issue with a checkbox list:
+
+   ```
+   ## Implementation Plan
+   - [ ] Step 1: description
+   - [ ] Step 2: description
+   - [ ] Step 3: description
+   - [ ] Verify: run tests, check linting
+   ```
+
+   Wait for approval (a reply containing "approved", "go ahead", "LGTM", or "looks good").
+
+3. **Implement**: Create a branch `claude/<issue-number>-<short-slug>` from `dev-claude`. Implement the plan step by step. Edit the plan comment to check off completed steps.
+
+4. **Open PR**: Create a PR targeting `dev-claude` with `Fixes #<number>` in the description. Include a summary of what was done and any decisions made.
+
+## Architecture Notes
+
+The core pipeline follows the sentence-transformers module pattern:
+
+```
+Input → MMContextModule → AdapterModule → [OmicsAttentionModule] → Pooling → Loss
+```
+
+- **MMContextModule**: Handles both text (tokenize → AutoModel) and omics (VectorStore lookup or direct input) modalities. Outputs a unified features dict with `token_embeddings`, `attention_mask`, and `modality_ids`.
+- **AdapterModule**: Projects omics vectors into the text model's embedding space.
+- **OmicsAttentionModule**: Optional self-attention layer for omics tokens.
+- Data stored in anndata format (`.h5ad`), converted via `adata-hf-datasets` for HuggingFace compatibility.
+
+## What NOT to Change
+
+- Do not modify files in `_legacy/` unless explicitly removing them
+- Do not change the sentence-transformers module interface contracts (features dict keys)
+- Do not add dependencies without discussion — the package already has heavy deps (torch, transformers, scanpy)
diff --git a/MIGRATION_PLAN.md b/MIGRATION_PLAN.md
new file mode 100644
index 0000000..d3f9d25
--- /dev/null
+++ b/MIGRATION_PLAN.md
@@ -0,0 +1,282 @@
+# Migration Plan: Split evaluation/benchmarking into `mmcontext-benchmark`
+
+**Status:** in progress — benchmark pipeline implemented (see §0); `mmcontext`
+slim-down (§3–§5) still pending.
+**Author:** Jonatan Menger (with Claude)
+**Date:** 2026-06-03 (updated 2026-06-09)
+**Goal:** Make `mmcontext` a focused model package (architecture + lightweight,
+single-model evaluation runnable from a notebook). Move all multi-model
+comparison, competitor integrations (SCSA, CellWhisperer), and the heavy
+embedding/eval orchestration into a new sibling repo, `mmcontext-benchmark`.
+Drop the `adata-hf-datasets` package dependency from `mmcontext`.
+
+---
+
+## 0. Implemented benchmark design (2026-06-09)
+
+The benchmark repo was built around an **adapter + artifact-contract**
+architecture rather than a straight file-move of the legacy orchestrators. What
+landed (and where it diverged from the original §2–§3 plan):
+
+- **Artifact contract** (`mmcontext_benchmark/artifacts.py`): one on-disk layout
+  per `(dataset, model)` — embedders write `embeddings.parquet` + `subset.h5ad`
+  - `*_label_embeddings_*`; direct classifiers write
+    `predictions_<label>.parquet`; both write `meta.yaml`. Adapters and eval only
+    talk through this, which is what makes different methods comparable.
+- **Adapters + registry** (`adapters/`): one adapter per model _kind_. The
+  `mmcontext` adapter **auto-detects** the old (registered-embeddings) vs new
+  (VectorStore) approach by inspecting the loaded model's first module — so a
+  single config entry works for both. External tools declare an explicit
+  `kind:` and are dispatched to their adapter. This replaces the planned
+  wholesale move of `model_utils`/`embed_pipeline` functions.
+- **External tools = containers** (`containers/`): instead of porting
+  `scsa_utils.py`/`cellwhisperer_utils.py` as in-process modules with their
+  bespoke venv bootstrapping, each competitor ships as its **own image**
+  (SCSA, CellWhisperer, and a separate **calmate** image for label
+  harmonisation) with a self-contained entrypoint. The host adapter prepares
+  inputs, runs the container (Docker locally / Apptainer on clusters), and reads
+  artifacts back. This is the clean replacement for the "unsatisfying" venv +
+  `subprocess` + interpreter-path approach.
+- **Two metric tracks** (`eval/`): embedders → kept-in-`mmcontext`
+  `LabelSimilarity`; classifiers (SCSA) → `eval/classification.py`, emitting the
+  **same** `LabelSimilarity/...` keys so both fuse in `collect/`. CellWhisperer's
+  dot+softmax+`logit_scale` is carried in `meta.yaml`. scIB is dropped (legacy).
+- **Config**: Hydra is retained. `conf/benchmark.yaml` composes `datasets/` and
+  `models/` lists; `scripts/run_benchmark.py` runs embed → eval → collect.
+
+Consequence for this plan: the §3.2 "MOVE" rows are realised as _adapters +
+container entrypoints_ in the benchmark repo, not as 1:1 file moves. The
+`mmcontext`-side slim-down below is unchanged and still to do.
+
+---
+
+## 1. Guiding principles
+
+1. **`mmcontext` = the model.** Modules, adapters, the VectorStore the model
+   depends on, hub/training utilities, and a _slim_ eval surface that a user can
+   run from `tutorials/evaluate_model.ipynb` on a single trained model.
+2. **`mmcontext-benchmark` = the science of comparison.** Everything that loads
+   _other_ models, runs multi-model embedding pipelines, computes comparative
+   metrics, and reproduces paper figures. It depends on `mmcontext` (and on
+   `adata-hf-datasets`), never the reverse.
+3. **No new dependency edges into `mmcontext`.** After the split, nothing in
+   `src/mmcontext` may import `adata_hf_datasets`, `scsa_utils`,
+   `cellwhisperer_utils`, or `eval_pipeline`.
+4. **Tests move with their code.** Each moved module takes its tests (or new
+   ones) to the benchmark repo; kept modules keep/keep-improving their tests.
+
+---
+
+## 2. Key findings from the current tree
+
+These shaped the boundaries below and are worth recording:
+
+- **No `src/` file imports `adata_hf_datasets`.** The only references are in
+  `pyproject.toml`, `egg-info`, docstrings, and _notebooks_. The package dep is
+  effectively a notebook-only dependency today, so removing it from
+  `pyproject.toml` is safe for the library. Users who need to _create_ test
+  datasets install `adata-hf-datasets` themselves (as the kept
+  `evaluate_model.ipynb` already instructs).
+- **The kept eval modules are self-contained.** `query_annotate`,
+  `label_similarity`, `ari`, `embedding_alignment`, `umap_plotter`,
+  `scib_wrapper`, `evaluate_scib`, `base`, `registry`, `utils` only import from
+  within `eval/`, plus `pl/` and `file_utils`. None imports `eval_pipeline.py`
+  or anything in `embed/`. The split is clean.
+- **`eval/eval_pipeline.py` is the multi-model orchestrator** (it special-cases
+  CellWhisperer logit-scale extraction and walks the registry over many models).
+  It is the natural seam — it moves; the registry + individual evaluators stay.
+- **`embed/model_utils.py` mixes two concerns:** single-model encode helpers
+  (`load_st_model`, `encode_adata`, `create_label_dataset`, `embed_labels`,
+  `HFIndexedDataset`) vs. multi-model orchestration (`prepare_model_and_embed`,
+  `prepare_model_and_ds`). It must be split, not moved wholesale.
+- **`evaluate_model.ipynb` imports** `mmcontext.eval` (incl. registry `get`),
+  `eval.query_annotate`, `pl`, `file_utils`, `utils` — all in the keep set. It
+  is the canonical "evaluate one model" notebook and stays.
+- **The two embed-data tutorials** are `tutorials/pretrained_inference.ipynb`
+  and `tutorials/train_new.ipynb`; both import `adata_hf_datasets`. These are
+  the two to remove. `tutorials/evaluate_model copy.ipynb` is a stray duplicate
+  and should also go.
+
+---
+
+## 3. File-by-file disposition
+
+### 3.1 KEEP in `mmcontext`
+
+| Path                                                                              | Notes                                                                                       |
+| --------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- |
+| `src/mmcontext/modules/*`                                                         | Core architecture — untouched.                                                              |
+| `src/mmcontext/io/vector_store.py`                                                | Required by `MMContextModule`.                                                              |
+| `src/mmcontext/io/prepare_store.py`                                               | Keep the VectorStore-from-HF-dataset path (no `adata_hf` import; only a docstring mention). |
+| `src/mmcontext/eval/{base,registry,utils}.py`                                     | Eval framework + label specs.                                                               |
+| `src/mmcontext/eval/query_annotate.py`                                            | `OmicsQueryAnnotator` — used by the notebook; has tests.                                    |
+| `src/mmcontext/eval/{ari,embedding_alignment,label_similarity}.py`                | Single-model metrics.                                                                       |
+| `src/mmcontext/eval/{scib_wrapper,evaluate_scib}.py`                              | scIB metrics for one embedding.                                                             |
+| `src/mmcontext/eval/umap_plotter.py`                                              | Single-model UMAP eval.                                                                     |
+| `src/mmcontext/embed/dataset_utils.py`                                            | `SentenceDataset`, `load_generic_dataset`, `collect_adata_subset` — minimal loading.        |
+| `src/mmcontext/embed/dataset_prep.py`                                             | `prepare_inference`, `InferenceData`, `prepare_dataset` — minimal data prep.                |
+| **NEW** `src/mmcontext/embed/encode.py`                                           | Single-model encode helpers split out of `model_utils.py` (see 3.3).                        |
+| `src/mmcontext/{callback,hub_utils,simulator,utils,file_utils,sanity_helpers}.py` | Core/training/util.                                                                         |
+| `src/mmcontext/pl/*`                                                              | Plotting (used by kept eval + notebook).                                                    |
+| `tutorials/evaluate_model.ipynb`                                                  | Canonical single-model evaluation notebook.                                                 |
+
+### 3.2 MOVE to `mmcontext-benchmark`
+
+| Path                                                                                                 | New home (benchmark)                          | Reason                                          |
+| ---------------------------------------------------------------------------------------------------- | --------------------------------------------- | ----------------------------------------------- |
+| `src/mmcontext/eval/eval_pipeline.py`                                                                | `mmcontext_benchmark/eval/pipeline.py`        | Multi-model comparison orchestrator.            |
+| `src/mmcontext/embed/embed_pipeline.py`                                                              | `mmcontext_benchmark/embed/pipeline.py`       | Multi-model embedding pipeline.                 |
+| `src/mmcontext/embed/scsa_utils.py`                                                                  | `mmcontext_benchmark/models/scsa.py`          | Competitor model (SCSA).                        |
+| `src/mmcontext/embed/cellwhisperer_utils.py`                                                         | `mmcontext_benchmark/models/cellwhisperer.py` | Competitor model (CellWhisperer).               |
+| `model_utils.py` → `prepare_model_and_embed`, `prepare_model_and_ds`, `process_single_dataset_model` | `mmcontext_benchmark/embed/orchestrate.py`    | Multi-model orchestration split.                |
+| `tutorials/pretrained_inference.ipynb`, `tutorials/train_new.ipynb`                                  | `mmcontext-benchmark/examples/` (or delete)   | Embed-data tutorials using `adata_hf_datasets`. |
+
+### 3.3 SPLIT `embed/model_utils.py`
+
+Stays in `mmcontext` as `embed/encode.py` (single-model):
+`load_st_model`, `encode_adata`, `create_label_dataset`, `embed_labels`,
+`HFIndexedDataset`, `SentenceDataset`.
+
+Moves to benchmark `embed/orchestrate.py` (multi-model):
+`prepare_model_and_embed`, `prepare_model_and_ds`.
+
+Update `src/mmcontext/embed/__init__.py` to export only kept symbols; remove
+`process_cellwhisperer_dataset_model`, `ensure_cellwhisperer_setup`,
+`embed_pipeline`, `process_single_dataset_model`, `prepare_model_and_embed`.
+
+### 3.4 REMOVE
+
+- `tutorials/pretrained_inference.ipynb`, `tutorials/train_new.ipynb`
+  (after their useful content is ported to benchmark `examples/`).
+- `tutorials/evaluate_model copy.ipynb` (stray duplicate).
+- `adata-hf-datasets` from `pyproject.toml` `dependencies` and the
+  `[tool.uv] sources.adata-hf-datasets` block.
+
+---
+
+## 4. Dependency changes (`pyproject.toml`)
+
+Remove from `mmcontext`:
+
+```toml
+# dependencies
+- "adata-hf-datasets @ git+https://github.com/mengerj/adata_hf_datasets.git",
+
+# [tool.uv]
+- sources.adata-hf-datasets = { git = "https://github.com/mengerj/adata_hf_datasets.git" }
+```
+
+**Training stays in `mmcontext`.** The package keeps its training code
+(`callback.py`, training utilities, hub upload). `wandb` is used for training
+logging, so it is **not removed** — instead it becomes an _optional_ extra so
+inference-only users get a lighter install:
+
+```toml
+# move wandb out of core dependencies into:
+optional-dependencies.train = [ "wandb>=0.27" ]
+# (alternatively a broader `train` extra bundling hydra-core, wandb, accelerate)
+```
+
+Training code must import `wandb` lazily (inside the training/callback path,
+guarded by a clear `ImportError` message pointing to `pip install
+mmcontext[train]`) so that `import mmcontext` and inference never require it.
+
+Candidates to also drop from `mmcontext` _once their only users have moved_
+(verify with a final grep — do not remove blindly): `scib`, `louvain` (if only
+used by scIB multi-model paths). Keep anything still imported by the kept
+eval/model code. `scib`/`louvain` are used by the kept
+`scib_wrapper`/`evaluate_scib`, so they **stay** unless those move too.
+`hydra-core` stays if training configs use it.
+
+`mmcontext-benchmark` `pyproject.toml` depends on:
+`mmcontext @ git+https://github.com/mengerj/mmcontext.git`,
+`adata-hf-datasets @ git+...`, plus the SCSA/CellWhisperer runtime deps.
+
+---
+
+## 5. Phased execution
+
+Each phase is a separate PR targeting `dev-claude`, one logical change each, CI
+green before merge (per `CLAUDE.md`).
+
+**Phase 0 — Scaffold benchmark repo (no `mmcontext` change).**
+Create `mmcontext-benchmark` skeleton (this plan ships a starter skeleton).
+Wire CI, pre-commit, pytest. Add `mmcontext` as a git dependency.
+
+**Phase 1 — Copy competitor + orchestration code into benchmark.**
+Port `scsa_utils`, `cellwhisperer_utils`, `embed_pipeline`, `eval_pipeline`,
+multi-model `model_utils` functions, and the two embed-data tutorials into the
+benchmark repo. Fix imports to `mmcontext.*` for shared helpers (`eval.utils`,
+`file_utils`, `pl`). Get benchmark tests passing against installed `mmcontext`.
+_No deletion in `mmcontext` yet — both repos temporarily overlap._
+
+**Phase 2 — Split `model_utils.py` in `mmcontext`.**
+Branch `claude/<issue>-split-model-utils`. Create `embed/encode.py` with the
+single-model helpers; update `embed/__init__.py`; update any internal imports.
+Run tests.
+
+**Phase 3 — Remove moved code from `mmcontext`.**
+Delete `eval/eval_pipeline.py`, `embed/embed_pipeline.py`, `embed/scsa_utils.py`,
+`embed/cellwhisperer_utils.py`, and the multi-model functions. Update
+`__init__.py` exports. Remove the two tutorials + the stray copy.
+
+**Phase 4 — Drop the dependency.**
+Remove `adata-hf-datasets` from `pyproject.toml` and `[tool.uv]`. Re-lock.
+Grep to confirm zero `adata_hf` / `scsa` / `cellwhisperer` references remain in
+`src/`. Update README references.
+
+**Phase 5 — Docs + notebook.**
+Ensure `evaluate_model.ipynb` runs end-to-end against the slimmed package.
+Add a README note pointing benchmarking/comparison users to
+`mmcontext-benchmark`. Update `CLAUDE.md` package-structure section.
+
+---
+
+## 6. Test strategy
+
+- **Kept code:** keep existing tests
+  (`test_query_annotate.py`, module/adapter/encoder tests, `test_io`,
+  `test_prepare_store.py`). Add a regression test asserting the slim public API:
+  `import mmcontext; mmcontext.eval.get(...)` works and the moved symbols are
+  _gone_ (`pytest.raises(ImportError)` / `AttributeError`).
+- **Import-isolation test:** a test that imports every `src/mmcontext` module and
+  asserts `adata_hf_datasets` was never imported (check `sys.modules`). This
+  locks in the dependency removal.
+- **Moved code:** create `tests/` in the benchmark repo. Reuse
+  `mmcontext.simulator` for synthetic fixtures (per project convention — prefer
+  the simulator over real datasets). Run on Python 3.11–3.13.
+- **Sandbox note:** `torch` is too large to install in the Cowork sandbox, so
+  the full test suite must run on Jo's machine / CI, not here.
+
+---
+
+## 7. Risks & mitigations
+
+| Risk                                                     | Mitigation                                                                                                                             |
+| -------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- |
+| Hidden runtime coupling not caught by static grep        | Phase 1 keeps both copies until benchmark tests pass; only Phase 3 deletes.                                                            |
+| `eval.utils.LabelKind` / `file_utils` used by moved code | These stay in `mmcontext`; benchmark imports them — intended one-way edge.                                                             |
+| Dropping `scib`/`louvain` breaks kept scIB eval          | Keep them; only drop deps whose _only_ importer moved (verify by grep).                                                                |
+| Notebook drift (Jo is mid-update)                        | Phase 5 is last; coordinate with the in-progress `evaluate_model.ipynb` rewrite.                                                       |
+| Users relying on old `mmcontext.embed.embed_pipeline`    | Add a deprecation shim that raises `ImportError` with a message pointing to `mmcontext-benchmark` (one release), or note in CHANGELOG. |
+
+---
+
+## 8. Rollback
+
+Each phase is an isolated PR. If a phase regresses, revert that PR; earlier
+phases (scaffold, copies) are additive and safe to leave in place. The
+benchmark repo is independent, so reverting `mmcontext` changes never breaks it
+beyond a version pin bump.
+
+---
+
+## 9. Definition of done
+
+- `grep -rn "adata_hf\|scsa\|cellwhisperer" src/mmcontext --include="*.py"`
+  returns nothing (docstrings cleaned too).
+- `pip install mmcontext` pulls no `adata-hf-datasets`.
+- `tutorials/evaluate_model.ipynb` runs against the slim package.
+- `mmcontext-benchmark` reproduces the multi-model comparison previously done by
+  `eval_pipeline.py`, with its own passing tests.
+- `CLAUDE.md` and `README` reflect the new boundary.
diff --git a/README.md b/README.md
index 4a4e8bc..f29d4ea 100644
--- a/README.md
+++ b/README.md
@@ -2,9 +2,9 @@
 
 [![Tests][badge-tests]][tests]
 
-Align embeddings across multiple modalities using context-aware embeddings at the sample level.
+Multimodal contrastive learning aligning text and omics embeddings via [sentence-transformers](https://www.sbert.net/).
 
-**mmcontext** is built upon the excellent [sentence-transformers](https://www.sbert.net/) framework maintained by [Hugging Face](https://huggingface.co/). By leveraging their comprehensive documentation and extensive capabilities for text embeddings, **mmcontext** enables you to efficiently generate multi-modal embeddings without reinventing the wheel.
+**mmcontext** enables joint embedding spaces where natural-language descriptions and biological data (single-cell gene expression) can be compared directly. Built on the sentence-transformers v5.4+ multimodal API.
 
 ![Conceptual Diagram](figs/concept.png)
 
@@ -14,210 +14,103 @@ Align embeddings across multiple modalities using context-aware embeddings at th
 Menger, Jonatan, Sonia Maria Krissmer, Clemens Kreutz, Harald Binder, and Maren Hackenberg. "mmContext: an open framework for multimodal contrastive learning of omics and text data." bioRxiv (2025): 2025-12.
 If you use **mmcontext** in your research, please cite our work.
 
-## Overview
-
-**mmcontext** provides tools for multi-modal embedding generation, with different workflows depending on your needs:
-
-- **[Using Pre-trained Models for Inference](#using-pre-trained-models-for-inference)** - Load and use pre-trained mmcontext models
-- **[Training New Models](#training-new-models)** - Train custom mmcontext models on your datasets
-- **[Reproducing Paper Results](#reproducing-paper-results)** - Reproduce experiments from the paper
-
-## Using Pre-trained Models for Inference
-
-You can use pre-trained models (currently limited to RNA-seq) to embed your own data into the joint latent space of an mmcontext model. You can then query the dataset with natural language, for annotation etc.
-
-### Installation
-
-For inference, you only need to install the package:
+## Installation
 
 ```bash
-pip install git+https://github.com/mengerj/mmcontext.git@main
+pip install mmcontext
 ```
 
-### Pre-trained Models
-
-Pre-trained mmcontext models are available on Hugging Face under the [jo-mengr](https://huggingface.co/jo-mengr) organization. Browse available models at: **https://huggingface.co/jo-mengr**
-
-### Tutorial
-
-See the [Using a Pre-trained Model for Inference](tutorials/pretrained_inference.ipynb) tutorial for a complete guide on:
-
-- Loading pre-trained models from Hugging Face
-- Processing your adata to be suitable
-- Generating embeddings for text and omics data
-- Using models for downstream tasks
-
-## Training New Models
-
-Train your own mmcontext models on custom datasets.
-
-### Tutorial
-
-See the [Training a New Model](tutorials/train_new.ipynb) tutorial for:
-
-- Step-by-step training guide
-- Configuration and hyperparameters
-- Best practices for multi-modal training
-
-### Dataset Preparation
-
-If you want to utilize the full pipeline used to process the training datasets, go to
-[adata-hf-datasets](https://github.com/mengerj/adata_hf_datasets).
-This package handles the conversion of AnnData objects to Hugging Face datasets with proper formatting for multimodal training.
-
-## Reproducing Paper Results
-
-To reproduce results from the paper, clone the repository and install in editable mode:
+Or install from source with development dependencies:
 
 ```bash
 git clone https://github.com/mengerj/mmcontext.git
 cd mmcontext
-
-# Create a virtual environment (however you like)
-# eg. with venv
-python -m venv mmcontext
-source mmcontext/bin/activate
-
-# And install the package with pip
-pip install -e .
-```
-
-### Training Models
-
-To train a model as done for the paper, use the [`scripts/train.py`](scripts/train.py) script.
-Try a small training run with
-
-```bash
-python scripts/train.py --config-name example_conf
+pip install -e ".[dev,test,eval,train]"
 ```
 
-Inspect the config at [`example_conf.yaml`](conf/training/example_conf.yaml). The configs to train the models presented in the paper, are [`basebert_numeric`](conf/training/basebert_numeric.yaml), for all models using numeric initial representations and [`basebert_text.yaml`](conf/training/basebert_text.yaml) for the model using cell_sentences (text only). All datasets used in training are hosted publically on the huggingface hub (with references to zenodo), therefore the training scripts can be launched without manually downloading any data.
-
-For HPC systems with CUDA support, the recommended approach is to use the [`scripts/run_training.slurm`](scripts/run_training.slurm) SLURM script to launch training jobs. Training also works on CPU or MPS devices if CUDA is not available.
-
-The SLURM script allows you to override configuration values from the command line, which is useful when launching several jobs with different configurations.
-
-```bash
-sbatch scripts/run_training.slurm
-```
+## Quick Start
 
-Before training, it is recommended to authenticate with Hugging Face (after activating the virtual environment used for mmcontext installation):
+### Using a Pre-trained Model
 
-```bash
-source .venv/bin/activate  # or your venv activation command
-hf auth login
-```
+```python
+from sentence_transformers import SentenceTransformer
 
-If you have a Weights & Biases (wandb) account, you can also log in so your training is tracked. From the command line, run:
+model = SentenceTransformer("jo-mengr/mmcontext-pubmedbert-gs10k", trust_remote_code=True)
 
-```bash
-wandb login
+# Encode text queries
+text_embeddings = model.encode(["CD4+ T cell", "B cell precursor"])
 ```
 
-This will prompt you to enter your wandb API key.
-
-Once training is complete, the finished models will be automatically uploaded to the Hugging Face Hub with metadata and model cards.
-
-### Evaluating Models
-
-Figure 1D of the paper investigates the latent space of one model in detail. This can be recreated with the [`evaluate_model.ipynb`](tutorials/evaluate_model.ipynb) notebook.
+Pre-trained models are available on Hugging Face under the [jo-mengr](https://huggingface.co/jo-mengr) organization.
 
-For Figure 1E, we evaluate several models on multiple datasets, with the [`scripts/embed_eval.py`](scripts/embed_eval.py) script. This runs both inference and evaluation pipelines in sequence.
+See the [pretrained inference tutorial](tutorials/pretrained_inference.ipynb) for a complete guide.
 
-The combined pipeline is configured using [`embed_eval_conf.yaml`](conf/eval/embed_eval_conf.yaml), which inherits from dataset and model configuration files that list the datasets and models to be evaluated. The configuration file contains additional parameters that are explained in the comments within the file itself.
-The models and datasets evaluated in the paper are referenced in [model_list_cxg_geo_all.yaml](conf/models/model_list_cxg_geo_all.yaml) and [dataset_list.yaml](conf/datasets/dataset_list.yaml). These configs are imported in [embed_eval_conf.yaml](conf/eval/embed_eval_conf.yaml). To jointly embed data and evaluate with CellWhisperer, set `run_cellwhisperer: true`. It is highly recommended to use CUDA for CellWhisperer. The mmcontext models also run in reasonable time on MPS or CPU.
+### Building a Pipeline
 
-Run it locally with
+```python
+from mmcontext.embed import build_pipeline
 
-```bash
-python scripts/embed_eval.py
+pipeline = build_pipeline(
+    text_model="microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract",
+    omics_dim=512,        # dimension of your omics vectors
+    shared_dim=256,       # joint embedding space dimension
+)
 ```
 
-For HPC systems, you can run the combined pipeline as array jobs using [`scripts/run_combined_cpu.slurm`](scripts/run_embed_eval_cpu.slurm):
-
-```bash
-sbatch scripts/run_embed_eval_cpu.slurm
-```
-
-This allows you to process multiple model configurations in parallel across different array job tasks, by spreading the models across several config files and passing them as a list to the array job.
-
-### SCSA Baseline
-
-[SCSA](https://github.com/bioinfo-ibms-pumc/SCSA) (Single-Cell RNA-seq Annotation) is included as a non-embedding baseline. Unlike the mmcontext models, SCSA annotates cell types by clustering cells, detecting marker genes per cluster, and matching them against the CellMarker database using Z-scores. It runs as a standalone script on the same datasets.
+### Evaluation
 
-SCSA requires `numpy<2.0` and therefore needs its own virtual environment:
+```python
+from mmcontext.eval import get
 
-```bash
-# 1. Clone the SCSA repository
-bash modules/prepare_scsa.sh
-
-# 2. Create a dedicated venv for SCSA and install its dependencies
-python -m venv modules/scsa_venv
-modules/scsa_venv/bin/pip install -r modules/requirements_scsa.txt
-
-# 3. (Optional) Create a dedicated venv for calmate label harmonisation
-#    This maps SCSA's predicted labels and ground-truth labels to Cell Ontology
-#    terms so that metrics are computed on normalised names.
-python -m venv modules/calmate_venv
-modules/calmate_venv/bin/pip install -r modules/requirements_calmate.txt
-
-# 4. Run SCSA on all datasets (venvs are auto-detected from modules/)
-python scripts/run_scsa.py
-
-# 5. (Optional) After reviewing calmate mappings, you can recompute\n#    metrics without re-running SCSA itself by using:\n+#    python scripts/run_scsa.py scsa.remap_only=true settings.skip_existing=false
+LabelSimilarity = get("LabelSimilarity")
+evaluator = LabelSimilarity(similarity="cosine")
+result = evaluator.compute(
+    omics_embeddings=omics_emb,
+    label_embeddings=label_emb,
+    query_labels=labels,
+    true_labels=true_labels,
+    label_key="cell_type",
+)
 ```
 
-When calmate is available, the pipeline runs label harmonisation automatically. On the first run, many labels will need manual review. The script prints instructions at the end — run `modules/calmate_venv/bin/calmate --store modules/.calmate/mappings.csv review` to interactively approve the suggested ontology mappings, then re-run with `settings.skip_existing=false`. If you only want to redo the mapping/metrics step and reuse existing SCSA predictions, set `scsa.remap_only=true`; in that mode `settings.skip_existing` is ignored for SCSA and only controls whether other datasets are skipped.
+See the [evaluate model 2.0 notebook](tutorials/evaluate_model_2.0.ipynb) for a complete evaluation workflow.
 
-SCSA is configured via [`scsa_conf.yaml`](conf/eval/scsa_conf.yaml), which shares the same dataset list as the embedding pipeline. Key parameters (fold-change threshold, p-value, clustering resolution, species) can be overridden from the command line. If the SCSA venv is in a non-default location, pass it explicitly: `python scripts/run_scsa.py scsa.scsa_python=/path/to/scsa_venv/bin/python`. Results are written to the same output directory tree as the other models, so `collect_metrics.py` picks them up automatically.
+## Architecture
 
-For HPC systems:
+mmcontext 2.0 uses the sentence-transformers module pipeline pattern:
 
-```bash
-sbatch scripts/run_scsa.slurm
+```
+Input -> MMContextModule -> AdapterModule -> Pooling -> Normalize
 ```
 
-Figure Panel 1E is created by collecting the metrics from the big evaluation run. The config [`collect_metrics_conf.yaml`](conf/eval/collect_metrics_conf.yaml) has to point to the directory where the results from all datasets/models were stored. Then to collect and plot the metrics, run
+The **MMContextModule** handles both text (tokenize -> AutoModel) and omics (VectorStore lookup) modalities. The **AdapterModule** projects omics vectors into the text model's embedding space.
 
-```bash
-python scripts/collect_metrics.py
-python scripts/plot_metrics.py
-```
+Data is stored in AnnData format (`.h5ad`/zarr), with a memory-efficient VectorStore for runtime omics lookups.
 
-## Core Components
+## Multi-Model Benchmarking
 
-The main model implementation is the **MMContextEncoder**, located in [`src/mmcontext/mmcontextencoder.py`](src/mmcontext/mmcontextencoder.py). This dual-tower encoder can process both text and omics data, enabling multi-modal embedding generation.
+For systematic comparisons across multiple models and datasets, see the companion repository: [mmcontext-benchmark](https://github.com/mengerj/mmcontext-benchmark).
 
-## Contributing
+## Dataset Preparation
 
-This package is under active development. Contributions and suggestions are very welcome! Please open an [issue](https://github.com/mengerj/mmcontext/issues) to:
+To prepare training datasets from AnnData objects, see [adata-hf-datasets](https://github.com/mengerj/adata_hf_datasets).
 
-- Propose enhancements
-- Report bugs
-- Discuss potential improvements
-- Ask questions or seek help
+## Contributing
 
-We encourage community contributions and are happy to help you get started.
+Contributions and suggestions are very welcome! Please open an [issue](https://github.com/mengerj/mmcontext/issues) to propose enhancements, report bugs, or ask questions.
 
 ## Citation
 
-If you find **mmcontext** useful for your research, please consider citing our paper (citation to be added upon publication):
-
 ```bibtex
-@misc{mmcontext,
-  author = {Jonatan Menger},
-  title = {mmcontext: Multi-modal Contextual Embeddings},
+@article{menger2025mmcontext,
+  author = {Menger, Jonatan and Krissmer, Sonia Maria and Kreutz, Clemens and Binder, Harald and Hackenberg, Maren},
+  title = {mmContext: an open framework for multimodal contrastive learning of omics and text data},
+  journal = {bioRxiv},
   year = {2025},
-  publisher = {GitHub},
-  journal = {GitHub Repository},
-  url = {https://github.com/mengerj/mmcontext}
+  doi = {10.1101/2025.12}
 }
 ```
 
-## Contact
-
-Encountered a bug or need help? Please use the [issue tracker](https://github.com/mengerj/mmcontext/issues). We are happy for any feedback.
-
 ---
 
 [badge-tests]: https://img.shields.io/github/actions/workflow/status/mengerj/mmcontext/test.yaml?branch=main
diff --git a/REFACTOR_ROADMAP.md b/REFACTOR_ROADMAP.md
new file mode 100644
index 0000000..3e0cd8f
--- /dev/null
+++ b/REFACTOR_ROADMAP.md
@@ -0,0 +1,289 @@
+# mmcontext Refactor Roadmap — ST v5.4 Alignment
+
+**Branch:** `dev-claude`
+**Goal:** Refactor mmcontext to align with the sentence-transformers v5.4 multimodal API while preserving all core functionality. Test-driven development throughout.
+
+## Architecture Overview
+
+### Current → New
+
+| Aspect            | Current                  | Refactored                                             |
+| ----------------- | ------------------------ | ------------------------------------------------------ |
+| Base class        | `Module`                 | `InputModule`                                          |
+| Tokenization      | `tokenize()`             | `preprocess()`                                         |
+| Omics storage     | `nn.Embedding` lookup    | `VectorStore` (memory-mapped)                          |
+| Omics key         | `pixel_values`           | `omics_values`                                         |
+| Adapters          | Inside encoder           | Separate `AdapterModule(Module)`                       |
+| Modality flag     | `omics_text_info`        | `modality_ids`                                         |
+| OneHotTextEncoder | Included                 | Removed                                                |
+| Data loading      | Coupled to encoder       | `mmcontext.io` module                                  |
+| Pipeline          | `[MMContextEncoder]`     | `[MMContextModule, AdapterModule, Pooling, Normalize]` |
+| Var support       | Same class, no attention | Optional `OmicsAttentionModule`                        |
+
+### Module Pipeline
+
+```
+SentenceTransformer(modules=[
+    MMContextModule,          # InputModule: text encoder + omics pass-through
+    OmicsAttentionModule,     # Module (OPTIONAL, var-only): self-attention on omics tokens
+    AdapterModule,            # Module: modality-aware projection to shared space
+    Pooling,                  # ST built-in: mean/cls/max pooling
+    Normalize,                # ST built-in: L2 normalization
+])
+```
+
+### Features Dict Contract
+
+All modules communicate through a features dict. After `MMContextModule.forward()`:
+
+```python
+{
+    "token_embeddings": Tensor(B, L, D_encoder),  # unified text + omics tokens
+    "attention_mask": Tensor(B, L),                # 1 = real, 0 = pad
+    "modality_ids": Tensor(B, L),                  # 0 = text, 1 = omics, 2 = pad
+    "sentence_embedding": Tensor(B, D)             # set after Pooling
+}
+```
+
+---
+
+## Phases
+
+### Phase 1: Foundation — VectorStore + IO Module
+
+**Files created/modified:**
+
+- `src/mmcontext/io/__init__.py` (new)
+- `src/mmcontext/io/vector_store.py` (new)
+- `src/mmcontext/io/adata_utils.py` (new — extracted from mmcontextencoder.py + file_utils.py)
+- `tests/test_vector_store.py` (new)
+
+**VectorStore responsibilities:**
+
+- Create from AnnData (obsm/varm), DataFrame, dict, or numpy array
+- Write embeddings to numpy memmap file + JSON index
+- Batch lookup by sample IDs → numpy array
+- Report dim, dtype, len
+- Support both obs (1 vector/sample) and var (N vectors/sample)
+
+**Tests (write FIRST):**
+
+1. `test_from_numpy` — round-trip: write memmap, read back, values match
+2. `test_from_adata_obs` — create from adata.obsm, lookup by obs index
+3. `test_from_adata_var` — create from adata.varm, lookup by var index
+4. `test_batch_lookup` — batch of IDs returns correct (N, D) array
+5. `test_unknown_id_raises` — KeyError for missing IDs
+6. `test_dim_property` — reports correct embedding dimension
+7. `test_memory_efficiency` — creating store doesn't load full matrix into RAM (check process RSS)
+8. `test_persistence` — store survives close/reopen cycle
+
+**adata_utils responsibilities:**
+
+- `load_embeddings_from_adata_link()` — extracted from `get_initial_embeddings_from_adata_link`
+- `create_token_dataframe_from_obsm()` — extracted from encoder
+- `build_embedding_df()` — extracted from file_utils
+
+---
+
+### Phase 2: Core Module — MMContextModule (InputModule)
+
+**Files created/modified:**
+
+- `src/mmcontext/modules/__init__.py` (new)
+- `src/mmcontext/modules/mmcontext_module.py` (new)
+- `tests/test_mmcontext_module.py` (new)
+
+**MMContextModule responsibilities:**
+
+- Extends `InputModule` from sentence-transformers
+- `modalities` property returns `["text", "omics"]`
+- `preprocess(inputs)` routes by input type:
+  - `str` → text tokenization via AutoTokenizer
+  - `dict` with `"omics_values"` key → omics vector packaging
+  - `str` starting with prefix → VectorStore lookup (if store attached)
+- `forward(features)` processes text through AutoModel, passes omics through as-is
+- `set_vector_store(store)` / `remove_vector_store()` for attaching data
+- `save()` / `load()` — saves config + text encoder weights (NOT omics data)
+- Handles mixed batches (text + omics in same batch)
+- `max_seq_length` property for ST compatibility
+- Text encoder freezing/unfreezing logic
+
+**Tests (write FIRST):**
+
+1. `test_preprocess_text_only` — text strings → input_ids, attention_mask, modality_ids
+2. `test_preprocess_omics_direct_vector` — raw vectors → omics_values, attention_mask, modality_ids
+3. `test_preprocess_omics_via_store` — prefixed IDs + VectorStore → resolved vectors
+4. `test_preprocess_mixed_batch` — text + omics in same batch → unified features
+5. `test_preprocess_var_multiple_vectors` — list of gene vectors → variable-length tokens
+6. `test_forward_text_only` — produces token_embeddings with correct shape
+7. `test_forward_omics_only` — omics vectors pass through to token_embeddings
+8. `test_forward_mixed_batch` — mixed batch produces unified token_embeddings + modality_ids
+9. `test_forward_preserves_omics_values` — omics vectors appear unchanged in token_embeddings
+10. `test_modalities_property` — returns ["text", "omics"]
+11. `test_save_load_roundtrip` — config + weights survive save/load
+12. `test_save_excludes_vector_store` — VectorStore data not in saved weights
+13. `test_text_encoder_freezing` — freeze/unfreeze text encoder parameters
+14. `test_max_seq_length` — property accessible and correct
+15. `test_no_store_omics_id_raises` — prefixed IDs without VectorStore → clear error
+
+---
+
+### Phase 3: Modality-Aware AdapterModule
+
+**Files created/modified:**
+
+- `src/mmcontext/modules/adapter_module.py` (new — replaces adapters.py as pipeline Module)
+- `tests/test_adapter_module.py` (new)
+
+**AdapterModule responsibilities:**
+
+- Extends `Module` from sentence-transformers
+- Reads `modality_ids` from features dict
+- Maintains separate projection weights: `text_proj` and `omics_proj`
+- Each projection: Linear → ReLU → Linear → BatchNorm (configurable)
+- Maps `D_text → D_shared` and `D_omics → D_shared`
+- Pad tokens (modality_id=2) pass through as zeros
+- `save()` / `load()` with config_keys for all dimensions
+- `get_sentence_embedding_dimension()` returns D_shared
+
+**Tests (write FIRST):**
+
+1. `test_forward_text_only` — text tokens projected correctly
+2. `test_forward_omics_only` — omics tokens projected correctly
+3. `test_forward_mixed_batch` — text and omics tokens get different projections
+4. `test_pad_tokens_stay_zero` — modality_id=2 tokens remain zero after projection
+5. `test_output_dimension` — output matches D_shared regardless of input modality
+6. `test_separate_weights` — text_proj and omics_proj have independent parameters
+7. `test_gradient_flow` — gradients flow through both projections
+8. `test_weights_update` — optimizer step changes both projection weights
+9. `test_save_load_roundtrip` — config + weights survive save/load
+10. `test_get_sentence_embedding_dimension` — returns D_shared
+
+---
+
+### Phase 4: OmicsAttentionModule (Optional, Var-only)
+
+**Files created/modified:**
+
+- `src/mmcontext/modules/omics_attention_module.py` (new)
+- `tests/test_omics_attention_module.py` (new)
+
+**OmicsAttentionModule responsibilities:**
+
+- Extends `Module` from sentence-transformers
+- Reads `modality_ids` from features dict
+- Applies multi-head self-attention ONLY to omics tokens (modality_id=1)
+- Text tokens (modality_id=0) pass through unchanged
+- Configurable: num_layers, num_heads, hidden_dim, dropout
+- Respects attention_mask for variable-length gene sequences
+- `save()` / `load()` for persistence
+
+**Tests (write FIRST):**
+
+1. `test_text_passthrough` — text tokens unchanged after module
+2. `test_omics_transformed` — omics tokens are modified by self-attention
+3. `test_attention_mask_respected` — padded positions don't influence real tokens
+4. `test_variable_length_sequences` — different-length gene sequences in same batch
+5. `test_output_shape_preserved` — (B, L, D) shape unchanged
+6. `test_gradient_flow` — gradients flow through attention layers
+7. `test_save_load_roundtrip` — config + weights survive save/load
+8. `test_single_token_sequence` — obs-like input (L=1) works without error
+
+---
+
+### Phase 5: SentenceTransformer Integration
+
+**Files created/modified:**
+
+- `tests/test_st_integration.py` (new — replaces test_sentence_transformer_integration.py)
+- Minor adjustments to modules for compatibility
+
+**Integration tests (write FIRST):**
+
+1. `test_pipeline_construction` — modules compose into SentenceTransformer
+2. `test_encode_text` — `model.encode(["text"])` produces correct shape
+3. `test_encode_omics_direct` — `model.encode([{"omics_values": vector}])` works
+4. `test_encode_omics_via_store` — `model.encode(["sample_idx:S1"])` with VectorStore
+5. `test_encode_mixed` — mixed text + omics batch produces unified embeddings
+6. `test_encode_normalize` — output vectors are L2-normalized
+7. `test_save_load_full_pipeline` — save + load produces identical encode() results
+8. `test_modules_json_structure` — saved modules.json has correct module chain
+9. `test_load_with_trust_remote_code` — model loadable via SentenceTransformer(..., trust_remote_code=True)
+10. `test_training_text_only` — SentenceTransformerTrainer runs with text data
+11. `test_training_bimodal` — SentenceTransformerTrainer runs with mixed data
+12. `test_precision_conversion` — fp32 → fp16 works across all modules
+13. `test_max_seq_length_propagation` — max_seq_length accessible from SentenceTransformer
+14. `test_obs_pipeline` — full obs pipeline: [MMContext, Adapter, Pooling, Normalize]
+15. `test_var_pipeline` — full var pipeline: [MMContext, OmicsAttn, Adapter, Pooling, Normalize]
+
+---
+
+### Phase 6: Documentation + Cleanup
+
+**Files modified:**
+
+- `src/mmcontext/__init__.py` — update public API exports
+- `src/mmcontext/modules/__init__.py` — export all modules
+- `src/mmcontext/io/__init__.py` — export VectorStore and utilities
+- Module docstrings — comprehensive docstrings with usage examples
+- `README.md` — update architecture description and usage examples
+
+**Cleanup:**
+
+- Remove `OneHotTextEncoder` (`onehot.py`)
+- Archive old `mmcontextencoder.py` (keep temporarily for reference, don't import)
+- Archive old `adapters.py` (replaced by modules/adapter_module.py)
+- Remove `test_adapter_callback.py` (empty)
+- Update `pyproject.toml` if new dependencies needed
+
+---
+
+## Implementation Order & Dependencies
+
+```
+Phase 1 (VectorStore)
+  ↓
+Phase 2 (MMContextModule) ← depends on VectorStore
+  ↓
+Phase 3 (AdapterModule)   ← depends on features dict contract from Phase 2
+  ↓
+Phase 4 (OmicsAttention)  ← depends on features dict contract from Phase 2
+  ↓
+Phase 5 (ST Integration)  ← depends on all modules
+  ↓
+Phase 6 (Docs + Cleanup)  ← final polish
+```
+
+## Test Strategy
+
+Each phase follows strict TDD:
+
+1. Write test file with all tests (they will fail)
+2. Implement the module until all tests pass
+3. Run full test suite to check for regressions
+
+**Stub strategy:** Tests use lightweight stubs (similar to existing `_TokStub`, `_TextEncStub`) to avoid downloading real models. The existing `conftest.py` pattern is extended for new modules.
+
+**What's preserved from current tests:**
+
+- Core encoding shapes and correctness (text, omics, mixed)
+- Save/load round-trips
+- Gradient flow through adapters
+- ST integration (encode, train, normalize)
+- Attention mask alignment
+- Freezing/unfreezing behavior
+
+**What's new:**
+
+- VectorStore tests (memmap, lookup, persistence)
+- Modality-aware adapter tests (separate projections)
+- OmicsAttentionModule tests (self-attention on omics only)
+- Direct vector input tests (no lookup table)
+- Pipeline composition tests (modules.json)
+
+## Open Items (Deferred)
+
+- Training script rewrite (Phase 7, separate effort)
+- Backward compatibility loading of old models (evaluate complexity later)
+- Cross-attention in OmicsAttentionModule (future extension)
+- Evaluation pipeline updates (depends on new model API)
diff --git a/ROADMAP.md b/ROADMAP.md
new file mode 100644
index 0000000..2669ca2
--- /dev/null
+++ b/ROADMAP.md
@@ -0,0 +1,342 @@
+# MMContext Refactor — Roadmap & Technical Notes
+
+This document captures the state of the codebase after the sentence-transformers
+v5.4+ refactor (Phases 1–5), identifies remaining work, and provides context for
+future development. It is structured as a reference for both human contributors
+and AI-assisted editing sessions.
+
+---
+
+## 1 Architecture overview (post-refactor)
+
+The new pipeline lives in two packages:
+
+| Package             | Role                                                                                          |
+| ------------------- | --------------------------------------------------------------------------------------------- |
+| `mmcontext.modules` | ST pipeline modules: `MMContextModule` (InputModule), `AdapterModule`, `OmicsAttentionModule` |
+| `mmcontext.io`      | `VectorStore` (mmap-backed lookup), `prepare_vector_store` (zarr → store builder)             |
+
+A trained model is a standard `SentenceTransformer` with this module chain:
+
+```
+MMContextModule → [OmicsAttentionModule] → AdapterModule → Pooling → Normalize
+```
+
+`OmicsAttentionModule` is optional — include it for var-level (gene-vector)
+inputs where the sequence length > 1; omit it for obs-level (single cell
+embedding) inputs where it would degenerate to a feedforward layer.
+
+### Key design decisions
+
+- **`input_values` preprocess key** — omics preprocess writes `input_values`
+  (not `token_embeddings`) so the ST training collator's `collect_features`
+  suffix matching detects omics columns. Forward reads `input_values` and
+  writes `token_embeddings` for downstream modules.
+- **`modality_ids` tensor** — 0 = text, 1 = omics, 2 = pad. The
+  `AdapterModule` uses this to route tokens through the correct projection head.
+- **No Trainer subclass** — training uses the standard
+  `SentenceTransformerTrainer` + `MultipleNegativesRankingLoss` with
+  `(anchor, positive)` column layout.
+
+---
+
+## 2 What was completed (Phases 1–5)
+
+| Phase | Deliverable                                                                                                                             | Tests                                                         |
+| ----- | --------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------- |
+| 1     | `VectorStore` — mmap-backed vector lookup with `from_numpy`, `from_adata`, `from_dict`, `load`                                          | `tests/test_vector_store.py` (in mmcontext_module tests)      |
+| 2     | `MMContextModule` — text encoding via AutoModel/AutoTokenizer, omics via VectorStore or direct vectors, `preprocess`/`forward` contract | `tests/test_mmcontext_module.py`                              |
+| 3     | `AdapterModule` — modality-aware projection (text head, omics head, shared dim), safetensors persistence                                | `tests/test_adapter_module.py`                                |
+| 4     | `OmicsAttentionModule` — optional self-attention over variable-length omics sequences                                                   | `tests/test_omics_attention_module.py`                        |
+| 5     | Full ST integration — pipeline construction, encode, save/load, training (text-only, bimodal, gene-list), `prepare_vector_store`        | `tests/test_st_integration.py`, `tests/test_prepare_store.py` |
+
+### Supporting files
+
+- `scripts/train_tiny.py` — end-to-end training script for `jo-mengr/cxg_schaefer_tiny` with wandb, MPS, and auto VectorStore preparation.
+- `src/mmcontext/io/prepare_store.py` — memory-efficient store builder that reads only needed obsm rows from zarr (never loads full AnnData).
+
+---
+
+## 3 Legacy code (`_legacy/`)
+
+The following modules were moved to `src/mmcontext/_legacy/` and are preserved
+for backward compatibility with previously trained models:
+
+| File                           | What it was                                                                  | Notes                                                                                                                 |
+| ------------------------------ | ---------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------- |
+| `mmcontextencoder.py`          | Dual-tower encoder (text + omics), `MMContextProcessor` tokenizer            | Central old architecture; `embed/model_utils.py` still imports `MMEnc.get_initial_embeddings_from_adata_link` from it |
+| `adapters.py`                  | MLP adapter with identity/linear/2-layer modes                               | Superseded by `modules/adapter_module.py` which adds modality-awareness                                               |
+| `omicsencoder.py`              | `MiniOmicsModel` — `nn.Embedding` lookup with HF `PreTrainedModel` interface | Superseded by `VectorStore` (no learnable embedding, just mmap lookup)                                                |
+| `onehot.py`                    | `OneHotTextEncoder` — learnable embedding per unique sentence                | Useful for ablation experiments; no new equivalent                                                                    |
+| `cell_sentence_transformer.py` | `OmicsEncoder` — full transformer over omics tokens with cross-attention     | Heavier than `OmicsAttentionModule`; cross-attention is not yet in new code                                           |
+
+### Legacy tests (still functional)
+
+These test files import from `_legacy` and exercise the old architecture. They
+should continue to pass for backward compatibility:
+
+- `test_mmcontext_encoder.py`
+- `test_mmcontext_tokenizer.py`
+- `test_miniOmicsEncoder.py`
+- `test_register_initial_embeddings.py`
+- `test_sentence_transformer_integration.py` (old integration tests)
+
+---
+
+## 4 Active non-module code — status & required changes
+
+### `callback.py` — trainer callbacks
+
+**Status:** Functional but coupled to old `MMContextEncoder` attribute paths
+(`model[0].text_encoder`, `model[0].text_adapter`, `model[0].omics_adapter`).
+
+**What needs to change:**
+
+- `UnfreezeTextEncoderCallback` should look for `model[0].auto_model` (the
+  `MMContextModule`'s text encoder) instead of `model[0].text_encoder`.
+- `UnfreezeAdapterCallback` should look for the `AdapterModule` by iterating
+  `model.children()` and checking `isinstance(m, AdapterModule)`, since it's
+  no longer at `model[0]` — it's `model[1]` (or `model[2]` if
+  OmicsAttentionModule is included).
+- Consider making freeze/unfreeze methods on the modules themselves
+  (`MMContextModule.freeze_text_encoder()` already exists;
+  `AdapterModule` could add `freeze_text_head()` / `freeze_omics_head()`).
+
+### `utils.py` — training utilities
+
+**Useful functions that work with new architecture as-is:**
+
+- `truncate_cell_sentences()` — gene-name truncation/filtering for HF datasets.
+  Domain-specific and actively needed for dataset preprocessing.
+- `truncate_semantic_cell_sentences_dataset()` — semantic-aware truncation.
+- `resolve_negative_indices_and_rename()` — resolves `negative_*_idx` columns
+  to actual text values. Needed for multiplet training with the real dataset.
+- `get_evaluator()` — instantiates ST evaluators (BinaryClassification,
+  Triplet) from dataset columns. Works with any ST model.
+- `consolidate_low_frequency_categories()` — groups rare labels into "other".
+- `get_device()` — simple CUDA/MPS/CPU device selection.
+
+**Functions that reference old architecture:**
+
+- `get_loss()` (both overloads) — hardcodes dataset_type → loss mapping.
+  Logic is fine, but should be reviewed for new column naming conventions.
+- `prepare_omics_resources()` — builds lookup dict with `"sample_idx:"` prefix.
+  Superseded by `prepare_vector_store()` with `"omics:"` prefix, but the old
+  function may still be needed for legacy model inference.
+
+### `hub_utils.py` — HuggingFace Hub upload
+
+**Status:** Template body references old module names. The upload logic
+(`HfApi` calls) is architecture-agnostic and works.
+
+**What needs to change:** Update the `MODEL_CARD_TEMPLATE` string to describe
+the new pipeline modules.
+
+### `file_utils.py` — I/O infrastructure
+
+**Status:** Actively used by both old and new code.
+
+**Key functions:**
+
+- `download_and_extract_links()` — full-featured download with caching,
+  Zenodo support, resume, retries. The new `io/prepare_store.py` has a
+  lighter `_download_zarr()` that doesn't support all features. Consider
+  converging on one implementation eventually.
+- `remove_corrupted_null_arrays()` — zarr repair utility; used by
+  `embed/dataset_utils.py`.
+- `collect_unique_links()` — deduplicates adata links from a HF dataset.
+- `load_test_adata()` — downloads + loads a test adata from HF.
+- `save_table()` — generic table saver (CSV/parquet/AnnData).
+
+---
+
+## 5 Evaluation framework — analysis & adaptation needs
+
+The `eval/` package is a **self-contained evaluation framework** with a
+decorator-based registry pattern. It is largely architecture-agnostic — most
+evaluators work on AnnData `.obsm` embeddings, not on the model directly.
+
+### Architecture
+
+```
+eval/registry.py      — @register decorator, get(name) lookup
+eval/base.py          — BaseEvaluator (abstract), EvalResult
+eval/eval_pipeline.py — orchestrator: discover evaluators, run on datasets
+eval/utils.py         — LabelKind, LabelSpec helpers
+```
+
+### Evaluators
+
+| Evaluator             | Registry name                | What it measures                                  | Architecture dependency                          |
+| --------------------- | ---------------------------- | ------------------------------------------------- | ------------------------------------------------ |
+| `ARI`                 | `"ARI"`                      | Adjusted Rand Index (KMeans clustering vs labels) | None — operates on embeddings                    |
+| `LabelSimilarity`     | (registered)                 | ROC-AUC of intra- vs inter-label cosine sim       | None                                             |
+| `ScibBundle`          | `"scib"`                     | scIB benchmark metrics (batch/bio)                | None                                             |
+| `UmapPlotter`         | (registered)                 | UMAP visualizations colored by labels             | Uses `pl/plotting.py`                            |
+| `OmicsQueryAnnotator` | (not registered, standalone) | Zero-shot annotation via cosine similarity        | Needs `model.encode()` — works with any ST model |
+
+### `embedding_alignment.py`
+
+Standalone module (not registered as an evaluator) that computes cross-modal
+alignment scores. Useful for measuring how well the shared space aligns
+text and omics embeddings. Works on raw numpy arrays.
+
+### What needs to change for new architecture
+
+1. **Embedding pipeline (`embed/`)** — `embed/model_utils.py` imports
+   `MMContextEncoder.get_initial_embeddings_from_adata_link` to register
+   omics vectors before embedding. This needs a new path:
+   - Load the ST model
+   - Detect if it has an `MMContextModule` at position 0
+   - Attach a `VectorStore` (built via `prepare_vector_store`) instead of
+     calling `register_initial_embeddings`
+   - The `prepare_model_and_embed()` function should be updated accordingly
+
+2. **Eval pipeline** — The evaluators themselves don't need changes (they
+   operate on `adata.obsm` arrays), but the `embed_pipeline.py` orchestrator
+   that feeds them embeddings does (see point 1).
+
+3. **`OmicsQueryAnnotator`** — Works with any model with `.encode()`.
+   No changes needed for the new pipeline.
+
+---
+
+## 6 Missing features & test gaps
+
+### Features not yet covered by new modules
+
+- [ ] **Callbacks for new architecture** — `UnfreezeTextEncoderCallback` and
+      `UnfreezeAdapterCallback` need to be updated for the new module layout
+      (see §4).
+
+- [ ] **Cross-attention** — The old `OmicsEncoder` in
+      `cell_sentence_transformer.py` supported cross-attention between text
+      and omics sequences. The new `OmicsAttentionModule` only does
+      self-attention on omics tokens. Cross-attention would enable richer
+      multimodal interaction but adds complexity.
+
+- [ ] **`OneHotTextEncoder` equivalent** — Useful for ablation experiments
+      (fast training without a real transformer). Currently only in `_legacy/`.
+      Could be a simple fixture/utility rather than a full module.
+
+- [ ] **Negative mining / hard negatives** — The training script currently
+      uses `(anchor, positive)` pairs with in-batch negatives (MNR loss).
+      The dataset has `negative_1_idx` and `negative_2_idx` columns.
+      `resolve_negative_indices_and_rename()` in `utils.py` handles
+      resolving these to actual text. Supporting `(anchor, positive, negative)`
+      triplets would improve training quality.
+
+- [ ] **Hub upload for new architecture** — `hub_utils.py` model card
+      template references old module names.
+
+- [ ] **Embed pipeline for new architecture** — `embed/model_utils.py`
+      uses `MMContextEncoder.get_initial_embeddings_from_adata_link()`.
+      Needs updating to use `VectorStore` + `prepare_vector_store`.
+
+### Test gaps
+
+- [ ] **`VectorStore` standalone tests** — Currently exercised through
+      `test_mmcontext_module.py` and `test_prepare_store.py` fixtures, but
+      there's no dedicated `test_vector_store.py` testing all construction
+      methods (`from_numpy`, `from_adata`, `from_dict`, `from_dataframe`),
+      edge cases, and persistence.
+
+- [ ] **AdapterModule standalone tests** — verify existence of
+      `tests/test_adapter_module.py` and confirm coverage.
+
+- [ ] **Mixed-modality batches** — No test currently sends a batch where
+      some samples are text and some are omics through the full pipeline
+      in a single forward pass. This is an important edge case for
+      training with heterogeneous datasets.
+
+- [ ] **Gradient flow end-to-end** — Training tests verify parameters change,
+      but don't check that gradients flow through specific module boundaries
+      (e.g., from loss through Pooling → AdapterModule → MMContextModule's
+      text encoder).
+
+- [ ] **Multi-GPU / DDP** — No tests for distributed training.
+
+- [ ] **Large-scale prepare_store** — `test_prepare_store.py` tests local
+      zarr files; no integration test for the Zenodo download path
+      (intentionally — network tests are fragile).
+
+### Potential improvements
+
+- [ ] **Converge download logic** — `file_utils.download_and_extract_links`
+      and `io/prepare_store._download_zarr` both handle zarr downloads with
+      different feature sets. Could share a common download backend.
+
+- [ ] **`simulator.py` integration** — The old `LOSS_PRESETS` dict and
+      `make_cluster_sampler` are valuable for generating synthetic test
+      datasets. Consider moving to a `testing/` subpackage.
+
+- [ ] **`sanity_helpers.py`** — Contains `plot_pca` and possibly other
+      debug helpers not in `pl/`. Consider merging into `pl/` or a
+      `debug/` module.
+
+---
+
+## 7 File map (post-cleanup)
+
+```
+src/mmcontext/
+├── __init__.py                  # exports: modules, io, eval, pl, embed
+├── _legacy/                     # old architecture (preserved for compat)
+│   ├── adapters.py
+│   ├── cell_sentence_transformer.py
+│   ├── mmcontextencoder.py
+│   ├── omicsencoder.py
+│   └── onehot.py
+├── modules/                     # NEW: ST v5.4+ pipeline modules
+│   ├── mmcontext_module.py      #   InputModule (text + omics routing)
+│   ├── adapter_module.py        #   modality-aware projection
+│   └── omics_attention_module.py#   optional self-attention (var-level)
+├── io/                          # NEW: data I/O
+│   ├── vector_store.py          #   mmap-backed vector lookup
+│   └── prepare_store.py         #   zarr → VectorStore builder
+├── eval/                        # evaluation framework (active, mostly arch-agnostic)
+│   ├── base.py, registry.py     #   BaseEvaluator + decorator registry
+│   ├── eval_pipeline.py         #   orchestrator
+│   ├── ari.py                   #   ARI evaluator
+│   ├── label_similarity.py      #   label similarity evaluator
+│   ├── embedding_alignment.py   #   cross-modal alignment scores
+│   ├── scib_wrapper.py          #   scIB benchmark
+│   ├── query_annotate.py        #   zero-shot annotation
+│   ├── umap_plotter.py          #   UMAP visualizations
+│   └── utils.py                 #   LabelKind, LabelSpec
+├── embed/                       # embedding generation pipeline (active, needs update)
+│   ├── embed_pipeline.py        #   orchestrator
+│   ├── model_utils.py           #   model loading + embedding (imports _legacy)
+│   ├── dataset_utils.py         #   adata/dataset loading
+│   ├── cellwhisperer_utils.py   #   CellWhisperer-specific helpers
+│   └── scsa_utils.py            #   scSA dataset helpers
+├── pl/                          # plotting (active)
+│   ├── plotting.py              #   UMAP + query-score plots
+│   └── metric_plots.py          #   benchmark result bar charts
+├── callback.py                  # trainer callbacks (needs update for new modules)
+├── file_utils.py                # I/O infrastructure (active)
+├── hub_utils.py                 # HF Hub upload (needs template update)
+├── sanity_helpers.py            # debug plotting (standalone)
+├── simulator.py                 # synthetic data generation (standalone, useful)
+├── utils.py                     # training utilities (mixed: some active, some legacy)
+└── models/                      # empty after cell_sentence_transformer.py moved
+    └── __init__.py
+
+scripts/
+├── train_tiny.py                # NEW: training script for new architecture
+├── train.py                     # old Hydra training script (uses _legacy)
+└── train_merged.py              # old merged training script (uses _legacy)
+
+tests/
+├── test_mmcontext_module.py     # NEW: MMContextModule unit tests
+├── test_adapter_module.py       # NEW: AdapterModule unit tests (verify exists)
+├── test_omics_attention_module.py # NEW: OmicsAttentionModule unit tests
+├── test_st_integration.py       # NEW: full pipeline integration + training
+├── test_prepare_store.py        # NEW: VectorStore preparation tests
+├── test_mmcontext_encoder.py    # legacy: old MMContextEncoder tests
+├── test_mmcontext_tokenizer.py  # legacy: old tokenizer tests
+├── test_miniOmicsEncoder.py     # legacy: old MiniOmicsModel tests
+├── test_register_initial_embeddings.py  # legacy
+├── test_sentence_transformer_integration.py  # legacy: old ST integration
+└── conftest.py                  # shared fixtures (stub encoder, both old+new)
+```
diff --git a/conf/eval/collect_metrics_conf.yaml b/conf/eval/collect_metrics_conf.yaml
index 4033b8b..47cc5a1 100644
--- a/conf/eval/collect_metrics_conf.yaml
+++ b/conf/eval/collect_metrics_conf.yaml
@@ -107,6 +107,7 @@ collect_metrics:
     label_kind_filter: bio # null = all, or "bio" / "batch"
     metrics:
       - "LabelSimilarity/balanced_accuracy"
+      - "LabelSimilarity/random_baseline_accuracy"
       - "LabelSimilarity/macro_f1"
       - "LabelSimilarity/mean_auc"
       - "LabelSimilarity/mrr"
@@ -115,6 +116,7 @@ collect_metrics:
       - "LabelSimilarity/topk_accuracy@5"
     metric_display_names:
       "LabelSimilarity/balanced_accuracy": "Bal.\\ Acc."
+      "LabelSimilarity/random_baseline_accuracy": "Baseline Acc. (1/n)"
       "LabelSimilarity/macro_f1": "Macro F1"
       "LabelSimilarity/mean_auc": "Mean AUC"
       "LabelSimilarity/mrr": "MRR"
diff --git a/conf/training/mmcontext_v2.yaml b/conf/training/mmcontext_v2.yaml
new file mode 100644
index 0000000..4657968
--- /dev/null
+++ b/conf/training/mmcontext_v2.yaml
@@ -0,0 +1,83 @@
+# Config-driven multi-dataset training (scripts/train_config.py).
+# Plain YAML — no Hydra. Run:
+#   python scripts/train_config.py --config conf/training/multi_example.yaml
+
+# --- Model ---
+model: null # HF repo / local path to continue fine-tuning, or null to build fresh
+tag: cxg_full # optional suffix appended to the auto-generated Hub model name
+push_to_hub: true
+
+# --- Shared embedding space ---
+obsm_key: X_gs10k # single obsm key used to build the VectorStore for ALL bimodal datasets
+shared_dim: 256 # output dim of the shared space
+adapter_hidden_dim: null # null/0 = Linear+LayerNorm head; >0 = MLP (Linear->ReLU->Linear->LayerNorm)
+
+# --- Text encoder ---
+text_encoder:
+  name: NeuML/pubmedbert-base-embeddings
+  max_seq_length: 512
+
+# --- Freezing ---
+freeze_text_encoder: true # static freeze at start
+unfreeze_last_n_layers: 1 # with freeze_text_encoder: keep top N layers trainable (0 = freeze all)
+unfreeze_epoch: 16 # unfreeze the WHOLE encoder at this epoch (UnfreezeTextEncoderCallback)
+# adapter_freezing:        # optional; omit to skip the per-modality adapter callback
+#   freeze_text_adapter: false
+#   freeze_omics_adapter: false
+#   unfreeze_text_adapter_epoch: null
+#   unfreeze_omics_adapter_epoch: null
+
+# --- Gene tokens (only affect text-modality truncation) ---
+gene_special_token: null # e.g. "[GENE]" — added to the tokenizer and used by truncation
+delimiter: " "
+
+# --- Caching ---
+cache_dir: data/huggingface # HF datasets cache (null = default)
+adata_cache_dir: data/anndata # zarr/AnnData download cache for the VectorStore
+vector_store_dir: data/vector_stores # where per-dataset + merged .mmap stores are written
+force_refresh_cache: false
+
+# --- Omics datasets (processed via prepare_dataset) ---
+# sample_idx is auto-namespaced to "{name}:{sample_idx}" so multiple datasets
+# never collide in the merged VectorStore. truncate/truncate_kwargs only apply
+# to modality: text (the omics-vector anchor of bimodal datasets is not text).
+omics_datasets:
+  - id: jo-mengr/cxg_schaefer_full
+    name: cxg_full
+    type: multiplets # -> MultipleNegativesRankingLoss + TripletEvaluator
+    modality: bimodal # bimodal (omics vectors) | text (cell-sentence text)
+    primary_cell_sentence: cell_sentence_1
+    positive_col: positive
+    use_hard_negatives: true
+    truncate: false
+
+# --- Bio datasets (already anchor/positive/negative; loaded raw) ---
+bio_datasets: []
+#  - id: jo-mengr/ncbi_cell_types_1000
+#    name: celltype_descriptions
+#    type: multiplets
+#    revision: hard_negatives
+
+# --- Trainer ---
+trainer:
+  output_dir: outputs/mmcontext_v2
+  num_train_epochs: 4
+  per_device_train_batch_size: 128
+  per_device_eval_batch_size: 128
+  learning_rate: 2e-5
+  warmup_ratio: 0.1
+  bf16: true
+  fp16: false
+  eval_strategy: steps
+  eval_steps: 200
+  save_strategy: steps
+  save_steps: 2000
+  save_total_limit: 2
+  logging_steps: 50
+  max_grad_norm: 1.0
+  gradient_checkpointing: false
+  dataloader_num_workers: 0
+  use_mps_device: true
+  use_cosine_scheduler: false
+  wandb_project: mmcontext # set to enable W&B logging
+  wandb_run_name: null
diff --git a/conf/training/mmcontext_v2_tune.yaml b/conf/training/mmcontext_v2_tune.yaml
new file mode 100644
index 0000000..070976e
--- /dev/null
+++ b/conf/training/mmcontext_v2_tune.yaml
@@ -0,0 +1,93 @@
+# Config-driven multi-dataset training (scripts/train_config.py).
+# Plain YAML — no Hydra. Run:
+#   python scripts/train_config.py --config conf/training/multi_example.yaml
+
+# --- Model ---
+model: outputs/tiny_model/final # HF repo / local path to continue fine-tuning, or null to build fresh
+tag: cxg_tiny_tuned # optional suffix appended to the auto-generated Hub model name
+push_to_hub: true
+
+# --- Shared embedding space ---
+obsm_key: X_scvi_fm # single obsm key used to build the VectorStore for ALL bimodal datasets
+shared_dim: 256 # output dim of the shared space
+adapter_hidden_dim: null # null/0 = Linear+LayerNorm head; >0 = MLP (Linear->ReLU->Linear->LayerNorm)
+
+# --- Text encoder ---
+text_encoder:
+  name: NeuML/pubmedbert-base-embeddings
+  max_seq_length: 512
+
+mining:
+  enabled: true
+  num_negatives: 3 # negatives to mine per anchor
+  range_min: 1 # skip the N closest candidates (closest is usually the true positive)
+  range_max: null # only consider candidates up to this rank (null = no bound)
+  max_score: 0.95 # drop candidates above this similarity (avoid false negatives)
+  relative_margin: null # negatives must score below relative_margin * positive_score
+  sampling_strategy: top # "top" | "random"
+  use_faiss: false # use FAISS for the similarity search (large datasets)
+
+# --- Freezing ---
+freeze_text_encoder: true # static freeze at start
+unfreeze_last_n_layers: 1 # with freeze_text_encoder: keep top N layers trainable (0 = freeze all)
+unfreeze_epoch: 16 # unfreeze the WHOLE encoder at this epoch (UnfreezeTextEncoderCallback)
+# adapter_freezing:        # optional; omit to skip the per-modality adapter callback
+#   freeze_text_adapter: false
+#   freeze_omics_adapter: false
+#   unfreeze_text_adapter_epoch: null
+#   unfreeze_omics_adapter_epoch: null
+
+# --- Gene tokens (only affect text-modality truncation) ---
+gene_special_token: null # e.g. "[GENE]" — added to the tokenizer and used by truncation
+delimiter: " "
+
+# --- Caching ---
+cache_dir: /Volumes/sandisk_2tb/dev/mmcontext/data/huggingface # HF datasets cache (null = default)
+adata_cache_dir: /Volumes/sandisk_2tb/dev/mmcontext/data/anndata # zarr/AnnData download cache for the VectorStore
+vector_store_dir: data/vector_stores # where per-dataset + merged .mmap stores are written
+force_refresh_cache: false
+
+# --- Omics datasets (processed via prepare_dataset) ---
+# sample_idx is auto-namespaced to "{name}:{sample_idx}" so multiple datasets
+# never collide in the merged VectorStore. truncate/truncate_kwargs only apply
+# to modality: text (the omics-vector anchor of bimodal datasets is not text).
+omics_datasets:
+  - id: jo-mengr/cxg_schaefer_tiny
+    name: cxg_tiny
+    type: multiplets # -> MultipleNegativesRankingLoss + TripletEvaluator
+    modality: bimodal # bimodal (omics vectors) | text (cell-sentence text)
+    primary_cell_sentence: cell_sentence_1
+    positive_col: positive
+    use_hard_negatives: false
+    truncate: false
+
+# --- Bio datasets (already anchor/positive/negative; loaded raw) ---
+bio_datasets: []
+#  - id: jo-mengr/ncbi_cell_types_1000
+#    name: celltype_descriptions
+#    type: multiplets
+#    revision: hard_negatives
+
+# --- Trainer ---
+trainer:
+  output_dir: outputs/mmcontext_v2_tuned
+  num_train_epochs: 4
+  per_device_train_batch_size: 64
+  per_device_eval_batch_size: 64
+  learning_rate: 2e-5
+  warmup_ratio: 0.1
+  bf16: true
+  fp16: false
+  eval_strategy: steps
+  eval_steps: 200
+  save_strategy: steps
+  save_steps: 2000
+  save_total_limit: 2
+  logging_steps: 50
+  max_grad_norm: 1.0
+  gradient_checkpointing: false
+  dataloader_num_workers: 0
+  use_mps_device: true
+  use_cosine_scheduler: false
+  wandb_project: mmcontext # set to enable W&B logging
+  wandb_run_name: null
diff --git a/conf/training/multi_example.yaml b/conf/training/multi_example.yaml
new file mode 100644
index 0000000..3d62926
--- /dev/null
+++ b/conf/training/multi_example.yaml
@@ -0,0 +1,106 @@
+# Config-driven multi-dataset training (scripts/train_config.py).
+# Plain YAML — no Hydra. Run:
+#   python scripts/train_config.py --config conf/training/multi_example.yaml
+
+# --- Model ---
+model: null # HF repo / local path to continue fine-tuning, or null to build fresh
+tag: cxg_geo_tiny # optional suffix appended to the auto-generated Hub model name
+push_to_hub: false
+
+# --- Shared embedding space ---
+obsm_key: X_gs10k # single obsm key used to build the VectorStore for ALL bimodal datasets
+shared_dim: 256 # output dim of the shared space
+adapter_hidden_dim: null # null/0 = Linear+LayerNorm head; >0 = MLP (Linear->ReLU->Linear->LayerNorm)
+
+# --- Text encoder ---
+text_encoder:
+  name: NeuML/pubmedbert-base-embeddings
+  max_seq_length: 512
+
+# --- Freezing ---
+freeze_text_encoder: true # static freeze at start
+unfreeze_last_n_layers: 1 # with freeze_text_encoder: keep top N layers trainable (0 = freeze all)
+unfreeze_epoch: 16 # unfreeze the WHOLE encoder at this epoch (UnfreezeTextEncoderCallback)
+# adapter_freezing:        # optional; omit to skip the per-modality adapter callback
+#   freeze_text_adapter: false
+#   freeze_omics_adapter: false
+#   unfreeze_text_adapter_epoch: null
+#   unfreeze_omics_adapter_epoch: null
+
+# --- Hard-negative mining (optional) ---
+# Continue fine-tuning by mining model-specific hard negatives instead of using
+# the dataset's own negative_*_idx columns (cf. scripts/finetune_tiny.py). Only
+# meaningful with `model:` set to an already-trained checkpoint — mining against
+# a freshly initialised bimodal model produces noise. When enabled, both the
+# train and eval splits are mined and the model card records these settings.
+# mining:
+#   enabled: true
+#   num_negatives: 3          # negatives to mine per anchor
+#   range_min: 1              # skip the N closest candidates (closest is usually the true positive)
+#   range_max: null           # only consider candidates up to this rank (null = no bound)
+#   max_score: 0.95           # drop candidates above this similarity (avoid false negatives)
+#   relative_margin: null     # negatives must score below relative_margin * positive_score
+#   sampling_strategy: top    # "top" | "random"
+#   use_faiss: false          # use FAISS for the similarity search (large datasets)
+
+# --- Gene tokens (only affect text-modality truncation) ---
+gene_special_token: null # e.g. "[GENE]" — added to the tokenizer and used by truncation
+delimiter: " "
+
+# --- Caching ---
+cache_dir: data/huggingface # HF datasets cache (null = default)
+adata_cache_dir: data/anndata # zarr/AnnData download cache for the VectorStore
+vector_store_dir: data/vector_stores # where per-dataset + merged .mmap stores are written
+force_refresh_cache: false
+
+# --- Omics datasets (processed via prepare_dataset) ---
+# sample_idx is auto-namespaced to "{name}:{sample_idx}" so multiple datasets
+# never collide in the merged VectorStore. truncate/truncate_kwargs only apply
+# to modality: text (the omics-vector anchor of bimodal datasets is not text).
+omics_datasets:
+  - id: jo-mengr/cxg_schaefer_tiny
+    name: cxg_tiny
+    type: multiplets # -> MultipleNegativesRankingLoss + TripletEvaluator
+    modality: bimodal # bimodal (omics vectors) | text (cell-sentence text)
+    primary_cell_sentence: cell_sentence_1
+    positive_col: positive
+    use_hard_negatives: true
+    truncate: false
+  - id: jo-mengr/geo_schaefer_tiny
+    name: geo_tiny
+    type: multiplets
+    modality: bimodal
+    primary_cell_sentence: cell_sentence_1
+    positive_col: positive
+    use_hard_negatives: true
+
+# --- Bio datasets (already anchor/positive/negative; loaded raw) ---
+bio_datasets:
+  - id: jo-mengr/ncbi_cell_types_1000
+    name: celltype_descriptions
+    type: multiplets
+    revision: hard_negatives
+
+# --- Trainer ---
+trainer:
+  output_dir: outputs/multi_gs10k
+  num_train_epochs: 2
+  per_device_train_batch_size: 64
+  per_device_eval_batch_size: 64
+  learning_rate: 2e-5
+  warmup_ratio: 0.1
+  bf16: true
+  fp16: false
+  eval_strategy: steps
+  eval_steps: 200
+  save_strategy: steps
+  save_steps: 2000
+  save_total_limit: 2
+  logging_steps: 50
+  max_grad_norm: 1.0
+  gradient_checkpointing: false
+  dataloader_num_workers: 0
+  use_mps_device: true
+  use_cosine_scheduler: false
+  wandb_project: mmcontext # set to enable W&B logging
+  wandb_run_name: null
diff --git a/docs/agentic_workflow_research.md b/docs/agentic_workflow_research.md
new file mode 100644
index 0000000..b3408c4
--- /dev/null
+++ b/docs/agentic_workflow_research.md
@@ -0,0 +1,239 @@
+# Agentic Workflow Plan: mmcontext Repository Automation
+
+## Decision: Claude-native approach
+
+We're going all-in on the Claude ecosystem: Claude Code Action for GitHub automation, gh-aw for scheduled audits, Cowork for interactive work. No multi-provider abstraction for now.
+
+## Current State
+
+- Repo: `github.com/mengerj/mmcontext` — multimodal contrastive learning (text + omics) via sentence-transformers
+- Existing workflows: `test.yaml`, `build.yaml`, `release.yaml`
+- Branch strategy: `dev-claude` exists as the integration branch for agent work
+- Auth: Need to run `claude setup-token` and store OAuth token as repo secret
+- No `CLAUDE.md`, no benchmarks, no agent workflows yet
+
+---
+
+## Phase 1: Foundation — CLAUDE.md
+
+**What**: Create `CLAUDE.md` in repo root with project context, conventions, and instructions for agent behavior.
+
+**Who does what**:
+
+- **Claude (Cowork)**: Drafts the file based on repo analysis
+- **Jo**: Reviews, adjusts tone/priorities, commits to `dev-claude`
+
+**Contents**: Package structure, key modules, test/lint commands, branch conventions (always branch from `dev-claude`), PR conventions (link to issues), coding standards (ruff, type hints, docstrings).
+
+---
+
+## Phase 2: Authentication Setup
+
+**What**: Generate OAuth token for subscription-based GitHub Actions.
+
+**Who does what**:
+
+- **Jo** (manual, ~5 min):
+  1. Run `claude setup-token` in terminal
+  2. Go to GitHub repo → Settings → Secrets → Actions
+  3. Add secret `CLAUDE_CODE_OAUTH_TOKEN` with the token value
+
+**Note**: Token expires if you log out of Claude Code. Regenerate with `claude setup-token` and update the secret. Starting June 15, 2026, programmatic usage draws from a separate credit pool capped at your subscription fee.
+
+---
+
+## Phase 3: PR Review Workflow
+
+**What**: Automated code review on every PR targeting `dev-claude` or `main`.
+
+**Who does what**:
+
+- **Claude (Cowork)**: Drafts `.github/workflows/claude-review.yaml`
+- **Jo**: Reviews, commits, tests with a real PR
+
+**Workflow triggers**: PR opened/updated against `dev-claude` or `main`, plus `@claude` mentions in PR comments.
+
+**Key config**:
+
+- OAuth token auth (no API key)
+- `actions: read` permission so Claude can see CI results
+- Custom instructions pointing to `CLAUDE.md`
+- Allowed tools: `Bash(pytest)`, `Bash(ruff check)`
+
+---
+
+## Phase 4: Issue-to-PR Workflow (plan-first protocol)
+
+**What**: When you mention `@claude` on an issue, the agent:
+
+1. Posts clarifying questions if the issue is ambiguous
+2. Posts an implementation plan as a checkbox list in a comment
+3. Waits for your approval (`@claude approved` or `@claude go ahead`)
+4. Creates a branch from `dev-claude`, implements the plan, ticks checkboxes as it goes
+5. Opens a PR linked to the issue
+
+**Who does what**:
+
+- **Claude (Cowork)**: Drafts `.github/workflows/claude-implement.yaml` and the plan-first protocol in `CLAUDE.md`
+- **Jo**: Reviews, commits, tests with a real issue
+
+**Branch naming**: `claude/<issue-number>-<short-slug>` (e.g., `claude/47-add-benchmark-scripts`)
+
+**Issue linking**: Claude Code Action can include `Fixes #47` or `Closes #47` in PR descriptions, which GitHub auto-links. The CLAUDE.md instructions will enforce this convention.
+
+**Plan-first enforcement**: Encoded in `CLAUDE.md` as the "Issue Implementation Protocol" — the agent reads this on every run and follows the plan-first flow.
+
+**First test**: Create an issue for the benchmark scripts (Phase 6) and let `@claude` implement it. This gives us a real test of the workflow AND gets the benchmarks built.
+
+---
+
+## Phase 5: Benchmark Scripts
+
+**What**: pytest-benchmark scripts measuring the three priority metrics:
+
+1. **Peak GPU memory** — `torch.cuda.max_memory_allocated()` during forward/training
+2. **Embedding quality** — kNN accuracy, retrieval metrics from `mmcontext.eval`
+3. **Training throughput** — samples/sec during a training loop
+
+**Who does what — two options**:
+
+- **Option A (GitHub workflow test)**: Create a GitHub issue with benchmark requirements, `@claude` it, and let the agent build the scripts. Good test of Phase 4.
+- **Option B (Cowork)**: Build here interactively if the GitHub workflow isn't set up yet or the task proves too complex for headless execution.
+
+**Recommendation**: Try Option A first. If the agent struggles (benchmarks need domain knowledge about mmcontext's model/data pipeline), fall back to Option B.
+
+**GPU constraint**: GitHub Actions runners lack GPUs. The benchmark workflow will run CPU-only smoke tests for regression detection. Real GPU benchmarks run locally or via SLURM (you already have `.slurm` scripts). Results can be pushed to the same tracking system.
+
+---
+
+## Phase 6: Benchmark GitHub Action
+
+**What**: `.github/workflows/benchmark.yaml` — runs benchmarks on push to `main`/`dev-claude`, weekly schedule, and tracks results via [github-action-benchmark](https://github.com/benchmark-action/github-action-benchmark).
+
+**Who does what**:
+
+- **Claude (Cowork or @claude)**: Drafts the workflow
+- **Jo**: Reviews, enables GitHub Pages for the benchmark chart
+
+**Features**: Publishes charts to GitHub Pages, comments on PRs when performance regresses >150%, stores historical data in `gh-pages` branch.
+
+---
+
+## Phase 7: gh-aw Scheduled Audits
+
+**What**: Use [GitHub Agentic Workflows](https://github.github.com/gh-aw/) for recurring, automated codebase analysis. Workflows are markdown files that compile to GitHub Actions YAML with built-in security guardrails (read-only tokens, network firewall, threat detection).
+
+**Who does what**:
+
+- **Jo** (manual, ~10 min):
+  1. `gh extension install github/gh-aw`
+  2. `gh aw init` in the repo
+  3. Set up engine secrets (Claude via OAuth)
+- **Claude (Cowork)**: Drafts workflow markdown files
+- **Jo**: Reviews, merges, monitors first runs
+
+**Planned workflows**:
+
+### 7a. Weekly Code Quality Audit
+
+```markdown
+## <!-- .github/workflows/code-quality-audit.md -->
+
+name: Code Quality Audit
+on:
+schedule: weekly on mondays
+engine: claude
+safe-outputs:
+create-discussion:
+title-prefix: "[audit] "
+category: "audits"
+max: 1
+tools:
+cache-memory: true
+github:
+toolsets: [default]
+
+---
+
+Analyze the mmcontext codebase for:
+
+- Test coverage gaps in src/mmcontext/
+- Type annotation completeness
+- Memory efficiency in data loading (io/, embed/)
+- API compatibility with sentence-transformers v5.4
+- Dead code in \_legacy/ that should be removed
+
+Compare against previous audit findings in cache memory.
+Create a discussion with prioritized recommendations.
+```
+
+### 7b. Test Coverage Monitor
+
+Runs on push, checks what changed, flags untested code paths.
+
+### 7c. ResearchPlanAssignOps cycle
+
+The audit discussion feeds into `/plan` → sub-issues → `@claude` implements. Full loop.
+
+**gh-aw engine options**: Can use Claude, Copilot, or Codex as engine — useful for future comparison without building custom abstractions.
+
+---
+
+## Phase 8: Cowork Scheduled Tasks — Comparison
+
+**What**: Set up a Cowork scheduled task that does similar work to the gh-aw audit, then compare the two approaches.
+
+**Differences**:
+| Aspect | gh-aw | Cowork Scheduled Task |
+|--------|-------|-----------------------|
+| Runs where | GitHub Actions (cloud) | Local (Cowork must be open) |
+| Security | 5-layer guardrails, read-only token | Full local access |
+| Output | GitHub Discussions/Issues | Chat notification in Cowork |
+| Data access | Repo only (sandboxed) | Full local filesystem |
+| Scheduling | Cron via Actions | Cron via Cowork app |
+| Cost | GitHub Actions minutes + LLM tokens | Subscription credit pool |
+
+**Verdict**: gh-aw for anything that should run reliably regardless of whether your laptop is open. Cowork tasks for things that need local context (your `.venv`, local data files, GPU) or personal briefings.
+
+---
+
+## Implementation Order
+
+| Step  | What                                       | Who                            | Blocked by |
+| ----- | ------------------------------------------ | ------------------------------ | ---------- |
+| **1** | Create `CLAUDE.md`                         | Claude (Cowork) → Jo reviews   | —          |
+| **2** | Run `claude setup-token`, add repo secret  | Jo (manual)                    | —          |
+| **3** | Draft + commit `claude-review.yaml`        | Claude (Cowork) → Jo reviews   | 1, 2       |
+| **4** | Draft + commit `claude-implement.yaml`     | Claude (Cowork) → Jo reviews   | 1, 2       |
+| **5** | Test: create benchmark issue, `@claude` it | Jo triggers, Claude implements | 3, 4       |
+| **6** | Draft + commit `benchmark.yaml`            | Claude (via @claude or Cowork) | 5          |
+| **7** | Install gh-aw, set up audit workflows      | Jo installs → Claude drafts    | 2          |
+| **8** | Set up Cowork scheduled task, compare      | Claude (Cowork)                | 7          |
+
+Steps 1–2 can happen in parallel. Steps 3–4 can happen in parallel after 1–2.
+
+---
+
+## Settings Jo Needs to Configure
+
+- [ ] `claude setup-token` → store as `CLAUDE_CODE_OAUTH_TOKEN` repo secret
+- [ ] Enable GitHub Pages on the repo (for benchmark charts) — Settings → Pages → Source: `gh-pages` branch
+- [ ] `gh extension install github/gh-aw` + `gh aw init`
+- [ ] Set up gh-aw engine secret for Claude (see [engine docs](https://github.github.com/gh-aw/reference/engines/))
+- [ ] Consider adding `pytest-benchmark` to `[project.optional-dependencies.test]` in `pyproject.toml`
+- [ ] Decide on GitHub Discussions category "audits" (create it in repo Settings → Discussions)
+
+---
+
+## Key Links
+
+- [Claude Code GitHub Actions docs](https://code.claude.com/docs/en/github-actions)
+- [claude-code-action repo](https://github.com/anthropics/claude-code-action)
+- [Claude Code Action with OAuth (Marketplace)](https://github.com/marketplace/actions/claude-code-action-with-oauth)
+- [Claude Code Routines docs](https://code.claude.com/docs/en/routines)
+- [claude-code-security-review](https://github.com/anthropics/claude-code-security-review)
+- [github-action-benchmark](https://github.com/benchmark-action/github-action-benchmark)
+- [pytest-benchmark](https://github.com/ionelmc/pytest-benchmark)
+- [GitHub Agentic Workflows](https://github.github.com/gh-aw/)
+- [gh-aw ResearchPlanAssignOps](https://github.github.com/gh-aw/patterns/research-plan-assign-ops/)
+- [June 2026 billing split](https://gist.github.com/yurukusa/7d854616809e673ca8d23353ed8267a6)
diff --git a/exploratory/notebooks/MMContextEncoder_quickstart.ipynb b/exploratory/notebooks/MMContextEncoder_quickstart.ipynb
index a5d83eb..013b8be 100644
--- a/exploratory/notebooks/MMContextEncoder_quickstart.ipynb
+++ b/exploratory/notebooks/MMContextEncoder_quickstart.ipynb
@@ -109,9 +109,9 @@
     "import numpy as np\n",
     "import pandas as pd\n",
     "import torch\n",
+    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from sentence_transformers import SentenceTransformer\n",
     "\n",
-    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from mmcontext.simulator import OmicsCaptionSimulator\n",
     "\n",
     "sim = OmicsCaptionSimulator(n_samples=100, n_genes=10).simulate()\n",
diff --git a/exploratory/notebooks/MMContextEncoder_training_demo.ipynb b/exploratory/notebooks/MMContextEncoder_training_demo.ipynb
index 6f4342f..3096ebf 100644
--- a/exploratory/notebooks/MMContextEncoder_training_demo.ipynb
+++ b/exploratory/notebooks/MMContextEncoder_training_demo.ipynb
@@ -108,9 +108,9 @@
     "import pandas as pd\n",
     "import torch\n",
     "from datasets import DatasetDict\n",
+    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from sentence_transformers import SentenceTransformer\n",
     "\n",
-    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from mmcontext.simulator import OmicsCaptionSimulator\n",
     "\n",
     "# simulate tiny dataset\n",
diff --git a/exploratory/notebooks/Training_sanity_check.ipynb b/exploratory/notebooks/Training_sanity_check.ipynb
index a9ae7ee..cfa67f9 100644
--- a/exploratory/notebooks/Training_sanity_check.ipynb
+++ b/exploratory/notebooks/Training_sanity_check.ipynb
@@ -72,8 +72,8 @@
    ],
    "source": [
     "import numpy as np\n",
-    "\n",
     "from mmcontext.mmcontextencoder import MMContextEncoder\n",
+    "\n",
     "from mmcontext.sanity_helpers import cluster_variances, plot_pca, stack_embeddings\n",
     "from mmcontext.simulator import OmicsCaptionSimulator, make_cluster_sampler\n",
     "\n",
diff --git a/exploratory/notebooks/clustering_issues.ipynb b/exploratory/notebooks/clustering_issues.ipynb
index 7acc426..ba1f00d 100644
--- a/exploratory/notebooks/clustering_issues.ipynb
+++ b/exploratory/notebooks/clustering_issues.ipynb
@@ -243,9 +243,10 @@
    ],
    "source": [
     "# Visualise the embeddings\n",
-    "from mmcontext.pl import plot_umap\n",
     "from mmcontext.pp.utils import consolidate_low_frequency_categories\n",
     "\n",
+    "from mmcontext.pl import plot_umap\n",
+    "\n",
     "adata_cut = consolidate_low_frequency_categories(adata, [\"cell_type\"], threshold=10)\n",
     "label_key = \"cell_type\"\n",
     "emb_key = \"mmcontext_emb\"\n",
@@ -364,9 +365,10 @@
    ],
    "source": [
     "# Visualise the embeddings\n",
-    "from mmcontext.pl import plot_umap\n",
     "from mmcontext.pp.utils import consolidate_low_frequency_categories\n",
     "\n",
+    "from mmcontext.pl import plot_umap\n",
+    "\n",
     "label_key = \"tissue\"\n",
     "adata_cut = consolidate_low_frequency_categories(adata, [\"tissue\"], threshold=0)\n",
     "emb_key = \"mmcontext_emb\"\n",
@@ -417,9 +419,10 @@
    ],
    "source": [
     "# Does maybe the original text embedding cluster after tissues?\n",
-    "from mmcontext.pl import plot_umap\n",
     "from mmcontext.pp.utils import consolidate_low_frequency_categories\n",
     "\n",
+    "from mmcontext.pl import plot_umap\n",
+    "\n",
     "label_key = \"tissue\"\n",
     "adata_cut = consolidate_low_frequency_categories(adata, [label_key], threshold=15)\n",
     "emb_key = \"og_text_emb\"\n",
diff --git a/exploratory/notebooks/load_and_push_model.ipynb b/exploratory/notebooks/load_and_push_model.ipynb
index 92bed6d..623d355 100644
--- a/exploratory/notebooks/load_and_push_model.ipynb
+++ b/exploratory/notebooks/load_and_push_model.ipynb
@@ -106,9 +106,8 @@
     }
    ],
    "source": [
-    "from sentence_transformers import SentenceTransformer\n",
-    "\n",
     "from mmcontext.mmcontextencoder import MMContextEncoder\n",
+    "from sentence_transformers import SentenceTransformer\n",
     "\n",
     "model_name = \"poster_models/scvi_model_ct_2048\"\n",
     "model_path = f\"../../{model_name}/0_MMContextEncoder\"  # Point to the encoder subdirectory\n",
diff --git a/exploratory/notebooks/mmcontext_concept_figs.ipynb b/exploratory/notebooks/mmcontext_concept_figs.ipynb
index 8d1f8cf..9881a80 100644
--- a/exploratory/notebooks/mmcontext_concept_figs.ipynb
+++ b/exploratory/notebooks/mmcontext_concept_figs.ipynb
@@ -48,9 +48,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sentence_transformers import SentenceTransformer\n",
-    "\n",
     "from mmcontext.mmcontextencoder import MMContextEncoder\n",
+    "from sentence_transformers import SentenceTransformer\n",
     "\n",
     "enc = MMContextEncoder(\n",
     "    \"NeuML/pubmedbert-base-embeddings\",\n",
diff --git a/exploratory/notebooks/omics_sentence_transformer_loss_demo.ipynb b/exploratory/notebooks/omics_sentence_transformer_loss_demo.ipynb
index 33e8be6..0260d28 100644
--- a/exploratory/notebooks/omics_sentence_transformer_loss_demo.ipynb
+++ b/exploratory/notebooks/omics_sentence_transformer_loss_demo.ipynb
@@ -32,6 +32,7 @@
     "from pathlib import Path\n",
     "\n",
     "import numpy as np\n",
+    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from sentence_transformers import (\n",
     "    SentenceTransformer,\n",
     "    SentenceTransformerTrainer,\n",
@@ -40,7 +41,6 @@
     ")\n",
     "from tqdm.autonotebook import tqdm\n",
     "\n",
-    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from mmcontext.sanity_helpers import cluster_variances, plot_pca, stack_embeddings\n",
     "from mmcontext.simulator import OmicsCaptionSimulator, make_cluster_sampler\n",
     "\n",
diff --git a/exploratory/notebooks/omics_sentence_transformer_loss_report.ipynb b/exploratory/notebooks/omics_sentence_transformer_loss_report.ipynb
index b4974ee..ef2f54b 100644
--- a/exploratory/notebooks/omics_sentence_transformer_loss_report.ipynb
+++ b/exploratory/notebooks/omics_sentence_transformer_loss_report.ipynb
@@ -69,6 +69,7 @@
     "import numpy as np\n",
     "import pandas as pd\n",
     "import torch\n",
+    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from sentence_transformers import (\n",
     "    SentenceTransformer,\n",
     "    SentenceTransformerTrainer,\n",
@@ -77,7 +78,6 @@
     ")\n",
     "from tqdm.autonotebook import tqdm\n",
     "\n",
-    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
     "from mmcontext.sanity_helpers import cluster_variances, plot_pca, stack_embeddings  # provided separately\n",
     "from mmcontext.simulator import OmicsCaptionSimulator, make_cluster_sampler\n",
     "\n",
diff --git a/exploratory/notebooks/text_batches.ipynb b/exploratory/notebooks/text_batches.ipynb
index ee02a81..1ea6b4f 100644
--- a/exploratory/notebooks/text_batches.ipynb
+++ b/exploratory/notebooks/text_batches.ipynb
@@ -186,9 +186,10 @@
    ],
    "source": [
     "# Visualise the embeddings\n",
-    "from mmcontext.pl import plot_umap\n",
     "from mmcontext.pp.utils import consolidate_low_frequency_categories\n",
     "\n",
+    "from mmcontext.pl import plot_umap\n",
+    "\n",
     "adata_cut = consolidate_low_frequency_categories(adata, [\"cell_type\"], threshold=10)\n",
     "label_key = \"_scvi_batch\"\n",
     "plot_umap(adata_cut, color_key=label_key, embedding_key=\"X_text\")"
diff --git a/exploratory/notebooks/text_only_concept_figs.ipynb b/exploratory/notebooks/text_only_concept_figs.ipynb
index 095a841..582bf86 100644
--- a/exploratory/notebooks/text_only_concept_figs.ipynb
+++ b/exploratory/notebooks/text_only_concept_figs.ipynb
@@ -78,9 +78,8 @@
     }
    ],
    "source": [
-    "from sentence_transformers import SentenceTransformer\n",
-    "\n",
     "from mmcontext.mmcontextencoder import MMContextEncoder\n",
+    "from sentence_transformers import SentenceTransformer\n",
     "\n",
     "enc = MMContextEncoder(\n",
     "    \"NeuML/pubmedbert-base-embeddings\",\n",
diff --git a/pyproject.toml b/pyproject.toml
index f9f8030..e6761f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,12 +1,11 @@
 [build-system]
 build-backend = "setuptools.build_meta"
-
-requires = [ "setuptools>=64", "wheel" ]
+requires = [ "setuptools>=64" ]
 
 [project]
 name = "mmcontext"
-version = "0.0.1"
-description = "This package allows to use the sentence-transfomers package to build multimodal embedding models for data in anndata format."
+version = "2.0.0"
+description = "Multimodal contrastive learning aligning text and omics embeddings via sentence-transformers."
 readme = "README.md"
 license = { file = "LICENSE" }
 maintainers = [
@@ -17,59 +16,61 @@ authors = [
 ]
 requires-python = ">=3.11,<3.14"
 classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Science/Research",
   "Programming Language :: Python :: 3 :: Only",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
+  "Topic :: Scientific/Engineering :: Bio-Informatics",
 ]
-
 dependencies = [
   "accelerate>=1.9",
-  "adata-hf-datasets @ git+https://github.com/mengerj/adata_hf_datasets.git",
-  "anndata>=0.11.3",
-  "datasets>=2.19.1,!=4",                                                     #Avoid: TypeError: Wrong key type: '0' of type '<class 'numpy.int64'>'. Expected one of int, slice, range, str or Iterable.
+  "anndata>=0.12",
+  "datasets>=2.19.1,!=4",
   "huggingface-hub>=0.34.3",
-  "hydra-core>=1.3.2",
-  "ipykernel>=6.30",
-  "louvain>=0.8.2",
-  "psutil>=7",
+  "numpy>=1.24",
   "python-dotenv>=1.1.1",
   "scanpy>=1.11.3",
-  "scib>=1.1.7",
-  "sentence-transformers>=5",
+  "sentence-transformers>=5.4",
   "torch>=2.5",
   "transformers>=4.57.1",
-  "wandb>=0.21",
-  "zarr<3",
+  "zarr>=3.1.6",
 ]
 optional-dependencies.dev = [
+  "build",
   "pre-commit",
+  "ruff",
   "twine",
-  # other dev dependencies...
+]
+optional-dependencies.eval = [
+  "scib>=1.1.7",
+  "umap-learn>=0.5",
 ]
 optional-dependencies.test = [
   "coverage",
   "pytest",
   "python-json-logger",
 ]
+optional-dependencies.train = [
+  "wandb>=0.27",
+]
 urls.Homepage = "https://github.com/mengerj/mmcontext"
+urls.Issues = "https://github.com/mengerj/mmcontext/issues"
 urls.Source = "https://github.com/mengerj/mmcontext"
 
 [tool.setuptools]
+packages.find.where = [ "src" ]
+packages.find.include = [ "mmcontext*" ]
 include-package-data = true
-[tool.setuptools.packages.find]
-where = [ "src" ]
-include = [ "mmcontext*" ]
-[tool.setuptools.package-data]
-#"mmcontext.conf" = [ "*.json" ]
+package-data = {}
 
+# "mmcontext.conf" = [ "*.json" ]
 [tool.ruff]
 line-length = 120
 src = [ "src" ]
 extend-include = [ "*.ipynb" ]
-
 format.docstring-code-format = true
-
 lint.select = [
   "B",      # flake8-bugbear
   "C4",     # flake8-comprehensions
@@ -82,7 +83,6 @@ lint.select = [
   "UP",     # pyupgrade
   "W",      # Warning detected by Pycodestyle
 ]
-
 lint.ignore = [
   "B008",   # function calls in argument defaults
   "BLE001", # allow blind expections
@@ -99,7 +99,6 @@ lint.ignore = [
   "E741",   # allow I, O, l as variable names
   "F401",   # dont remove unused imports, because of fiass flag
 ]
-
 lint.per-file-ignores."*/__init__.py" = [ "F401" ]
 lint.per-file-ignores."tests/*" = [ "D" ]
 lint.pydocstyle.convention = "numpy"
@@ -110,18 +109,18 @@ ignore = [
   "UP031",
 ]
 
-[tool.pytest.ini_options]
-testpaths = [ "tests" ]
-xfail_strict = true
-addopts = [
+[tool.pytest]
+ini_options.testpaths = [ "tests" ]
+ini_options.addopts = [
   "--import-mode=importlib",
 ]
+ini_options.xfail_strict = true
 
-[tool.coverage.run]
-source = [ "mmcontext" ]
-omit = [
+[tool.coverage]
+run.omit = [
   "**/test_*.py",
 ]
+run.source = [ "mmcontext" ]
 
 [tool.cruft]
 skip = [
@@ -129,6 +128,3 @@ skip = [
   "src/**/__init__.py",
   "src/**/basic.py",
 ]
-
-[tool.uv.sources]
-adata-hf-datasets = { git = "https://github.com/mengerj/adata_hf_datasets.git" }
diff --git a/scripts/embed_eval.py b/scripts/embed_eval.py
deleted file mode 100644
index 060be9a..0000000
--- a/scripts/embed_eval.py
+++ /dev/null
@@ -1,482 +0,0 @@
-import logging
-import shutil
-from pathlib import Path
-from typing import Any
-
-import hydra
-from omegaconf import DictConfig
-
-from mmcontext.embed.cellwhisperer_utils import process_cellwhisperer_dataset_model as cellwhisperer_single
-from mmcontext.embed.embed_pipeline import process_single_dataset_model as embed_single
-from mmcontext.eval.eval_pipeline import process_single_dataset_model as eval_single
-from mmcontext.eval.eval_pipeline import run_scib_evaluation
-from mmcontext.utils import copy_resolved_config
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
-    datefmt="%H:%M:%S",
-)
-
-
-def process_combined_dataset_model(
-    ds_cfg: Any,
-    model_cfg: Any,
-    cfg: DictConfig,
-) -> tuple[str, str, bool, str, int]:
-    """
-    Process a single dataset/model combination through both embedding and evaluation steps.
-
-    Parameters
-    ----------
-    ds_cfg : dict
-        Dataset configuration
-    model_cfg : dict
-        Model configuration
-    cfg : DictConfig
-        Full configuration object with embed, eval, and settings sections
-
-    Returns
-    -------
-    tuple[str, str, bool, str, int]
-        Tuple of (dataset_name, model_name, success, error_msg, eval_result_count)
-    """
-    dataset_name = ds_cfg.name
-    model_name = model_cfg.get("name", model_cfg.source)
-    text_only = model_cfg.get("text_only", False)
-
-    if text_only:
-        model_name_for_path = model_name + "_text_only"
-    else:
-        model_name_for_path = model_name
-
-    logger.info(f"Processing combined pipeline for {dataset_name} + {model_name_for_path}")
-
-    try:
-        # ═══════════════════════════════════════════════════════════════
-        # STEP 0: CHECK IF EVALUATION ALREADY EXISTS (BEFORE EMBEDDING)
-        # ═══════════════════════════════════════════════════════════════
-        final_eval_dir = Path(cfg.settings.final_root) / dataset_name / model_name_for_path / "eval"
-        skip_existing_eval = cfg.settings.get("skip_existing_evaluations", False)
-
-        if skip_existing_eval and final_eval_dir.exists():
-            logger.info(
-                f"Evaluation already exists for {dataset_name}/{model_name_for_path} in final directory. Skipping entire pipeline."
-            )
-            return dataset_name, model_name_for_path, True, "evaluation_skipped_existing", 0
-
-        # ═══════════════════════════════════════════════════════════════
-        # STEP 1: EMBEDDING GENERATION
-        # ═══════════════════════════════════════════════════════════════
-        logger.info(f"Step 1/2: Generating embeddings for {dataset_name}/{model_name_for_path}")
-
-        embed_dataset_name, embed_model_name, embed_success, embed_error = embed_single(
-            ds_cfg=ds_cfg,
-            model_cfg=model_cfg,
-            run_cfg=cfg.embed,
-            output_root=cfg.settings.computation_root,
-            output_format=cfg.settings.output_format,
-            adata_cache=cfg.settings.adata_cache,
-            hf_cache=cfg.settings.hf_cache,
-        )
-
-        if not embed_success:
-            if embed_error == "skipped_existing":
-                logger.info(f"Embeddings already exist for {dataset_name}/{model_name_for_path}")
-            else:
-                logger.error(f"Embedding failed for {dataset_name}/{model_name_for_path}: {embed_error}")
-                return dataset_name, model_name_for_path, False, f"Embedding failed: {embed_error}", 0
-        else:
-            logger.info(f"✓ Embedding completed for {dataset_name}/{model_name_for_path}")
-
-        # ═══════════════════════════════════════════════════════════════
-        # STEP 1.5: CELLWHISPERER PROCESSING (if enabled)
-        # ═══════════════════════════════════════════════════════════════
-        if cfg.settings.get("run_cellwhisperer", False):
-            logger.info(f"Step 1.5/2: Running CellWhisperer for {dataset_name}")
-
-            # Create CellWhisperer model config
-            from omegaconf import DictConfig
-
-            cellwhisperer_model_cfg = DictConfig({"name": "cellwhisperer", "source": "cellwhisperer"})
-
-            modules_dir = Path(
-                cfg.settings.get("cellwhisperer_modules_dir", cfg.settings.computation_root + "/../modules")
-            )
-
-            cw_dataset_name, cw_model_name, cw_success, cw_error = cellwhisperer_single(
-                ds_cfg=ds_cfg,
-                model_cfg=cellwhisperer_model_cfg,
-                run_cfg=cfg.embed,
-                output_root=cfg.settings.computation_root,
-                output_format=cfg.settings.output_format,
-                adata_cache=cfg.settings.adata_cache,
-                modules_dir=modules_dir,
-                hf_cache=cfg.settings.hf_cache,
-            )
-
-            if not cw_success:
-                if cw_error == "skipped_existing":
-                    logger.info(f"CellWhisperer embeddings already exist for {dataset_name}")
-                else:
-                    logger.warning(f"CellWhisperer failed for {dataset_name}: {cw_error}")
-                    # Don't fail the entire pipeline if CellWhisperer fails
-            else:
-                logger.info(f"✓ CellWhisperer completed for {dataset_name}")
-
-                # Also run evaluation for CellWhisperer
-                logger.info(f"Running evaluation for CellWhisperer: {dataset_name}")
-                try:
-                    _cw_eval_results = eval_single(
-                        ds_cfg=ds_cfg,
-                        model_cfg=cellwhisperer_model_cfg,
-                        eval_cfg=cfg.eval,
-                        output_root=cfg.settings.computation_root,
-                        modules_dir=modules_dir,
-                    )
-
-                    # Move CellWhisperer evaluation results to final directory
-                    cw_computation_eval_dir = (
-                        Path(cfg.settings.computation_root) / dataset_name / "cellwhisperer" / "eval"
-                    )
-                    cw_final_eval_dir = Path(cfg.settings.final_root) / dataset_name / "cellwhisperer" / "eval"
-
-                    if cw_computation_eval_dir.exists():
-                        logger.info(f"Moving CellWhisperer evaluation results to {cw_final_eval_dir}")
-                        cw_final_eval_dir.parent.mkdir(parents=True, exist_ok=True)
-                        if cw_final_eval_dir.exists():
-                            shutil.rmtree(cw_final_eval_dir)
-                        shutil.move(str(cw_computation_eval_dir), str(cw_final_eval_dir))
-                        logger.info("✓ Moved CellWhisperer evaluation results")
-
-                except Exception as e:
-                    logger.warning(f"CellWhisperer evaluation failed for {dataset_name}: {e}")
-                    # Continue with main pipeline even if CellWhisperer evaluation fails
-
-        # ═══════════════════════════════════════════════════════════════
-        # STEP 2: EVALUATION
-        # ═══════════════════════════════════════════════════════════════
-        logger.info(f"Step 2/2: Running evaluation for {dataset_name}/{model_name_for_path}")
-
-        # Run evaluation using computation directory as source
-        # Get modules_dir from settings if available (for CellWhisperer logit_scale extraction)
-        modules_dir = (
-            Path(cfg.settings.get("cellwhisperer_modules_dir", cfg.settings.computation_root + "/../modules"))
-            if cfg.settings.get("cellwhisperer_modules_dir")
-            else None
-        )
-
-        eval_results = eval_single(
-            ds_cfg=ds_cfg,
-            model_cfg=model_cfg,
-            eval_cfg=cfg.eval,
-            output_root=cfg.settings.computation_root,  # Read embeddings from computation directory
-            modules_dir=modules_dir,
-        )
-
-        # ═══════════════════════════════════════════════════════════════
-        # STEP 3: MOVE EVALUATION RESULTS TO FINAL DIRECTORY
-        # ═══════════════════════════════════════════════════════════════
-        computation_eval_dir = Path(cfg.settings.computation_root) / dataset_name / model_name_for_path / "eval"
-
-        if computation_eval_dir.exists():
-            logger.info(f"Moving evaluation results from {computation_eval_dir} to {final_eval_dir}")
-
-            # Create final directory structure
-            final_eval_dir.parent.mkdir(parents=True, exist_ok=True)
-
-            # Move evaluation directory to final location
-            if final_eval_dir.exists():
-                shutil.rmtree(final_eval_dir)  # Remove existing if present
-            shutil.move(str(computation_eval_dir), str(final_eval_dir))
-
-            logger.info(f"✓ Moved evaluation results to {final_eval_dir}")
-        else:
-            logger.warning(f"No evaluation directory found at {computation_eval_dir}")
-
-        logger.info(f"✓ Combined pipeline completed for {dataset_name}/{model_name_for_path}")
-        return dataset_name, model_name_for_path, True, None, len(eval_results)
-
-    except Exception as e:
-        logger.error(f"✗ Error in combined pipeline for {dataset_name}/{model_name_for_path}: {e}")
-        return dataset_name, model_name_for_path, False, str(e), 0
-
-
-def _process_combined_worker(args):
-    """
-    Worker function for parallel processing of combined pipeline.
-
-    Parameters
-    ----------
-    args : tuple
-        Tuple of (ds_cfg, model_cfg, cfg)
-
-    Returns
-    -------
-    tuple[str, str, bool, str, int]
-        Result from process_combined_dataset_model
-    """
-    import os
-    import sys
-    import time
-    import traceback
-
-    ds_cfg, model_cfg, cfg = args
-
-    dataset_name = ds_cfg.name
-    model_name = model_cfg.get("name", model_cfg.source)
-    text_only = model_cfg.get("text_only", False)
-    if text_only:
-        model_name_for_path = model_name + "_text_only"
-    else:
-        model_name_for_path = model_name
-
-    # Configure logging for this worker process
-    worker_logger = logging.getLogger(f"combined_worker_{os.getpid()}")
-    worker_logger.setLevel(logging.INFO)
-
-    # Create a handler that writes to stdout so we can see worker logs
-    if not worker_logger.handlers:
-        handler = logging.StreamHandler(sys.stdout)
-        handler.setFormatter(logging.Formatter(f"[Worker {os.getpid()}] %(levelname)s - %(message)s"))
-        worker_logger.addHandler(handler)
-
-    worker_logger.info(f"Starting combined processing: {dataset_name}/{model_name_for_path}")
-    start_time = time.time()
-
-    try:
-        result = process_combined_dataset_model(ds_cfg, model_cfg, cfg)
-
-        elapsed_time = time.time() - start_time
-        worker_logger.info(
-            f"Completed combined processing: {dataset_name}/{model_name_for_path} in {elapsed_time:.1f}s"
-        )
-
-        return result
-
-    except Exception as e:
-        # Get full traceback for debugging
-        elapsed_time = time.time() - start_time
-        error_msg = f"{str(e)}\n{traceback.format_exc()}"
-        worker_logger.error(
-            f"Error in combined processing {dataset_name}/{model_name_for_path} after {elapsed_time:.1f}s: {error_msg}"
-        )
-        return dataset_name, model_name_for_path, False, str(e), 0
-
-
-def combined_pipeline(cfg: DictConfig) -> None:
-    """
-    Orchestrate the combined embedding + evaluation pipeline.
-
-    Parameters
-    ----------
-    cfg : omegaconf.DictConfig
-        Instantiated Hydra config matching combined_conf.yaml schema.
-    """
-    import concurrent.futures
-    import multiprocessing as mp
-    import time
-    from concurrent.futures import ProcessPoolExecutor, as_completed
-
-    print("\n=== Running Combined Embedding + Evaluation Pipeline ===")
-
-    # Create list of all dataset/model combinations
-    tasks = []
-    for ds_cfg in cfg.datasets:
-        for model_cfg in cfg.models:
-            tasks.append((ds_cfg, model_cfg, cfg))
-
-    print(f"Total dataset/model combinations to process: {len(tasks)}")
-
-    # Create output directories
-    Path(cfg.settings.computation_root).mkdir(parents=True, exist_ok=True)
-    Path(cfg.settings.final_root).mkdir(parents=True, exist_ok=True)
-
-    # Check if parallel processing is enabled
-    enable_parallel = cfg.settings.get("enable_parallel", False)
-    max_workers = cfg.settings.get("max_workers", None)
-
-    if max_workers is None:
-        max_workers = mp.cpu_count()
-
-    start_time = time.time()
-
-    if not enable_parallel or len(tasks) == 1:
-        # Sequential processing
-        print("Running in sequential mode...")
-        successful = 0
-        failed = 0
-        skipped = 0
-
-        for i, (ds_cfg, model_cfg, task_cfg) in enumerate(tasks):
-            dataset_name = ds_cfg.name
-            model_name = model_cfg.get("name", model_cfg.source)
-            text_only = model_cfg.get("text_only", False)
-            if text_only:
-                model_name_for_path = model_name + "_text_only"
-            else:
-                model_name_for_path = model_name
-
-            task_start = time.time()
-            print(f"\n=== [{i + 1}/{len(tasks)}] Processing: {dataset_name}/{model_name_for_path} ===")
-
-            try:
-                dataset_name_result, model_name_result, success, error_msg, eval_count = process_combined_dataset_model(
-                    ds_cfg, model_cfg, task_cfg
-                )
-                task_time = time.time() - task_start
-
-                if success:
-                    if error_msg in ["skipped_existing", "evaluation_skipped_existing"]:
-                        skipped += 1
-                        print(f"↷ Skipped: {dataset_name_result}/{model_name_result} ({task_time:.1f}s)")
-                    else:
-                        successful += 1
-                        print(f"✓ Completed: {dataset_name_result}/{model_name_result} ({task_time:.1f}s)")
-                        print(f"  → {eval_count} evaluation metrics")
-                else:
-                    failed += 1
-                    print(f"✗ Failed: {dataset_name_result}/{model_name_result} ({task_time:.1f}s)")
-                    print(f"  Error: {error_msg}")
-
-                # Estimate remaining time
-                if i > 0:
-                    avg_time = (time.time() - start_time) / (i + 1)
-                    remaining_time = avg_time * (len(tasks) - i - 1)
-                    print(f"  → Estimated time remaining: {remaining_time:.1f}s")
-
-            except Exception as e:
-                failed += 1
-                print(f"✗ Error processing {dataset_name}/{model_name_for_path}: {e}")
-                print("  Continuing with next combination...")
-                continue
-
-        total_time = time.time() - start_time
-        print("\n=== Summary ===")
-        print(f"Total tasks: {len(tasks)}")
-        print(f"Successful: {successful}")
-        print(f"Skipped: {skipped}")
-        print(f"Failed: {failed}")
-        print(f"Success rate: {successful / len(tasks) * 100:.1f}%")
-        print(f"Total time: {total_time:.1f}s")
-        print(f"Average time per task: {total_time / len(tasks):.1f}s")
-
-    else:
-        # Parallel processing
-        print(f"Running in parallel mode with {max_workers} workers...")
-
-        # Get timeout from config
-        task_timeout = cfg.settings.get("task_timeout", 10800)  # 3 hours default
-        print(f"Task timeout: {task_timeout}s ({task_timeout / 60:.1f} minutes)")
-
-        # Use ProcessPoolExecutor for CPU processing
-        print("Using ProcessPoolExecutor for CPU processing")
-        executor_class = ProcessPoolExecutor
-        worker_func = _process_combined_worker
-        task_args_func = lambda task: (task,)
-
-        with executor_class(max_workers=max_workers) as executor:
-            # Submit all tasks
-            future_to_task = {}
-            for i, (ds_cfg, model_cfg, task_cfg) in enumerate(tasks):
-                dataset_name = ds_cfg.name
-                model_name = model_cfg.get("name", model_cfg.source)
-                text_only = model_cfg.get("text_only", False)
-                if text_only:
-                    model_name_for_path = model_name + "_text_only"
-                else:
-                    model_name_for_path = model_name
-
-                future = executor.submit(worker_func, *task_args_func((ds_cfg, model_cfg, task_cfg)))
-                future_to_task[future] = (dataset_name, model_name_for_path, i, time.time())
-
-            # Process completed tasks with timeout
-            completed = 0
-            successful = 0
-            failed = 0
-            skipped = 0
-            timed_out = 0
-
-            try:
-                for future in as_completed(future_to_task, timeout=task_timeout + 60):
-                    dataset_name, model_name_for_path, task_idx, submit_time = future_to_task[future]
-                    completed += 1
-                    task_time = time.time() - submit_time
-
-                    try:
-                        dataset_name_result, model_name_result, success, error_msg, eval_count = future.result(
-                            timeout=10
-                        )
-
-                        if success:
-                            if error_msg in ["skipped_existing", "evaluation_skipped_existing"]:
-                                skipped += 1
-                                print(
-                                    f"↷ [{completed}/{len(tasks)}] Skipped: {dataset_name_result}/{model_name_result} ({task_time:.1f}s)"
-                                )
-                            else:
-                                successful += 1
-                                print(
-                                    f"✓ [{completed}/{len(tasks)}] Completed: {dataset_name_result}/{model_name_result} ({task_time:.1f}s)"
-                                )
-                                print(f"  → {eval_count} evaluation metrics")
-                        else:
-                            failed += 1
-                            print(
-                                f"✗ [{completed}/{len(tasks)}] Failed: {dataset_name_result}/{model_name_result} ({task_time:.1f}s)"
-                            )
-                            print(f"  Error: {error_msg}")
-
-                        # Progress update
-                        if completed % max(1, len(tasks) // 10) == 0:
-                            elapsed = time.time() - start_time
-                            if completed > 0:
-                                eta = (elapsed / completed) * (len(tasks) - completed)
-                                print(f"  → Progress: {completed / len(tasks) * 100:.1f}% | ETA: {eta:.1f}s")
-
-                    except concurrent.futures.TimeoutError:
-                        timed_out += 1
-                        failed += 1
-                        print(
-                            f"✗ [{completed}/{len(tasks)}] Timeout: {dataset_name}/{model_name_for_path} (>{task_timeout}s)"
-                        )
-                        future.cancel()
-
-                    except Exception as e:
-                        failed += 1
-                        print(
-                            f"✗ [{completed}/{len(tasks)}] Exception in worker for {dataset_name}/{model_name_for_path}: {e}"
-                        )
-
-            except concurrent.futures.TimeoutError:
-                print("⚠ Warning: Overall timeout reached. Some tasks may not have completed.")
-                remaining_tasks = len(tasks) - completed
-                if remaining_tasks > 0:
-                    print(f"⚠ Warning: {remaining_tasks} tasks did not complete")
-                    failed += remaining_tasks
-
-            total_time = time.time() - start_time
-            print("\n=== Summary ===")
-            print(f"Total tasks: {len(tasks)}")
-            print(f"Successful: {successful}")
-            print(f"Skipped: {skipped}")
-            print(f"Failed: {failed}")
-            if timed_out > 0:
-                print(f"Timed out: {timed_out}")
-            print(f"Success rate: {successful / len(tasks) * 100:.1f}%")
-            print(f"Total time: {total_time:.1f}s")
-            print(f"Average time per task: {total_time / len(tasks):.1f}s")
-            print(f"Speedup vs sequential: ~{max_workers}x (theoretical)")
-
-    print("\n=== Combined pipeline completed ===")
-
-
-@hydra.main(config_path="../conf/eval", config_name="embed_eval_conf")
-def main(cfg: DictConfig) -> None:
-    """Entry point launched by Hydra."""
-    combined_pipeline(cfg)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/explore_hard_negatives.py b/scripts/explore_hard_negatives.py
new file mode 100644
index 0000000..71ca5f9
--- /dev/null
+++ b/scripts/explore_hard_negatives.py
@@ -0,0 +1,210 @@
+#!/usr/bin/env python
+"""Explore which hard negatives a trained model mines for each config dataset.
+
+This is an *offline inspection* tool, not a training script. It reuses the
+``scripts/train_config.py`` machinery (``load_config``, ``load_raw_datasets``,
+``build_merged_vector_store``) to set up exactly the datasets + merged
+VectorStore a real run would use, then mines model-specific hard negatives per
+dataset with :func:`mmcontext.embed.mine_negatives` and writes a readable report
+(:func:`mmcontext.embed.mining_report`) of what got chosen and how similar it is
+to the true positive.
+
+Mining is run **per dataset** (not over a merged corpus) because the
+sentence-transformers trainer samples each batch from a single dataset, so
+in-batch and mined negatives should both stay within-dataset.
+
+You must point at an *already-trained* model — mining against a freshly
+initialised bimodal model produces noise.
+
+Usage::
+
+    python scripts/explore_hard_negatives.py \
+        --config conf/training/multi_example.yaml \
+        --model jo-mengr/<some-trained-model> \
+        --output-dir outputs/hard_neg_exploration \
+        --num-negatives 3 --range-min 1 --max-score 0.95
+
+Outputs (under --output-dir):
+    <dataset>__mined.csv     per-anchor positive + mined negatives + sims/margins
+    summary.md               score distributions + a few examples per dataset
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+from pathlib import Path
+
+import pandas as pd
+from dotenv import load_dotenv
+from sentence_transformers import SentenceTransformer
+
+from mmcontext.embed import MiningConfig, mine_negatives, mining_report, prepare_dataset
+from mmcontext.training import (
+    DatasetConfig,
+    TrainConfig,
+    _namespace_sample_ids,
+    _select_training_columns,
+    build_merged_vector_store,
+    load_config,
+    load_raw_datasets,
+)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Explore mined hard negatives for the config's datasets")
+    parser.add_argument("--config", required=True, help="Path to the YAML training config")
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="HF repo / local path of an ALREADY-TRAINED model to mine with "
+        "(mining a fresh bimodal model produces noise).",
+    )
+    parser.add_argument("--output-dir", default="outputs/hard_neg_exploration", help="Where to write the report")
+    parser.add_argument(
+        "--max-pairs",
+        type=int,
+        default=None,
+        help="Subsample each dataset to the first N anchor/positive pairs before "
+        "mining (also shrinks the corpus). Default: use the full dataset.",
+    )
+    parser.add_argument(
+        "--report-rows",
+        type=int,
+        default=50,
+        help="Rows per dataset to recompute similarities for in the CSV (default: 50).",
+    )
+    parser.add_argument("--overwrite-store", action="store_true", help="Rebuild cached vector stores")
+
+    mining = parser.add_argument_group("hard-negative mining")
+    mining.add_argument("--num-negatives", type=int, default=3)
+    mining.add_argument("--range-min", type=int, default=1)
+    mining.add_argument("--range-max", type=int, default=None)
+    mining.add_argument("--max-score", type=float, default=0.95)
+    mining.add_argument("--min-score", type=float, default=None)
+    mining.add_argument("--absolute-margin", type=float, default=None)
+    mining.add_argument("--relative-margin", type=float, default=None)
+    mining.add_argument("--sampling-strategy", choices=["top", "random"], default="top")
+    mining.add_argument("--use-faiss", action=argparse.BooleanOptionalAction, default=False)
+    mining.add_argument("--batch-size", type=int, default=32)
+    return parser.parse_args()
+
+
+def _prepare_omics_pairs(dcfg: DatasetConfig, dsdict):
+    """Return an anchor/positive ``datasets.Dataset`` for one omics dataset."""
+    train_split = dsdict["train"]
+    if dcfg.modality == "bimodal":
+        train_split = _namespace_sample_ids(train_split, dcfg.name)
+    # Drop the dataset's own pre-baked negatives so we mine fresh, model-specific ones.
+    return prepare_dataset(
+        train_split,
+        purpose="train",
+        modality=dcfg.modality,
+        primary_cell_sentence=dcfg.primary_cell_sentence,
+        positive_col=dcfg.positive_col,
+        use_hard_negatives=False,
+    )
+
+
+def _summarise(name: str, report: pd.DataFrame) -> dict:
+    """Compute score/margin stats for one dataset's report."""
+    neg_sim_cols = [c for c in report.columns if c.startswith("neg_") and c.endswith("_sim")]
+    margin_cols = [c for c in report.columns if c.startswith("margin_")]
+    neg_sims = report[neg_sim_cols].to_numpy().ravel() if neg_sim_cols else []
+    margins = report[margin_cols].to_numpy().ravel() if margin_cols else []
+    stat = {"dataset": name, "n_reported": len(report)}
+    if len(report):
+        stat["pos_sim_mean"] = round(float(report["pos_sim"].mean()), 4)
+    if len(neg_sims):
+        stat["neg_sim_mean"] = round(float(neg_sims.mean()), 4)
+        stat["neg_sim_max"] = round(float(neg_sims.max()), 4)
+    if len(margins):
+        stat["margin_mean"] = round(float(margins.mean()), 4)
+        stat["margin_min"] = round(float(margins.min()), 4)
+    return stat
+
+
+def main() -> None:
+    """Mine and report hard negatives for every dataset in the config."""
+    args = _parse_args()
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    cfg: TrainConfig = load_config(args.config)
+    logger.info("Loaded config from %s", args.config)
+
+    omics_raw, bio_raw = load_raw_datasets(cfg)
+    overwrite = args.overwrite_store or cfg.overwrite_vector_store
+    merged_store, _ = build_merged_vector_store(cfg, omics_raw, overwrite=overwrite)
+
+    logger.info("Loading model %s", args.model)
+    model = SentenceTransformer(args.model, trust_remote_code=True)
+    if merged_store is not None:
+        model[0].set_vector_store(merged_store)
+        logger.info("Attached merged VectorStore: %d vectors, dim=%d", len(merged_store), merged_store.dim)
+
+    mining_cfg = MiningConfig(
+        num_negatives=args.num_negatives,
+        range_min=args.range_min,
+        range_max=args.range_max,
+        max_score=args.max_score,
+        min_score=args.min_score,
+        absolute_margin=args.absolute_margin,
+        relative_margin=args.relative_margin,
+        sampling_strategy=args.sampling_strategy,
+        use_faiss=args.use_faiss,
+        batch_size=args.batch_size,
+    )
+    logger.info("Mining config: %s", mining_cfg)
+
+    # (name, anchor/positive pairs) for every dataset in the config.
+    jobs: list[tuple[str, object]] = []
+    for dcfg in cfg.omics_datasets:
+        jobs.append((dcfg.name, _prepare_omics_pairs(dcfg, omics_raw[dcfg.name])))
+    for dcfg in cfg.bio_datasets:
+        if "train" not in bio_raw[dcfg.name]:
+            logger.warning("Bio dataset '%s' has no train split, skipping", dcfg.name)
+            continue
+        jobs.append((dcfg.name, _select_training_columns(bio_raw[dcfg.name]["train"], dcfg.name)))
+
+    stats: list[dict] = []
+    for name, pairs in jobs:
+        if args.max_pairs is not None and len(pairs) > args.max_pairs:
+            pairs = pairs.select(range(args.max_pairs))
+        logger.info("[%s] mining over %d pairs", name, len(pairs))
+        mined = mine_negatives(pairs, model, mining_cfg, output_format="n-tuple", verbose=True)
+        report = mining_report(mined, model, max_rows=args.report_rows)
+        csv_path = out_dir / f"{name}__mined.csv"
+        report.to_csv(csv_path, index=False)
+        logger.info("[%s] wrote %s (%d rows; %d mined total)", name, csv_path, len(report), len(mined))
+        stats.append({**_summarise(name, report), "n_pairs": len(pairs), "n_mined": len(mined)})
+
+    _write_summary(out_dir / "summary.md", mining_cfg, stats)
+    logger.info("Done. Report written to %s", out_dir)
+
+
+def _write_summary(path: Path, mining_cfg: MiningConfig, stats: list[dict]) -> None:
+    """Write a markdown summary with the score distributions and example rows."""
+    lines = ["# Hard-negative mining exploration", "", "## Mining config", "", "```", repr(mining_cfg), "```", ""]
+    summary_table = pd.DataFrame(stats).to_string(index=False) if stats else "(no datasets)"
+    lines += ["## Per-dataset summary", "", "```", summary_table, "```", ""]
+    lines += [
+        "## Reading the numbers",
+        "",
+        "- `pos_sim`: cosine(anchor, true positive). Higher = the model already aligns the pair.",
+        "- `neg_*_sim`: cosine(anchor, mined negative). Close to `pos_sim` ⇒ a genuinely hard negative.",
+        "- `margin_*` = `pos_sim − neg_sim`. Small/negative margins are the hardest (and the most",
+        "  likely false negatives — check that `max_score` is filtering those out).",
+        "",
+        "Per-dataset CSVs with every reported anchor/positive/negative are in this directory.",
+    ]
+    path.write_text("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/finetune_tiny.py b/scripts/finetune_tiny.py
new file mode 100644
index 0000000..cb202ac
--- /dev/null
+++ b/scripts/finetune_tiny.py
@@ -0,0 +1,326 @@
+#!/usr/bin/env python
+"""Finetune an already-trained MMContext model with mined hard negatives.
+
+Unlike :mod:`scripts.train_tiny`, this script does **not** use the
+pre-provided hard negatives baked into the dataset (the ``negative_*_idx``
+columns). Instead it loads a model that was already trained once and uses
+:func:`sentence_transformers.util.mine_hard_negatives` to mine model-specific
+hard negatives, then continues training on those. This only makes sense with a
+model that already has a meaningful joint embedding space — mining against a
+freshly initialised bimodal model would produce noise.
+
+How mining works here: the dataset's ``anchor`` column (omics ids for the
+bimodal modality, resolved through the attached VectorStore) is encoded as the
+query, while the ``positive`` texts form the corpus. The closest-but-not-true
+positives become the mined negatives — confusable text descriptions in the same
+shared embedding space.
+
+Usage::
+
+    # 1) Produce a pretrained model with train_tiny first
+    python scripts/train_tiny.py --epochs 1 --output-dir outputs/tiny_model
+
+    # 2) Finetune it with mined hard negatives (bimodal, default)
+    python scripts/finetune_tiny.py \
+        --model-path outputs/tiny_model/final \
+        --output-dir outputs/tiny_finetuned
+
+    # Text modality
+    python scripts/finetune_tiny.py --modality text \
+        --model-path outputs/tiny_text/final \
+        --output-dir outputs/tiny_text_finetuned
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
+from sentence_transformers.util import mine_hard_negatives
+
+from mmcontext.embed import prepare_dataset
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Dataset helpers
+# ---------------------------------------------------------------------------
+HF_DATASET = "jo-mengr/cxg_schaefer_tiny"
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    """Finetune a pretrained MMContext model on mined hard negatives."""
+    parser = argparse.ArgumentParser(description="Finetune MMContext with mined hard negatives")
+    parser.add_argument(
+        "--modality",
+        choices=["text", "bimodal"],
+        default="bimodal",
+        help="Training modality (default: bimodal)",
+    )
+    parser.add_argument(
+        "--model-path",
+        default="outputs/tiny_model/final",
+        help="Path to the pretrained model to finetune (default: outputs/tiny_model/final).",
+    )
+    parser.add_argument(
+        "--vector-store",
+        default=None,
+        help="Path to .mmap VectorStore file. For bimodal modality: if omitted, "
+        "the store is built automatically from adata_link + sample_idx "
+        "columns using --obsm-key.",
+    )
+    parser.add_argument(
+        "--obsm-key",
+        default="X_scvi_fm",
+        help="obsm key to extract when building VectorStore (default: X_scvi_fm). "
+        "Other common choices: X_pca, X_geneformer, X_gs10k",
+    )
+    parser.add_argument("--output-dir", default="outputs/tiny_finetuned", help="Output directory")
+    parser.add_argument("--epochs", type=int, default=3, help="Training epochs (default: 3)")
+    parser.add_argument("--batch-size", type=int, default=32, help="Batch size (default: 32)")
+    parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate (default: 2e-5)")
+    parser.add_argument("--dataset", default=HF_DATASET, help="HuggingFace dataset name")
+    parser.add_argument(
+        "--use-mps-device",
+        action=argparse.BooleanOptionalAction,
+        default=torch.backends.mps.is_available(),
+        help="Train on Apple Metal (MPS); default: on when MPS is available",
+    )
+    parser.add_argument(
+        "--wandb-project",
+        default=None,
+        help="Weights & Biases project name. Enables wandb logging when set. "
+        "You can also set WANDB_PROJECT env var instead.",
+    )
+    parser.add_argument(
+        "--wandb-run-name",
+        default=None,
+        help="Optional W&B run name (auto-generated if omitted)",
+    )
+    parser.add_argument(
+        "--freeze-text-encoder",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Freeze the text encoder to cut gradient/optimizer memory. "
+        "On by default for finetuning; combine with --unfreeze-last-n to keep "
+        "the top N layers trainable. Disable with --no-freeze-text-encoder.",
+    )
+    parser.add_argument(
+        "--unfreeze-last-n",
+        type=int,
+        default=1,
+        help="With --freeze-text-encoder, keep the top N transformer layers "
+        "(plus pooler) trainable. 0 freezes the whole encoder (default: 1).",
+    )
+    parser.add_argument(
+        "--bf16",
+        action=argparse.BooleanOptionalAction,
+        default=torch.backends.mps.is_available(),
+        help="Use bf16 mixed precision. Supported on MPS (macOS 14+) and CUDA; "
+        "roughly halves activation memory. Default: on when MPS is available.",
+    )
+
+    # --- Hard-negative mining parameters ---
+    mining = parser.add_argument_group("hard-negative mining")
+    mining.add_argument("--num-negatives", type=int, default=3, help="Negatives to mine per anchor (default: 3)")
+    mining.add_argument(
+        "--range-min",
+        type=int,
+        default=1,
+        help="Skip the N closest candidates (the closest is usually the true positive). Default: 1.",
+    )
+    mining.add_argument(
+        "--range-max",
+        type=int,
+        default=None,
+        help="Only consider candidates up to this rank (default: None = no upper bound).",
+    )
+    mining.add_argument(
+        "--max-score",
+        type=float,
+        default=0.95,
+        help="Drop candidates whose similarity exceeds this ceiling to avoid mining "
+        "false negatives (descriptions that are actually correct). Default: 0.95.",
+    )
+    mining.add_argument(
+        "--relative-margin",
+        type=float,
+        default=None,
+        help="Negatives must score below relative_margin * positive_score (e.g. 0.95). Default: None.",
+    )
+    mining.add_argument(
+        "--sampling-strategy",
+        choices=["top", "random"],
+        default="top",
+        help="Which qualifying candidates to keep as negatives (default: top).",
+    )
+    mining.add_argument(
+        "--use-faiss",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Use FAISS for the similarity search (recommended for large datasets). Default: off.",
+    )
+    args = parser.parse_args()
+
+    # --- Dataset (raw) ---
+    ds_raw = load_dataset(args.dataset, split="train")
+
+    # --- Load pretrained model ---
+    # Dimensions (text/omics/shared) come from the saved config; no rebuild needed.
+    logger.info("Loading pretrained model from %s", args.model_path)
+    model = SentenceTransformer(args.model_path)
+
+    # --- Attach VectorStore (bimodal only), required before mining ---
+    if args.modality == "bimodal":
+        from mmcontext.io import VectorStore, prepare_vector_store
+
+        if args.vector_store is not None:
+            store = VectorStore.load(args.vector_store)
+        else:
+            store_path = os.path.join(args.output_dir, "vector_store.mmap")
+            store = prepare_vector_store(
+                ds_raw,
+                obsm_key=args.obsm_key,
+                output_path=store_path,
+            )
+        model[0].set_vector_store(store)
+        logger.info("VectorStore: %d vectors, dim=%d", len(store), store.dim)
+
+    # --- Freezing (memory savings) ---
+    # Freezing the text encoder removes its gradients + Adam optimizer state;
+    # keeping only the top N layers trainable is the usual fine-tuning sweet
+    # spot. On by default here since we are continuing to train an aligned model.
+    if args.freeze_text_encoder:
+        text_module = model[0]
+        # unfreeze_last_n == 0 freezes the whole encoder; > 0 keeps the top N trainable.
+        text_module.freeze_all_but_top_layers(args.unfreeze_last_n)
+        logger.info("Froze text encoder, keeping top %d layers trainable", args.unfreeze_last_n)
+        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in model.parameters())
+        logger.info("Trainable params: %d / %d (%.1f%%)", trainable, total, 100.0 * trainable / total)
+
+    # --- Prepare anchor/positive pairs (drop the dataset's own negatives) ---
+    modality = "bimodal" if args.modality == "bimodal" else "text"
+    primary_cell_sentence = "cell_sentence_1" if modality == "bimodal" else "cell_sentence_2"
+    pairs = prepare_dataset(
+        ds_raw,
+        purpose="train",
+        modality=modality,
+        primary_cell_sentence=primary_cell_sentence,
+        use_hard_negatives=False,
+    )
+    logger.info("Prepared %d anchor/positive pairs, columns=%s", len(pairs), pairs.column_names)
+
+    # --- Mine hard negatives with the pretrained model ---
+    model.eval()
+    logger.info(
+        "Mining hard negatives: num_negatives=%d, range_min=%d, range_max=%s, max_score=%s, "
+        "relative_margin=%s, sampling=%s, faiss=%s",
+        args.num_negatives,
+        args.range_min,
+        args.range_max,
+        args.max_score,
+        args.relative_margin,
+        args.sampling_strategy,
+        args.use_faiss,
+    )
+    mined = mine_hard_negatives(
+        pairs,
+        model,
+        anchor_column_name="anchor",
+        positive_column_name="positive",
+        num_negatives=args.num_negatives,
+        range_min=args.range_min,
+        range_max=args.range_max,
+        max_score=args.max_score,
+        relative_margin=args.relative_margin,
+        sampling_strategy=args.sampling_strategy,
+        output_format="n-tuple",
+        batch_size=args.batch_size,
+        use_faiss=args.use_faiss,
+        verbose=True,
+    )
+    logger.info("Mined dataset: %d rows (from %d pairs), columns=%s", len(mined), len(pairs), mined.column_names)
+
+    # --- Wandb setup ---
+    wandb_project = args.wandb_project or os.environ.get("WANDB_PROJECT")
+    use_wandb = wandb_project is not None
+    if use_wandb:
+        os.environ["WANDB_PROJECT"] = wandb_project
+        if args.wandb_run_name:
+            os.environ["WANDB_NAME"] = args.wandb_run_name
+        logger.info("W&B enabled: project=%s, run=%s", wandb_project, args.wandb_run_name or "(auto)")
+
+    # --- Finetuning ---
+    model.train()
+    loss = MultipleNegativesRankingLoss(model)
+    training_args = SentenceTransformerTrainingArguments(
+        output_dir=args.output_dir,
+        num_train_epochs=args.epochs,
+        per_device_train_batch_size=args.batch_size,
+        learning_rate=args.lr,
+        warmup_ratio=0.1,
+        # bf16 (MPS macOS 14+ / CUDA) takes precedence; fall back to fp16 on CUDA only.
+        bf16=args.bf16,
+        fp16=(not args.bf16) and torch.cuda.is_available() and not args.use_mps_device,
+        use_mps_device=args.use_mps_device,
+        report_to="wandb" if use_wandb else "none",
+        logging_steps=10,
+        save_strategy="epoch",
+        save_total_limit=2,
+        run_name=args.wandb_run_name,
+    )
+
+    trainer = SentenceTransformerTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=mined,
+        loss=loss,
+    )
+
+    logger.info(
+        "Starting finetuning: modality=%s, epochs=%d, batch_size=%d", args.modality, args.epochs, args.batch_size
+    )
+    trainer.train()
+
+    # --- Save final model ---
+    save_path = os.path.join(args.output_dir, "final")
+    model.save(save_path)
+    logger.info("Model saved to %s", save_path)
+
+    # --- Verify reload ---
+    logger.info("Verifying save/load roundtrip...")
+    model.eval()
+    test_inputs = ["MALAT1 MT-CO3 GNAS SYT1 CALM1", "A cortical neuron expressing synaptic markers."]
+    original_embs = model.encode(test_inputs)
+
+    loaded = SentenceTransformer(save_path)
+    loaded.eval()
+    loaded_embs = loaded.encode(test_inputs)
+
+    max_diff = np.abs(original_embs - loaded_embs).max()
+    logger.info("Save/load verification: max_diff=%.2e (should be < 1e-5)", max_diff)
+    if max_diff > 1e-4:
+        logger.warning("Save/load roundtrip difference is unexpectedly large!")
+    else:
+        logger.info("Save/load roundtrip OK")
+
+    logger.info("Done.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/run_scsa.py b/scripts/run_scsa.py
deleted file mode 100644
index 756f824..0000000
--- a/scripts/run_scsa.py
+++ /dev/null
@@ -1,116 +0,0 @@
-"""
-Standalone SCSA evaluation pipeline.
-
-Runs the SCSA (Single-Cell RNA-seq Annotation) baseline method on the same
-datasets used by the embedding pipeline and writes metrics in the identical
-format so that ``collect_metrics.py`` can fuse them with other model results.
-
-Usage::
-
-    python scripts/run_scsa.py                     # use defaults from scsa_conf.yaml
-    python scripts/run_scsa.py settings.output_root=/path/to/results
-"""
-
-import logging
-import time
-from pathlib import Path
-
-import hydra
-from omegaconf import DictConfig
-
-from mmcontext.embed.scsa_utils import ensure_scsa_setup, process_scsa_dataset
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
-    datefmt="%H:%M:%S",
-)
-
-
-@hydra.main(config_path="../conf/eval", config_name="scsa_conf", version_base=None)
-def main(cfg: DictConfig) -> None:
-    """Run SCSA baseline on every dataset in the shared dataset list."""
-    print("\n=== Running SCSA Baseline Evaluation ===")
-
-    modules_dir = Path(cfg.settings.modules_dir)
-    if not ensure_scsa_setup(modules_dir):
-        print("ERROR: SCSA setup failed. Run modules/prepare_scsa.sh manually.")
-        return
-
-    datasets = list(cfg.datasets)
-    print(f"Datasets to process: {len(datasets)}")
-    for ds in datasets:
-        print(f"  - {ds.name}")
-
-    Path(cfg.settings.output_root).mkdir(parents=True, exist_ok=True)
-
-    start_time = time.time()
-    successful = 0
-    failed = 0
-    skipped = 0
-    needs_review = 0
-    all_calmate_messages: list[str] = []
-
-    for i, ds_cfg in enumerate(datasets):
-        task_start = time.time()
-        print(f"\n=== [{i + 1}/{len(datasets)}] Processing: {ds_cfg.name} ===")
-
-        dataset_name, success, error_msg, calmate_msgs = process_scsa_dataset(
-            ds_cfg=ds_cfg,
-            run_cfg=cfg.embed,
-            output_root=cfg.settings.output_root,
-            modules_dir=modules_dir,
-            scsa_cfg=cfg.scsa,
-            adata_cache=cfg.settings.adata_cache,
-            hf_cache=cfg.settings.get("hf_cache", None),
-            skip_existing=cfg.settings.get("skip_existing", True),
-            temp_dir=cfg.settings.get("temp_dir", None),
-        )
-
-        task_time = time.time() - task_start
-
-        if calmate_msgs:
-            needs_review += 1
-            all_calmate_messages.extend([f"[{dataset_name}]"] + calmate_msgs + [""])
-
-        if success:
-            if error_msg == "skipped_existing":
-                skipped += 1
-                print(f"  Skipped (already exists): {dataset_name} ({task_time:.1f}s)")
-            else:
-                successful += 1
-                print(f"  Completed: {dataset_name} ({task_time:.1f}s)")
-        else:
-            failed += 1
-            print(f"  FAILED: {dataset_name} ({task_time:.1f}s)")
-            print(f"    Error: {error_msg}")
-
-    total_time = time.time() - start_time
-    print("\n=== SCSA Evaluation Summary ===")
-    print(f"Total datasets: {len(datasets)}")
-    print(f"Successful:     {successful}")
-    print(f"Skipped:        {skipped}")
-    print(f"Failed:         {failed}")
-    print(f"Total time:     {total_time:.1f}s")
-    print(f"Output root:    {cfg.settings.output_root}")
-
-    if all_calmate_messages:
-        calmate_store = cfg.scsa.get("calmate_store", "modules/.calmate/mappings.csv")
-        calmate_cli = "modules/calmate_venv/bin/calmate"
-        print(f"\n{'=' * 60}")
-        print("  LABEL MAPPING REVIEW REQUIRED")
-        print(f"{'=' * 60}")
-        print(f"\n{needs_review} dataset(s) have unreviewed or missing label mappings.")
-        print("\nDetails:")
-        for msg in all_calmate_messages:
-            print(f"  {msg}")
-        print("To review interactively:")
-        print(f"  {calmate_cli} --store {calmate_store} review")
-        print("\nThen re-run:")
-        print("  python scripts/run_scsa.py settings.skip_existing=false")
-        print(f"{'=' * 60}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/scripts/train.py b/scripts/train.py
index b54c0ba..13db283 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -19,11 +19,11 @@
 from sentence_transformers.evaluation import SequentialEvaluator
 from transformers.integrations import WandbCallback
 
+from mmcontext._legacy.mmcontextencoder import MMContextEncoder
 from mmcontext.callback import UnfreezeAdapterCallback, UnfreezeTextEncoderCallback
 
 # from mmcontext.eval import SystemMonitor
 from mmcontext.hub_utils import prepare_model_for_hub, upload_model_to_hub
-from mmcontext.mmcontextencoder import MMContextEncoder
 
 # from mmcontext.pp.utils import consolidate_low_frequency_categories
 from mmcontext.utils import (  # , load_test_adata_from_hf_dataset
diff --git a/scripts/train_config.py b/scripts/train_config.py
new file mode 100644
index 0000000..dba7aa3
--- /dev/null
+++ b/scripts/train_config.py
@@ -0,0 +1,304 @@
+#!/usr/bin/env python
+"""Config-driven multi-dataset training for the MMContext ST pipeline.
+
+Unlike ``train_tiny.py`` (single dataset, argparse), this script reads a plain
+YAML config (no Hydra) and supports multiple omics + bio datasets, mixing
+bimodal (omics-vector) and text anchors in one run. It reuses the existing
+multi-dataset machinery: ``prepare_dataset``, ``prepare_vector_store``,
+``get_loss`` / ``get_evaluator``, the unfreeze callbacks, and the Hub helpers.
+
+When continuing from an already-trained model (``model:`` set), an optional
+``mining:`` block enables hard-negative mining (as in ``scripts/finetune_tiny.py``):
+each dataset's own negatives are dropped and model-specific negatives are mined
+with :func:`sentence_transformers.util.mine_hard_negatives`. The mining settings
+are recorded in the pushed model card.
+
+Usage::
+
+    python scripts/train_config.py --config conf/training/multi_example.yaml
+
+See ``conf/training/multi_example.yaml`` for a documented config.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+from dotenv import load_dotenv
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.evaluation import SequentialEvaluator
+from transformers.integrations import WandbCallback
+
+from mmcontext.callback import UnfreezeAdapterCallback, UnfreezeTextEncoderCallback
+from mmcontext.embed import build_pipeline
+from mmcontext.modules import MMContextModule
+from mmcontext.training import (
+    TrainConfig,
+    assemble_training_data,
+    build_merged_vector_store,
+    format_training_details,
+    load_config,
+    load_raw_datasets,
+)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+
+def build_or_load_model(cfg: TrainConfig, omics_dim: int | None) -> SentenceTransformer:
+    """Build a fresh pipeline or load one to continue fine-tuning.
+
+    Parameters
+    ----------
+    cfg : TrainConfig
+        Run config.
+    omics_dim : int or None
+        Dimensionality of the merged VectorStore (``None`` for text-only runs).
+
+    Returns
+    -------
+    SentenceTransformer
+    """
+    if cfg.model:
+        logger.info("Loading existing model to continue training: %s", cfg.model)
+        model = SentenceTransformer(cfg.model, trust_remote_code=True)
+        if not isinstance(model[0], MMContextModule):
+            raise TypeError(f"Loaded model's first module is {type(model[0]).__name__}, expected MMContextModule")
+        if omics_dim is not None:
+            adapter_dim = model[1].omics_input_dim
+            if adapter_dim != omics_dim:
+                raise ValueError(
+                    f"Loaded adapter expects omics_input_dim={adapter_dim} but the merged "
+                    f"VectorStore has dim={omics_dim}. Check obsm_key matches the loaded model."
+                )
+        return model
+
+    return build_pipeline(
+        text_model=cfg.text_encoder.name,
+        omics_dim=omics_dim,
+        shared_dim=cfg.shared_dim,
+        adapter_hidden_dim=cfg.adapter_hidden_dim,
+        max_seq_length=cfg.text_encoder.max_seq_length,
+    )
+
+
+def register_gene_special_token(model: SentenceTransformer, token: str) -> None:
+    """Add *token* to the tokenizer and resize the encoder embeddings."""
+    tokenizer = model[0].tokenizer
+    num_added = tokenizer.add_tokens([token])
+    if num_added > 0:
+        model[0].auto_model.resize_token_embeddings(len(tokenizer))
+        logger.info("Added special token %r and resized embeddings to %d", token, len(tokenizer))
+    else:
+        logger.info("Special token %r already present in tokenizer", token)
+
+
+def main() -> None:
+    """Run config-driven multi-dataset training."""
+    parser = argparse.ArgumentParser(description="Config-driven MMContext training")
+    parser.add_argument("--config", required=True, help="Path to the YAML training config")
+    parser.add_argument(
+        "--overwrite-store",
+        action="store_true",
+        help="Rebuild per-dataset vector stores even if cached .mmap files exist "
+        "(also settable via overwrite_vector_store in the config)",
+    )
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    logger.info("Loaded config from %s", args.config)
+
+    # --- W&B ---
+    use_wandb = cfg.trainer.wandb_project is not None
+    if use_wandb:
+        os.environ["WANDB_PROJECT"] = cfg.trainer.wandb_project
+        if cfg.trainer.wandb_run_name:
+            os.environ["WANDB_NAME"] = cfg.trainer.wandb_run_name
+        logger.info("W&B enabled: project=%s", cfg.trainer.wandb_project)
+
+    # --- Data + merged vector store ---
+    omics_raw, bio_raw = load_raw_datasets(cfg)
+    overwrite_store = args.overwrite_store or cfg.overwrite_vector_store
+    merged_store, omics_dim = build_merged_vector_store(cfg, omics_raw, overwrite=overwrite_store)
+
+    # --- Model ---
+    model = build_or_load_model(cfg, omics_dim)
+    if merged_store is not None:
+        model[0].set_vector_store(merged_store)
+    if cfg.gene_special_token:
+        register_gene_special_token(model, cfg.gene_special_token)
+
+    # --- Hard-negative mining sanity check ---
+    mining_enabled = cfg.mining is not None and cfg.mining.enabled
+    if mining_enabled and not cfg.model:
+        logger.warning(
+            "Hard-negative mining is enabled but no `model:` was provided. Mining against a "
+            "freshly initialised model produces near-random negatives — set `model:` to a "
+            "pretrained checkpoint to continue fine-tuning, or disable mining."
+        )
+    if mining_enabled:
+        logger.info("Hard-negative mining enabled; dataset-provided negatives will be replaced by mined ones.")
+
+    # --- Static text-encoder freezing ---
+    if cfg.freeze_text_encoder:
+        model[0].freeze_all_but_top_layers(cfg.unfreeze_last_n_layers)
+        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in model.parameters())
+        logger.info("Trainable params: %d / %d (%.1f%%)", trainable, total, 100.0 * trainable / total)
+
+    # --- Assemble datasets / losses / evaluators ---
+    assembled = assemble_training_data(
+        cfg,
+        model,
+        omics_raw,
+        bio_raw,
+        log_backend="wandb" if use_wandb else "auto",
+    )
+    if not assembled.train_datasets:
+        raise ValueError("No training datasets were assembled — check omics_datasets / bio_datasets in the config.")
+    logger.info("Training datasets: %s", list(assembled.train_datasets))
+
+    dev_evaluator = SequentialEvaluator(assembled.evaluators) if assembled.evaluators else None
+    eval_datasets = assembled.eval_datasets or None
+    eval_strategy = cfg.trainer.eval_strategy
+    eval_steps = cfg.trainer.eval_steps
+    if not eval_datasets or dev_evaluator is None:
+        logger.info("No eval datasets — disabling evaluation during training")
+        eval_strategy, eval_steps = "no", None
+
+    # --- Training arguments ---
+    training_args = SentenceTransformerTrainingArguments(
+        output_dir=cfg.trainer.output_dir,
+        num_train_epochs=cfg.trainer.num_train_epochs,
+        per_device_train_batch_size=cfg.trainer.per_device_train_batch_size,
+        per_device_eval_batch_size=cfg.trainer.per_device_eval_batch_size,
+        learning_rate=cfg.trainer.learning_rate,
+        warmup_ratio=cfg.trainer.warmup_ratio,
+        bf16=cfg.trainer.bf16,
+        fp16=cfg.trainer.fp16,
+        use_mps_device=cfg.trainer.use_mps_device,
+        eval_strategy=eval_strategy,
+        eval_steps=eval_steps,
+        save_strategy=cfg.trainer.save_strategy,
+        save_steps=cfg.trainer.save_steps,
+        save_total_limit=cfg.trainer.save_total_limit,
+        max_grad_norm=cfg.trainer.max_grad_norm,
+        logging_steps=cfg.trainer.logging_steps,
+        dataloader_num_workers=cfg.trainer.dataloader_num_workers,
+        report_to="wandb" if use_wandb else "none",
+        run_name=cfg.trainer.wandb_run_name,
+    )
+    if cfg.trainer.use_cosine_scheduler:
+        training_args.set_lr_scheduler(name="cosine", warmup_ratio=cfg.trainer.warmup_ratio)
+
+    if cfg.trainer.gradient_checkpointing:
+        # Enable gradient flow for input_ids so partial unfreezing still trains.
+        model[0].auto_model.enable_input_require_grads()
+        model[0].auto_model.gradient_checkpointing_enable()
+
+    # --- Callbacks ---
+    callbacks = [UnfreezeTextEncoderCallback(unfreeze_epoch=cfg.unfreeze_epoch)]
+    if use_wandb:
+        callbacks.append(WandbCallback())
+    if cfg.adapter_freezing is not None:
+        callbacks.append(
+            UnfreezeAdapterCallback(
+                freeze_text_adapter=cfg.adapter_freezing.freeze_text_adapter,
+                freeze_omics_adapter=cfg.adapter_freezing.freeze_omics_adapter,
+                unfreeze_text_adapter_epoch=cfg.adapter_freezing.unfreeze_text_adapter_epoch,
+                unfreeze_omics_adapter_epoch=cfg.adapter_freezing.unfreeze_omics_adapter_epoch,
+            )
+        )
+
+    # --- Train ---
+    trainer = SentenceTransformerTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=assembled.train_datasets,
+        eval_dataset=eval_datasets,
+        loss=assembled.losses,
+        evaluator=dev_evaluator,
+        callbacks=callbacks,
+    )
+    logger.info("Starting training on %d dataset(s)", len(assembled.train_datasets))
+    trainer.train()
+
+    # --- Save ---
+    # Tag the model card before saving so the tag lands in the README frontmatter.
+    if mining_enabled:
+        model.model_card_data.add_tags(["hard-negative-mining"])
+    save_path = Path(cfg.trainer.output_dir) / "final"
+    model.save(str(save_path))
+    logger.info("Model saved to %s", save_path)
+
+    # --- Save/load roundtrip sanity check ---
+    verify_roundtrip(model, save_path)
+
+    # --- Hub upload ---
+    if cfg.push_to_hub:
+        push_to_hub(cfg, model, save_path)
+
+    logger.info("Done.")
+
+
+def verify_roundtrip(model: SentenceTransformer, save_path: Path) -> None:
+    """Encode a probe with the in-memory and reloaded model; warn on drift."""
+    model.eval()
+    probe = ["MALAT1 MT-CO3 GNAS SYT1 CALM1", "A cortical neuron expressing synaptic markers."]
+    original = model.encode(probe)
+    loaded = SentenceTransformer(str(save_path), trust_remote_code=True)
+    loaded.eval()
+    reloaded = loaded.encode(probe)
+    max_diff = float(np.abs(original - reloaded).max())
+    if max_diff > 1e-4:
+        logger.warning("Save/load roundtrip difference is large: max_diff=%.2e", max_diff)
+    else:
+        logger.info("Save/load roundtrip OK (max_diff=%.2e)", max_diff)
+
+
+def push_to_hub(cfg: TrainConfig, model: SentenceTransformer, save_path: Path) -> None:
+    """Upload the trained model to the Hub under a unique name.
+
+    The saved model directory's ``README.md`` is augmented with a training-details
+    section (including any hard-negative mining settings) and that local folder is
+    uploaded verbatim, so the card on the Hub records how the model was finetuned.
+    """
+    from mmcontext.hub_utils import append_model_card_section, generate_unique_model_name, get_hf_username
+
+    try:
+        username = get_hf_username()
+        model_name = generate_unique_model_name(
+            cfg.text_encoder.name,
+            obsm_key=cfg.obsm_key if cfg.omics_datasets else None,
+            tag=cfg.tag,
+            username=username,
+        )
+        repo_id = f"{username}/{model_name}"
+
+        # Record how the model was (fine)tuned in the model card before pushing.
+        append_model_card_section(save_path, format_training_details(cfg))
+
+        logger.info("Pushing model to Hub: %s", repo_id)
+        model.push_to_hub(repo_id, private=True, exist_ok=True, local_model_path=str(save_path))
+        logger.info("Uploaded model to https://huggingface.co/%s", repo_id)
+    except Exception as e:
+        logger.warning("Hub upload failed: %s. The model is saved locally and can be pushed manually.", e)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as e:
+        logger.exception(e)
+        sys.exit(1)
diff --git a/scripts/train_merged.py b/scripts/train_merged.py
index 06fc0b9..3b2433c 100644
--- a/scripts/train_merged.py
+++ b/scripts/train_merged.py
@@ -29,10 +29,10 @@
 )
 from transformers.integrations import WandbCallback
 
+from mmcontext._legacy.mmcontextencoder import MMContextEncoder
 from mmcontext.callback import UnfreezeAdapterCallback, UnfreezeTextEncoderCallback
 from mmcontext.eval import SystemMonitor
 from mmcontext.hub_utils import get_model_info_from_config, upload_model_to_hub
-from mmcontext.mmcontextencoder import MMContextEncoder
 from mmcontext.utils import (
     get_evaluator,
     get_loss,
diff --git a/scripts/train_tiny.py b/scripts/train_tiny.py
new file mode 100644
index 0000000..923cf67
--- /dev/null
+++ b/scripts/train_tiny.py
@@ -0,0 +1,267 @@
+#!/usr/bin/env python
+"""Train an MMContext model on the cxg_schaefer_tiny dataset.
+
+Quick-start training script using sentence-transformers v5.4+ pipeline.
+Supports two modalities:
+
+  1. **text** — gene-name strings as anchors, text
+     descriptions as positives.  No VectorStore needed.
+  2. **bimodal** — omics vectors from a VectorStore as anchors, text as
+     positives.  Requires ``--vector-store`` pointing to a ``.mmap`` file
+     that covers the ``sample_idx`` column of the dataset.
+
+Usage::
+
+    # Text modality
+    python scripts/train_tiny.py --output-dir outputs/tiny_text
+
+    # Bimodal modality (default)
+    python scripts/train_tiny.py --modality bimodal \
+        --vector-store /path/to/store.mmap \
+        --output-dir outputs/tiny_bimodal
+
+    # Custom text encoder
+    python scripts/train_tiny.py --text-model dmis-lab/biobert-base-cased-v1.2
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
+
+from mmcontext.embed import build_pipeline, prepare_dataset
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Dataset helpers
+# ---------------------------------------------------------------------------
+HF_DATASET = "jo-mengr/cxg_schaefer_tiny"
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    """Train MMContext on cxg_schaefer_tiny dataset."""
+    parser = argparse.ArgumentParser(description="Train MMContext on cxg_schaefer_tiny")
+    parser.add_argument(
+        "--modality",
+        choices=["text", "bimodal"],
+        default="bimodal",
+        help="Training modality (default: bimodal)",
+    )
+    parser.add_argument(
+        "--text-model",
+        default="NeuML/pubmedbert-base-embeddings",
+        help="HuggingFace text encoder (default: NeuML/pubmedbert-base-embeddings)",
+    )
+    parser.add_argument(
+        "--vector-store",
+        default=None,
+        help="Path to .mmap VectorStore file. For bimodal modality: if omitted, "
+        "the store is built automatically from adata_link + sample_idx "
+        "columns using --obsm-key.",
+    )
+    parser.add_argument(
+        "--obsm-key",
+        default="X_scvi_fm",
+        help="obsm key to extract when building VectorStore (default: X_scvi_fm). "
+        "Other common choices: X_pca, X_geneformer, X_gs10k",
+    )
+    parser.add_argument("--omics-dim", type=int, default=None, help="Omics vector dimension")
+    parser.add_argument("--shared-dim", type=int, default=256, help="Shared embedding dim (default: 256)")
+    parser.add_argument("--output-dir", default="outputs/tiny_model", help="Output directory")
+    parser.add_argument("--epochs", type=int, default=3, help="Training epochs (default: 3)")
+    parser.add_argument("--batch-size", type=int, default=32, help="Batch size (default: 8)")
+    parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate (default: 2e-5)")
+    parser.add_argument("--dataset", default=HF_DATASET, help="HuggingFace dataset name")
+    parser.add_argument(
+        "--use-mps-device",
+        action=argparse.BooleanOptionalAction,
+        default=torch.backends.mps.is_available(),
+        help="Train on Apple Metal (MPS); default: on when MPS is available",
+    )
+    parser.add_argument(
+        "--wandb-project",
+        default=None,
+        help="Weights & Biases project name. Enables wandb logging when set. "
+        "You can also set WANDB_PROJECT env var instead.",
+    )
+    parser.add_argument(
+        "--wandb-run-name",
+        default=None,
+        help="Optional W&B run name (auto-generated if omitted)",
+    )
+    parser.add_argument(
+        "--hard-negatives",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Use hard negatives from negative_*_idx dataset columns (default: True).",
+    )
+    parser.add_argument(
+        "--freeze-text-encoder",
+        action="store_true",
+        help="Freeze the text encoder to cut gradient/optimizer memory. "
+        "Combine with --unfreeze-last-n to keep the top N layers trainable.",
+    )
+    parser.add_argument(
+        "--unfreeze-last-n",
+        type=int,
+        default=0,
+        help="With --freeze-text-encoder, keep the top N transformer layers "
+        "(plus pooler) trainable. 0 freezes the whole encoder (default: 0).",
+    )
+    parser.add_argument(
+        "--bf16",
+        action=argparse.BooleanOptionalAction,
+        default=torch.backends.mps.is_available(),
+        help="Use bf16 mixed precision. Supported on MPS (macOS 14+) and CUDA; "
+        "roughly halves activation memory. Default: on when MPS is available.",
+    )
+    args = parser.parse_args()
+
+    # --- Dataset ---
+    ds_raw = load_dataset(args.dataset, split="train")
+    # genelist mode uses cell-sentence text anchors; bimodal uses omics ids.
+    # The primary cell sentence is chosen by modality: cell_sentence_1
+    # (gene-list) for bimodal, cell_sentence_2 (text description) for text.
+    modality = "bimodal" if args.modality == "bimodal" else "text"
+    primary_cell_sentence = "cell_sentence_1" if modality == "bimodal" else "cell_sentence_2"
+    ds = prepare_dataset(
+        ds_raw,
+        purpose="train",
+        modality=modality,
+        primary_cell_sentence=primary_cell_sentence,
+        use_hard_negatives=args.hard_negatives,
+    )
+
+    # --- Pipeline ---
+    pipeline = build_pipeline(
+        text_model=args.text_model,
+        omics_dim=args.omics_dim,
+        shared_dim=args.shared_dim,
+    )
+
+    # Attach VectorStore for bimodal modality
+    if args.modality == "bimodal":
+        from mmcontext.io import VectorStore, prepare_vector_store
+
+        if args.vector_store is not None:
+            store = VectorStore.load(args.vector_store)
+        else:
+            # Auto-build from adata_link column
+            store_path = os.path.join(args.output_dir, "vector_store.mmap")
+            store = prepare_vector_store(
+                ds_raw,
+                obsm_key=args.obsm_key,
+                output_path=store_path,
+            )
+
+        first_module = list(pipeline.children())[0]
+        first_module.set_vector_store(store)
+
+        # Infer omics_dim if not set
+        if args.omics_dim is None:
+            args.omics_dim = store.dim
+            # Rebuild pipeline with correct omics_dim
+            pipeline = build_pipeline(
+                text_model=args.text_model,
+                omics_dim=args.omics_dim,
+                shared_dim=args.shared_dim,
+            )
+            first_module = list(pipeline.children())[0]
+            first_module.set_vector_store(store)
+
+        logger.info("VectorStore: %d vectors, dim=%d", len(store), store.dim)
+
+    # --- Freezing (memory savings) ---
+    # Applied to the final pipeline's MMContextModule. Freezing the text
+    # encoder removes its gradients + Adam optimizer state; keeping only the
+    # top N layers trainable is the usual fine-tuning sweet spot.
+    if args.freeze_text_encoder:
+        text_module = list(pipeline.children())[0]
+        # unfreeze_last_n == 0 freezes the whole encoder; > 0 keeps the top N trainable.
+        text_module.freeze_all_but_top_layers(args.unfreeze_last_n)
+        logger.info("Froze text encoder, keeping top %d layers trainable", args.unfreeze_last_n)
+        trainable = sum(p.numel() for p in pipeline.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in pipeline.parameters())
+        logger.info("Trainable params: %d / %d (%.1f%%)", trainable, total, 100.0 * trainable / total)
+
+    # --- Wandb setup ---
+    wandb_project = args.wandb_project or os.environ.get("WANDB_PROJECT")
+    use_wandb = wandb_project is not None
+    if use_wandb:
+        os.environ["WANDB_PROJECT"] = wandb_project
+        if args.wandb_run_name:
+            os.environ["WANDB_NAME"] = args.wandb_run_name
+        logger.info("W&B enabled: project=%s, run=%s", wandb_project, args.wandb_run_name or "(auto)")
+
+    # --- Training ---
+    loss = MultipleNegativesRankingLoss(pipeline)
+    training_args = SentenceTransformerTrainingArguments(
+        output_dir=args.output_dir,
+        num_train_epochs=args.epochs,
+        per_device_train_batch_size=args.batch_size,
+        learning_rate=args.lr,
+        warmup_ratio=0.1,
+        # bf16 (MPS macOS 14+ / CUDA) takes precedence; fall back to fp16 on CUDA only.
+        bf16=args.bf16,
+        fp16=(not args.bf16) and torch.cuda.is_available() and not args.use_mps_device,
+        use_mps_device=args.use_mps_device,
+        report_to="wandb" if use_wandb else "none",
+        logging_steps=10,
+        save_strategy="epoch",
+        save_total_limit=2,
+        run_name=args.wandb_run_name,
+    )
+
+    trainer = SentenceTransformerTrainer(
+        model=pipeline,
+        args=training_args,
+        train_dataset=ds,
+        loss=loss,
+    )
+
+    logger.info("Starting training: modality=%s, epochs=%d, batch_size=%d", args.modality, args.epochs, args.batch_size)
+    trainer.train()
+
+    # --- Save final model ---
+    save_path = os.path.join(args.output_dir, "final")
+    pipeline.save(save_path)
+    logger.info("Model saved to %s", save_path)
+
+    # --- Verify reload ---
+    logger.info("Verifying save/load roundtrip...")
+    pipeline.eval()
+    test_inputs = ["MALAT1 MT-CO3 GNAS SYT1 CALM1", "A cortical neuron expressing synaptic markers."]
+    original_embs = pipeline.encode(test_inputs)
+
+    loaded = SentenceTransformer(save_path)
+    loaded.eval()
+    loaded_embs = loaded.encode(test_inputs)
+
+    max_diff = np.abs(original_embs - loaded_embs).max()
+    logger.info("Save/load verification: max_diff=%.2e (should be < 1e-5)", max_diff)
+    if max_diff > 1e-4:
+        logger.warning("Save/load roundtrip difference is unexpectedly large!")
+    else:
+        logger.info("Save/load roundtrip OK")
+
+    logger.info("Done.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/mmcontext/__init__.py b/src/mmcontext/__init__.py
index 6cc09d2..ac4b245 100644
--- a/src/mmcontext/__init__.py
+++ b/src/mmcontext/__init__.py
@@ -1,7 +1,7 @@
 from importlib.metadata import version
 
-from . import embed, eval, models, pl
+from . import embed, eval, io, modules, pl
 
-__all__ = ["models", "eval", "pl", "embed"]
+__all__ = ["embed", "eval", "io", "modules", "pl"]
 
 __version__ = version("mmcontext")
diff --git a/src/mmcontext/_legacy/__init__.py b/src/mmcontext/_legacy/__init__.py
new file mode 100644
index 0000000..582fdeb
--- /dev/null
+++ b/src/mmcontext/_legacy/__init__.py
@@ -0,0 +1,9 @@
+"""Legacy modules from the pre-v5.4 architecture.
+
+These modules are preserved for reference and backward compatibility with
+existing trained models. New code should use :mod:`mmcontext.modules` and
+:mod:`mmcontext.io` instead.
+
+.. deprecated::
+    This package will be removed in a future version.
+"""
diff --git a/src/mmcontext/adapters.py b/src/mmcontext/_legacy/adapters.py
similarity index 100%
rename from src/mmcontext/adapters.py
rename to src/mmcontext/_legacy/adapters.py
diff --git a/src/mmcontext/models/cell_sentence_transformer.py b/src/mmcontext/_legacy/cell_sentence_transformer.py
similarity index 100%
rename from src/mmcontext/models/cell_sentence_transformer.py
rename to src/mmcontext/_legacy/cell_sentence_transformer.py
diff --git a/src/mmcontext/mmcontextencoder.py b/src/mmcontext/_legacy/mmcontextencoder.py
similarity index 99%
rename from src/mmcontext/mmcontextencoder.py
rename to src/mmcontext/_legacy/mmcontextencoder.py
index c4de365..38ee909 100644
--- a/src/mmcontext/mmcontextencoder.py
+++ b/src/mmcontext/_legacy/mmcontextencoder.py
@@ -42,14 +42,14 @@
 from sentence_transformers.models import Module, Pooling
 from transformers import AutoModel, AutoTokenizer
 
-from .adapters import AdapterModule
-
-# Import local dependencies using relative imports (works with package structure)
-from .file_utils import (
+# file_utils remains at the package root (not moved to _legacy)
+from mmcontext.file_utils import (
     build_embedding_df,
     collect_unique_links,
     download_and_extract_links,
 )
+
+from .adapters import AdapterModule
 from .omicsencoder import MiniOmicsModel
 from .onehot import OneHotTextEncoder
 
diff --git a/src/mmcontext/omicsencoder.py b/src/mmcontext/_legacy/omicsencoder.py
similarity index 100%
rename from src/mmcontext/omicsencoder.py
rename to src/mmcontext/_legacy/omicsencoder.py
diff --git a/src/mmcontext/onehot.py b/src/mmcontext/_legacy/onehot.py
similarity index 100%
rename from src/mmcontext/onehot.py
rename to src/mmcontext/_legacy/onehot.py
diff --git a/src/mmcontext/callback.py b/src/mmcontext/callback.py
index 8456c09..5864b62 100644
--- a/src/mmcontext/callback.py
+++ b/src/mmcontext/callback.py
@@ -2,6 +2,9 @@
 
 from transformers import TrainerCallback
 
+from mmcontext.modules.adapter_module import AdapterModule
+from mmcontext.modules.mmcontext_module import MMContextModule
+
 logger = logging.getLogger(__name__)
 
 
@@ -29,42 +32,48 @@ def on_epoch_begin(self, args, state, control, **kwargs):
         """
         if not self.unfrozen and state.epoch >= self.unfreeze_epoch:
             model = kwargs["model"]
-            # Assuming model[0] is your MMContextEncoder containing a text_encoder attribute
-            if hasattr(model[0], "text_encoder"):
-                for param in model[0].text_encoder.parameters():
+            if isinstance(model[0], MMContextModule):
+                model[0].unfreeze_text_encoder()
+                self.unfrozen = True
+                logger.info(f"Text encoder unfrozen at epoch {state.epoch:.2f}")
+            elif hasattr(model[0], "auto_model"):
+                for param in model[0].auto_model.parameters():
                     param.requires_grad = True
                 self.unfrozen = True
                 logger.info(f"Text encoder unfrozen at epoch {state.epoch:.2f}")
             else:
-                logger.warning("Model does not have a 'text_encoder' attribute at model[0].")
+                logger.warning("Model does not have a compatible text encoder at model[0].")
         return control
 
 
 class UnfreezeAdapterCallback(TrainerCallback):
     """
-    A TrainerCallback to freeze/unfreeze text and omics adapters at specified epochs.
+    A TrainerCallback to freeze/unfreeze text and omics adapter projections at specified epochs.
 
     This callback provides fine-grained control over adapter training by allowing you to:
-    1. Start with frozen adapters and unfreeze them at specific epochs
-    2. Control text and omics adapters independently
-    3. Handle cases where only one type of adapter is present
+    1. Start with frozen adapter projection heads and unfreeze them at specific epochs
+    2. Control text and omics projections independently
+    3. Handle cases where the AdapterModule is at any position in the pipeline
+
+    The callback searches ``model.children()`` for an :class:`AdapterModule` instance,
+    so it works regardless of whether an OmicsAttentionModule is included in the pipeline.
 
     Parameters
     ----------
     freeze_text_adapter : bool, optional
-        Whether to start with the text adapter frozen. Defaults to False.
+        Whether to start with the text projection head frozen. Defaults to False.
     freeze_omics_adapter : bool, optional
-        Whether to start with the omics adapter frozen. Defaults to False.
+        Whether to start with the omics projection head frozen. Defaults to False.
     unfreeze_text_adapter_epoch : float, optional
-        The epoch at which to unfreeze the text adapter. If None and freeze_text_adapter
-        is True, the adapter remains frozen. Defaults to None.
+        The epoch at which to unfreeze the text projection head. If None and freeze_text_adapter
+        is True, the projection remains frozen. Defaults to None.
     unfreeze_omics_adapter_epoch : float, optional
-        The epoch at which to unfreeze the omics adapter. If None and freeze_omics_adapter
-        is True, the adapter remains frozen. Defaults to None.
+        The epoch at which to unfreeze the omics projection head. If None and freeze_omics_adapter
+        is True, the projection remains frozen. Defaults to None.
 
     Examples
     --------
-    >>> # Freeze both adapters initially, unfreeze text at epoch 1, omics at epoch 2
+    >>> # Freeze both projections initially, unfreeze text at epoch 1, omics at epoch 2
     >>> callback = UnfreezeAdapterCallback(
     ...     freeze_text_adapter=True,
     ...     freeze_omics_adapter=True,
@@ -72,7 +81,7 @@ class UnfreezeAdapterCallback(TrainerCallback):
     ...     unfreeze_omics_adapter_epoch=2.0,
     ... )
 
-    >>> # Only control omics adapter (useful for text-only datasets)
+    >>> # Only control omics projection (useful for text-only datasets)
     >>> callback = UnfreezeAdapterCallback(freeze_omics_adapter=True, unfreeze_omics_adapter_epoch=1.5)
     """
 
@@ -88,59 +97,37 @@ def __init__(
         self.unfreeze_text_adapter_epoch = unfreeze_text_adapter_epoch
         self.unfreeze_omics_adapter_epoch = unfreeze_omics_adapter_epoch
 
-        # Track unfreezing state
         self.text_adapter_unfrozen = False
         self.omics_adapter_unfrozen = False
-
-        # Track initialization state
         self.adapters_initialized = False
+        self._adapter: AdapterModule | None = None
 
-    def _freeze_adapter(self, adapter, adapter_name: str):
-        """Freeze all parameters in an adapter."""
-        if adapter is not None:
-            for param in adapter.parameters():
-                param.requires_grad = False
-            logger.info(f"{adapter_name} adapter frozen")
-        else:
-            logger.debug(f"{adapter_name} adapter not present, skipping freeze")
-
-    def _unfreeze_adapter(self, adapter, adapter_name: str):
-        """Unfreeze all parameters in an adapter."""
-        if adapter is not None:
-            for param in adapter.parameters():
-                param.requires_grad = True
-            logger.info(f"{adapter_name} adapter unfrozen")
-        else:
-            logger.debug(f"{adapter_name} adapter not present, skipping unfreeze")
+    def _find_adapter(self, model) -> AdapterModule | None:
+        """Find the AdapterModule by iterating model.children()."""
+        for m in model.children():
+            if isinstance(m, AdapterModule):
+                return m
+        return None
 
     def on_train_begin(self, args, state, control, **kwargs):
         """Called at the beginning of training to initialize adapter freezing state."""
         model = kwargs["model"]
 
-        # Check if we have MMContextEncoder at model[0]
-        if not hasattr(model[0], "text_adapter") and not hasattr(model[0], "omics_adapter"):
-            logger.warning("Model does not have adapter attributes. Adapter callback will have no effect.")
-            # Don't set adapters_initialized = True when no adapters are present
+        adapter = self._find_adapter(model)
+        if adapter is None:
+            logger.warning("No AdapterModule found in model. Adapter callback will have no effect.")
             return control
 
-        # Initialize freezing state for adapters
-        if self.freeze_text_adapter and hasattr(model[0], "text_adapter"):
-            self._freeze_adapter(model[0].text_adapter, "Text")
+        self._adapter = adapter
 
-        if self.freeze_omics_adapter and hasattr(model[0], "omics_adapter"):
-            self._freeze_adapter(model[0].omics_adapter, "Omics")
+        if self.freeze_text_adapter:
+            adapter.freeze_text_proj()
 
-        self.adapters_initialized = True
-
-        # Log current adapter status
-        text_adapter_present = hasattr(model[0], "text_adapter") and model[0].text_adapter is not None
-        omics_adapter_present = hasattr(model[0], "omics_adapter") and model[0].omics_adapter is not None
-
-        logger.info(
-            f"Adapter callback initialized - Text adapter: {'present' if text_adapter_present else 'absent'}, "
-            f"Omics adapter: {'present' if omics_adapter_present else 'absent'}"
-        )
+        if self.freeze_omics_adapter:
+            adapter.freeze_omics_proj()
 
+        self.adapters_initialized = True
+        logger.info("Adapter callback initialized — AdapterModule found and freeze state applied.")
         return control
 
     def on_epoch_begin(self, args, state, control, **kwargs):
@@ -148,28 +135,22 @@ def on_epoch_begin(self, args, state, control, **kwargs):
         if not self.adapters_initialized:
             return control
 
-        model = kwargs["model"]
-
-        # Check text adapter unfreezing
         if (
             not self.text_adapter_unfrozen
             and self.unfreeze_text_adapter_epoch is not None
             and state.epoch >= self.unfreeze_text_adapter_epoch
         ):
-            if hasattr(model[0], "text_adapter"):
-                self._unfreeze_adapter(model[0].text_adapter, "Text")
-                self.text_adapter_unfrozen = True
-                logger.info(f"Text adapter unfrozen at epoch {state.epoch:.2f}")
+            self._adapter.unfreeze_text_proj()
+            self.text_adapter_unfrozen = True
+            logger.info(f"Text adapter projection unfrozen at epoch {state.epoch:.2f}")
 
-        # Check omics adapter unfreezing
         if (
             not self.omics_adapter_unfrozen
             and self.unfreeze_omics_adapter_epoch is not None
             and state.epoch >= self.unfreeze_omics_adapter_epoch
         ):
-            if hasattr(model[0], "omics_adapter"):
-                self._unfreeze_adapter(model[0].omics_adapter, "Omics")
-                self.omics_adapter_unfrozen = True
-                logger.info(f"Omics adapter unfrozen at epoch {state.epoch:.2f}")
+            self._adapter.unfreeze_omics_proj()
+            self.omics_adapter_unfrozen = True
+            logger.info(f"Omics adapter projection unfrozen at epoch {state.epoch:.2f}")
 
         return control
diff --git a/src/mmcontext/embed/__init__.py b/src/mmcontext/embed/__init__.py
index fb4b6bf..00ad7be 100644
--- a/src/mmcontext/embed/__init__.py
+++ b/src/mmcontext/embed/__init__.py
@@ -1,19 +1,22 @@
-from .cellwhisperer_utils import ensure_cellwhisperer_setup, process_cellwhisperer_dataset_model
-from .dataset_utils import SentenceDataset, collect_adata_subset, load_generic_dataset
-from .embed_pipeline import embed_pipeline, process_single_dataset_model
-from .model_utils import HFIndexedDataset, create_label_dataset, embed_labels, load_st_model, prepare_model_and_embed
+"""mmcontext.embed — Pipeline building, dataset preparation, embedding, and hard-negative mining."""
+
+from .dataset_prep import InferenceData, prepare_dataset, prepare_inference
+from .dataset_utils import load_generic_dataset
+from .encode import create_label_dataset, embed_labels, load_st_model, prepare_model_and_embed
+from .mining import MiningConfig, mine_negatives, mining_report
+from .pipeline import build_pipeline
 
 __all__ = [
-    "SentenceDataset",
+    "build_pipeline",
+    "create_label_dataset",
+    "embed_labels",
+    "InferenceData",
     "load_generic_dataset",
-    "collect_adata_subset",
-    "HFIndexedDataset",
     "load_st_model",
+    "MiningConfig",
+    "mine_negatives",
+    "mining_report",
+    "prepare_dataset",
+    "prepare_inference",
     "prepare_model_and_embed",
-    "create_label_dataset",
-    "embed_labels",
-    "embed_pipeline",
-    "process_single_dataset_model",
-    "process_cellwhisperer_dataset_model",
-    "ensure_cellwhisperer_setup",
 ]
diff --git a/src/mmcontext/embed/cellwhisperer_utils.py b/src/mmcontext/embed/cellwhisperer_utils.py
deleted file mode 100644
index 9ebed2d..0000000
--- a/src/mmcontext/embed/cellwhisperer_utils.py
+++ /dev/null
@@ -1,500 +0,0 @@
-"""
-CellWhisperer embedding utilities for integration with mmcontext pipeline.
-
-This module provides functions to generate embeddings using CellWhisperer models
-and integrate them into the mmcontext embedding/evaluation pipeline.
-"""
-
-import json
-import logging
-import os
-import subprocess
-import sys
-from pathlib import Path
-from typing import Any
-
-import anndata as ad
-import numpy as np
-import pandas as pd
-import torch
-from tqdm import tqdm
-
-logger = logging.getLogger(__name__)
-
-
-def ensure_cellwhisperer_setup(modules_dir: Path) -> bool:
-    """
-    Ensure CellWhisperer is properly set up.
-
-    This is done by cloning the repository, setting up the model, creating a virtual environment, and installing dependencies.
-
-    Parameters
-    ----------
-    modules_dir : Path
-        Directory where CellWhisperer should be cloned/installed
-
-    Returns
-    -------
-    bool
-        True if setup successful, False otherwise
-    """
-    cellwhisperer_dir = modules_dir / "CellWhisperer"
-    cellwhisperer_requirements = modules_dir / "requirements_cellwhisperer.txt"
-    cellwhisperer_venv = modules_dir / "cellwhisperer_venv"
-    cellwhisperer_python = cellwhisperer_venv / "bin" / "python"
-
-    # Check if CellWhisperer is already set up (repository + model + venv + python)
-    setup_complete = (
-        cellwhisperer_dir.exists()
-        and (cellwhisperer_dir / "resources" / "geneformer-12L-30M").exists()
-        and cellwhisperer_venv.exists()
-        and cellwhisperer_python.exists()
-    )
-
-    if setup_complete:
-        logger.info("CellWhisperer already set up with virtual environment")
-        return True
-
-    try:
-        # Step 1: Run prepare script to clone repository and set up model
-        prepare_script = modules_dir / "prepare_cellwhisperer.sh"
-        if not prepare_script.exists():
-            logger.error(f"Prepare script not found at {prepare_script}")
-            return False
-
-        logger.info("Running CellWhisperer prepare script...")
-        result = subprocess.run(
-            ["bash", str(prepare_script)],
-            cwd=modules_dir.parent,  # Run from repo root
-            capture_output=True,
-            text=True,
-            timeout=600,  # 10 minute timeout
-        )
-
-        if result.returncode != 0:
-            logger.error(f"CellWhisperer prepare script failed with return code {result.returncode}")
-            if result.stdout:
-                logger.error(f"Prepare stdout: {result.stdout}")
-            if result.stderr:
-                logger.error(f"Prepare stderr: {result.stderr}")
-            return False
-
-        logger.info("CellWhisperer prepare script completed successfully")
-        if result.stdout:
-            logger.info(f"Prepare output: {result.stdout}")
-
-        # Step 2: Create virtual environment
-        logger.info("Creating CellWhisperer virtual environment...")
-        if cellwhisperer_venv.exists():
-            logger.info("Virtual environment already exists, removing...")
-            import shutil
-
-            shutil.rmtree(cellwhisperer_venv)
-
-        result = subprocess.run(
-            [sys.executable, "-m", "venv", str(cellwhisperer_venv)], capture_output=True, text=True, timeout=60
-        )
-
-        if result.returncode != 0:
-            logger.error(f"Failed to create virtual environment: {result.stderr}")
-            return False
-
-        logger.info("✓ Virtual environment created successfully")
-
-        # Step 3: Install requirements in the virtual environment
-        if not cellwhisperer_requirements.exists():
-            logger.error(f"Requirements file not found at {cellwhisperer_requirements}")
-            return False
-
-        logger.info("Installing CellWhisperer dependencies...")
-        result = subprocess.run(
-            [str(cellwhisperer_python), "-m", "pip", "install", "-r", str(cellwhisperer_requirements)],
-            capture_output=True,
-            text=True,
-            timeout=1200,  # 20 minute timeout for package installation
-        )
-
-        if result.returncode != 0:
-            logger.error(f"Failed to install requirements: {result.stderr}")
-            if result.stdout:
-                logger.error(f"Install stdout: {result.stdout}")
-            return False
-
-        logger.info("✓ CellWhisperer dependencies installed successfully")
-        if result.stdout:
-            logger.info(f"Install output: {result.stdout}")
-
-        # Verify setup is complete
-        if not cellwhisperer_python.exists():
-            logger.error("CellWhisperer venv Python not found after setup")
-            return False
-
-        logger.info("✓ CellWhisperer setup completed successfully")
-        return True
-
-    except subprocess.TimeoutExpired:
-        logger.error("CellWhisperer setup timed out")
-        return False
-    except Exception as e:
-        logger.error(f"Error running CellWhisperer setup: {e}")
-        return False
-
-
-def create_label_embeddings(labels, model, tokenizer, batch_size):
-    """Create embeddings for label strings using CellWhisperer text encoder."""
-    import pandas as pd
-
-    logger.info(f"Creating label embeddings for {len(labels)} unique labels...")
-
-    embeddings_list = []
-    sample_indices = []
-
-    # Process labels in batches
-    for start_idx in tqdm(range(0, len(labels), batch_size), desc="Creating label embeddings"):
-        end_idx = min(start_idx + batch_size, len(labels))
-        batch_labels = labels[start_idx:end_idx]
-
-        # Get text embeddings using model's embed_texts method
-        with torch.no_grad():
-            text_embeddings = model.embed_texts(batch_labels, chunk_size=len(batch_labels))
-            text_embeddings = text_embeddings.cpu().numpy()
-
-        # Store embeddings
-        for i, embedding in enumerate(text_embeddings):
-            sample_idx = start_idx + i
-            sample_indices.append(sample_idx)
-            embeddings_list.append(embedding.tolist())
-
-    embeddings_df = pd.DataFrame({"sample_idx": sample_indices, "embedding": embeddings_list})
-
-    logger.info(f"✓ Created {len(embeddings_df)} label embeddings with dimension {len(embeddings_list[0])}")
-    return embeddings_df
-
-
-def create_cellwhisperer_embeddings_venv(
-    adata_path: Path, modules_dir: Path, cache_dir: Path, batch_size: int = 32, annotation_keys: list[str] = None
-) -> tuple[pd.DataFrame, dict, float | None]:
-    """
-    Create embeddings using CellWhisperer model running in separate venv.
-
-    Parameters
-    ----------
-    adata_path : Path
-        Path to saved AnnData file
-    modules_dir : Path
-        Directory containing CellWhisperer installation
-    cache_dir : Path
-        Cache directory for model checkpoints
-    batch_size : int
-        Batch size for processing
-    annotation_keys : list[str]
-        Annotation keys to process for similarity scoring
-
-    Returns
-    -------
-    tuple[pd.DataFrame, dict, float | None]
-        (embeddings_df, similarity_results, logit_scale)
-    """
-    import pickle
-    import tempfile
-
-    logger.info("Running CellWhisperer in separate virtual environment...")
-
-    # Create temporary files for communication
-    with tempfile.TemporaryDirectory() as temp_dir:
-        temp_dir = Path(temp_dir)
-        results_file = temp_dir / "cellwhisperer_results.pkl"
-
-        # Get path to standalone script
-        script_path = Path(__file__).parent.parent.parent.parent / "scripts" / "run_cellwhisperer_standalone.py"
-        if not script_path.exists():
-            raise RuntimeError(f"CellWhisperer standalone script not found: {script_path}")
-
-        # Prepare arguments
-        args = [
-            "--adata_path",
-            str(adata_path),
-            "--cache_dir",
-            str(cache_dir),
-            "--output_file",
-            str(results_file),
-            "--batch_size",
-            str(batch_size),
-        ]
-
-        if annotation_keys:
-            args.extend(["--annotation_keys", ",".join(annotation_keys)])
-            # Pass project cache directory for label embeddings (not HuggingFace cache)
-            current_dir = Path.cwd()
-            if "runs" in str(current_dir):
-                # We're in a Hydra run directory, go up to project root
-                project_root = current_dir
-                while project_root.name != "mmcontext" and project_root.parent != project_root:
-                    project_root = project_root.parent
-                project_cache_dir = project_root / "cache"
-            else:
-                # We're already in project root
-                project_cache_dir = current_dir / "cache"
-            args.extend(["--output_dir", str(project_cache_dir)])
-
-        # Run CellWhisperer standalone script in venv
-        cellwhisperer_python = modules_dir / "cellwhisperer_venv" / "bin" / "python"
-
-        if not cellwhisperer_python.exists():
-            raise RuntimeError(f"CellWhisperer venv Python not found at {cellwhisperer_python}")
-
-        cmd = [str(cellwhisperer_python), str(script_path)] + args
-        logger.info(f"Running command: {' '.join(cmd)}")
-
-        result = subprocess.run(
-            cmd,
-            capture_output=True,
-            text=True,
-            timeout=259200,  # 3 day timeout
-        )
-
-        logger.info(f"CellWhisperer script return code: {result.returncode}")
-        if result.stdout:
-            logger.info(f"CellWhisperer stdout: {result.stdout}")
-        if result.stderr:
-            logger.info(f"CellWhisperer stderr: {result.stderr}")
-
-        if result.returncode != 0:
-            logger.error(f"CellWhisperer script failed with return code {result.returncode}")
-            raise RuntimeError(f"CellWhisperer processing failed: {result.stderr}")
-
-        # Load results
-        if results_file.exists():
-            with open(results_file, "rb") as f:
-                results = pickle.load(f)
-
-            if results["success"]:
-                logger.info("✓ CellWhisperer processing completed successfully")
-                logger.info(
-                    f"Processed {results['n_cells']} cells with {results['n_genes']} genes on {results['device']}"
-                )
-                logit_scale = results.get("logit_scale")
-                if logit_scale is not None:
-                    try:
-                        logit_scale = float(logit_scale)
-                        logger.info(f"✓ Received logit_scale from CellWhisperer run: {logit_scale}")
-                    except Exception:
-                        logger.warning(f"CellWhisperer returned non-numeric logit_scale={logit_scale!r}")
-                        logit_scale = None
-                return results["embeddings_df"], results["similarity_results"], logit_scale
-            else:
-                raise RuntimeError(f"CellWhisperer processing failed: {results['error']}")
-        else:
-            raise RuntimeError("CellWhisperer results file not found")
-
-    return pd.DataFrame(), {}, None
-
-
-def process_cellwhisperer_dataset_model(
-    ds_cfg: Any,
-    model_cfg: Any,
-    run_cfg: Any,
-    output_root: str,
-    output_format: str,
-    adata_cache: str,
-    modules_dir: Path,
-    hf_cache: str = None,
-) -> tuple[str, str, bool, str]:
-    """
-    Process a single dataset with CellWhisperer model for embedding generation.
-
-    This function follows the same pattern as process_single_dataset_model but
-    uses CellWhisperer instead of sentence transformers.
-
-    Parameters
-    ----------
-    ds_cfg : dict
-        Dataset configuration
-    model_cfg : dict
-        Model configuration (should specify CellWhisperer)
-    run_cfg : dict
-        Run configuration
-    output_root : str
-        Output root directory
-    output_format : str
-        Output format (parquet/csv)
-    adata_cache : str
-        AnnData cache directory
-    modules_dir : Path
-        Directory containing CellWhisperer installation
-    hf_cache : str, optional
-        HuggingFace cache directory
-
-    Returns
-    -------
-    tuple[str, str, bool, str]
-        (dataset_name, model_name, success, error_msg)
-    """
-    from mmcontext.embed.dataset_utils import collect_adata_subset, load_generic_dataset
-    from mmcontext.file_utils import save_table
-
-    dataset_name = ds_cfg.name
-    model_name = model_cfg.get("name", "cellwhisperer")
-
-    logger.info(f"Processing CellWhisperer: {dataset_name} + {model_name}")
-
-    # Determine output directory
-    out_dir = Path(output_root) / ds_cfg.name / model_name
-
-    # Check if files already exist and skip if overwrite is False
-    from mmcontext.embed.embed_pipeline import check_required_files_exist
-
-    if not run_cfg.overwrite and check_required_files_exist(out_dir, output_format):
-        logger.info(f"CellWhisperer embeddings already exist for {ds_cfg.name} + {model_name}. Skipping.")
-        return dataset_name, model_name, True, "skipped_existing"
-
-    try:
-        # Ensure CellWhisperer is set up
-        if not ensure_cellwhisperer_setup(modules_dir):
-            return dataset_name, model_name, False, "CellWhisperer setup failed"
-
-        # Load dataset
-        adata_download_dir = Path(adata_cache) / ds_cfg.name
-        cache_dir = hf_cache if hf_cache is not None else getattr(run_cfg, "hf_cache", None)
-
-        # Load raw dataset
-        raw_ds = load_generic_dataset(
-            source=ds_cfg.source,
-            fmt=ds_cfg.format,
-            split=ds_cfg.get("split", "test"),
-            max_rows=run_cfg.n_rows,
-            seed=run_cfg.seed,
-            cache_dir=cache_dir,
-        )
-
-        # Check if numeric data is available
-        numeric_data_available = ("adata_link" in raw_ds.column_names) or ("share_link" in raw_ds.column_names)
-
-        if not numeric_data_available:
-            logger.warning("CellWhisperer requires numeric data (share_link or adata_link), but none found")
-            return dataset_name, model_name, False, "No numeric data available for CellWhisperer"
-
-        # Get sample IDs and collect AnnData subset
-        sample_ids = raw_ds[ds_cfg.index_col]
-        if isinstance(sample_ids[0], str) and "sample_idx:" in sample_ids[0]:
-            sample_ids = [sid.split(":")[1] for sid in sample_ids]
-
-        adata_subset = collect_adata_subset(
-            download_dir=adata_download_dir,
-            sample_ids=sample_ids,
-        )
-
-        # Save AnnData to temporary file for CellWhisperer venv processing
-        import tempfile
-
-        with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
-            temp_adata_path = Path(tmp_file.name)
-            adata_subset.write_h5ad(temp_adata_path)
-
-        try:
-            # Get annotation keys for similarity scoring
-            annotation_keys = []
-            if hasattr(ds_cfg, "bio_label_list") and ds_cfg.bio_label_list:
-                annotation_keys.extend(ds_cfg.bio_label_list)
-
-            # Create embeddings using CellWhisperer in separate venv
-            embeddings_df, similarity_results, logit_scale = create_cellwhisperer_embeddings_venv(
-                temp_adata_path,
-                modules_dir,
-                Path(cache_dir) if cache_dir else Path.cwd(),
-                batch_size=run_cfg.batch_size,
-                annotation_keys=annotation_keys,
-            )
-
-            # Create output directory
-            out_dir.mkdir(parents=True, exist_ok=True)
-
-            # Persist CellWhisperer metadata for downstream eval (no cellwhisperer dependency)
-            if logit_scale is not None:
-                meta_path = out_dir / "cellwhisperer_meta.json"
-                meta_payload = {"logit_scale": float(logit_scale)}
-                meta_path.write_text(json.dumps(meta_payload, indent=2, sort_keys=True) + "\n")
-                logger.info(f"Wrote CellWhisperer metadata → {meta_path}")
-
-            # Copy label embeddings to output directory if they were generated
-            if annotation_keys:
-                import shutil
-
-                # Use project cache directory for label embeddings
-                current_dir = Path.cwd()
-                if "runs" in str(current_dir):
-                    # We're in a Hydra run directory, go up to project root
-                    project_root = current_dir
-                    while project_root.name != "mmcontext" and project_root.parent != project_root:
-                        project_root = project_root.parent
-                    project_cache_dir = project_root / "cache"
-                else:
-                    # We're already in project root
-                    project_cache_dir = current_dir / "cache"
-
-                for annotation_key in annotation_keys:
-                    # Look for label embedding files in project cache directory
-                    label_emb_file = project_cache_dir / f"bio_label_embeddings_{annotation_key}.parquet"
-                    mapping_file = project_cache_dir / f"bio_label_embeddings_{annotation_key}_mapping.json"
-
-                    if label_emb_file.exists():
-                        # Copy to output directory
-                        shutil.copy2(label_emb_file, out_dir / label_emb_file.name)
-                        logger.info(f"Copied label embeddings for {annotation_key}")
-                    else:
-                        logger.warning(f"Label embeddings file not found: {label_emb_file}")
-
-                    if mapping_file.exists():
-                        # Copy to output directory
-                        shutil.copy2(mapping_file, out_dir / mapping_file.name)
-                        logger.info(f"Copied label mapping for {annotation_key}")
-                    else:
-                        logger.warning(f"Label mapping file not found: {mapping_file}")
-        finally:
-            # Clean up temporary file
-            if temp_adata_path.exists():
-                temp_adata_path.unlink()
-
-        # Save embeddings (output directory already created above)
-        save_table(
-            embeddings_df,
-            out_path=out_dir / "embeddings",
-            fmt=output_format,
-        )
-
-        # Save metadata
-        (out_dir / "meta.yaml").write_text(f"model: {model_name}\ndataset: {ds_cfg.name}\nrows: {len(embeddings_df)}\n")
-
-        # Save AnnData subset
-        subset_out = out_dir / "subset.h5ad"
-        adata_subset.write_h5ad(subset_out)
-        logger.info("Wrote subset AnnData → %s", subset_out)
-
-        # Save similarity results if available
-        if similarity_results:
-            for annotation_key, results in similarity_results.items():
-                # Save predictions as CSV (separate from label embeddings)
-                predictions_df = pd.DataFrame(
-                    {
-                        "cell_id": range(len(results["true_labels"])),
-                        "true_labels": results["true_labels"],
-                        "predicted_labels": results["predicted_labels"],
-                    }
-                )
-
-                save_table(
-                    predictions_df,
-                    out_path=out_dir / f"cellwhisperer_predictions_{annotation_key}",
-                    fmt=output_format,
-                )
-
-                # Save cell type labels
-                with open(out_dir / f"cellwhisperer_labels_{annotation_key}.json", "w") as f:
-                    json.dump(results["cell_type_labels"], f, indent=2)
-
-        logger.info(f"✓ Completed CellWhisperer processing {dataset_name} + {model_name}")
-        return dataset_name, model_name, True, None
-
-    except Exception as e:
-        logger.error(f"✗ Error processing CellWhisperer {dataset_name} + {model_name}: {e}")
-        return dataset_name, model_name, False, str(e)
diff --git a/src/mmcontext/embed/dataset_prep.py b/src/mmcontext/embed/dataset_prep.py
new file mode 100644
index 0000000..73e8801
--- /dev/null
+++ b/src/mmcontext/embed/dataset_prep.py
@@ -0,0 +1,449 @@
+"""Unified dataset preparation for training and inference.
+
+mmcontext aligns text and omics embeddings with a sentence-transformers
+pipeline whose first module is an :class:`~mmcontext.modules.MMContextModule`.
+That module resolves omics samples at runtime from an attached
+:class:`~mmcontext.io.VectorStore`: any input string starting with the omics
+prefix (``"omics:"`` by default) is looked up in the store, while plain strings
+are tokenised as text.
+
+This module provides a single, reusable preparation API that turns a raw
+HuggingFace dataset into the column shape the trainer / encoder expects:
+
+* :func:`prepare_dataset` — the pure ``dataset -> dataset`` core. It builds the
+  ``anchor`` column (omics id for ``modality="bimodal"`` or cell-sentence text
+  for ``modality="text"``) and, for ``purpose="train"``, resolves the
+  ``positive`` / hard-negative columns.
+* :func:`prepare_inference` — an orchestrator for the evaluation/inference path.
+  It loads the referenced AnnData chunk, subsets the dataset to it, builds an
+  anchor-ready dataset, and (for bimodal) builds and attaches a VectorStore to
+  the model.
+
+Training datasets carry ``positive`` and ``negative_*_idx`` columns; test
+datasets carry only ``sample_idx``, ``cell_sentence_*`` and ``adata_link``.
+The same :func:`prepare_dataset` call handles both via the ``purpose`` switch.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+from datasets import Dataset, DatasetDict
+
+from mmcontext.utils import resolve_negative_indices_and_rename, truncate_cell_sentences
+
+if TYPE_CHECKING:
+    import anndata as ad
+    from sentence_transformers import SentenceTransformer
+
+    from mmcontext.io import VectorStore
+
+logger = logging.getLogger(__name__)
+
+Purpose = Literal["train", "inference"]
+Modality = Literal["bimodal", "text"]
+
+
+def _has_negative_idx_columns(ds: Dataset) -> bool:
+    """Return True if *ds* exposes ``negative_*_idx`` columns."""
+    return any(c.startswith("negative_") and c.endswith("_idx") for c in ds.column_names)
+
+
+def _prepare_split(
+    ds: Dataset,
+    *,
+    purpose: Purpose,
+    modality: Modality,
+    primary_cell_sentence: str,
+    sample_id_col: str,
+    positive_col: str,
+    omics_prefix: str,
+    use_hard_negatives: bool,
+    truncate: bool,
+    truncate_kwargs: dict | None,
+) -> Dataset:
+    """Prepare a single :class:`~datasets.Dataset` split.
+
+    See :func:`prepare_dataset` for the meaning of the parameters.
+    """
+    if modality not in ("bimodal", "text"):
+        raise ValueError(f"modality must be 'bimodal' or 'text', got {modality!r}")
+    if purpose not in ("train", "inference"):
+        raise ValueError(f"purpose must be 'train' or 'inference', got {purpose!r}")
+
+    if primary_cell_sentence not in ds.column_names:
+        raise KeyError(
+            f"Primary cell-sentence column {primary_cell_sentence!r} not found. Available columns: {ds.column_names}"
+        )
+    if modality == "bimodal" and sample_id_col not in ds.column_names:
+        raise KeyError(
+            f"modality='bimodal' requires the sample-id column {sample_id_col!r}. Available columns: {ds.column_names}"
+        )
+
+    # ------------------------------------------------------------------
+    # text modality: optionally truncate the cell sentence before it
+    # becomes the anchor (gene filtering, max length, …).
+    # ------------------------------------------------------------------
+    if modality == "text" and truncate:
+        tkwargs = dict(truncate_kwargs or {})
+        max_length = tkwargs.pop("max_length", 64)
+        ds = truncate_cell_sentences(ds, primary_cell_sentence, max_length=max_length, **tkwargs)
+
+    if purpose == "train":
+        return _prepare_train_split(
+            ds,
+            modality=modality,
+            primary_cell_sentence=primary_cell_sentence,
+            sample_id_col=sample_id_col,
+            positive_col=positive_col,
+            omics_prefix=omics_prefix,
+            use_hard_negatives=use_hard_negatives,
+        )
+    return _prepare_inference_split(
+        ds,
+        modality=modality,
+        primary_cell_sentence=primary_cell_sentence,
+        sample_id_col=sample_id_col,
+        omics_prefix=omics_prefix,
+    )
+
+
+def _prepare_train_split(
+    ds: Dataset,
+    *,
+    modality: Modality,
+    primary_cell_sentence: str,
+    sample_id_col: str,
+    positive_col: str,
+    omics_prefix: str,
+    use_hard_negatives: bool,
+) -> Dataset:
+    """Build a training-ready split: ``anchor`` + ``positive`` (+ ``negative_*``)."""
+    if positive_col not in ds.column_names:
+        raise KeyError(
+            f"purpose='train' requires a positive column {positive_col!r}. Available columns: {ds.column_names}"
+        )
+
+    resolve = use_hard_negatives and _has_negative_idx_columns(ds)
+    if resolve:
+        # Resolves negative_*_idx -> text, renames primary_cell_sentence -> 'anchor'
+        # and positive_col -> 'positive'. Keep sample_idx so we can build the
+        # omics anchor for the bimodal case afterwards.
+        ds = resolve_negative_indices_and_rename(
+            ds,
+            primary_cell_sentence_col=primary_cell_sentence,
+            positive_col=positive_col,
+            negative_prefix="negative",
+            index_col=sample_id_col,
+            remove_index_col=False,
+        )
+    else:
+        if primary_cell_sentence != "anchor":
+            ds = ds.rename_column(primary_cell_sentence, "anchor")
+        if positive_col != "positive":
+            ds = ds.rename_column(positive_col, "positive")
+
+    if modality == "bimodal":
+        ds = ds.map(
+            lambda row: {"anchor": f"{omics_prefix}{row[sample_id_col]}"},
+            desc="Building omics anchors",
+        )
+
+    # Keep only resolved negatives (named "negative_N"); never raw "_idx" columns.
+    neg_cols = sorted(c for c in ds.column_names if c.startswith("negative") and not c.endswith("_idx"))
+    keep = ["anchor", "positive", *neg_cols]
+    ds = ds.select_columns(keep)
+    logger.info("Prepared train split: %d rows, columns=%s", len(ds), ds.column_names)
+    return ds
+
+
+def _prepare_inference_split(
+    ds: Dataset,
+    *,
+    modality: Modality,
+    primary_cell_sentence: str,
+    sample_id_col: str,
+    omics_prefix: str,
+) -> Dataset:
+    """Build an inference-ready split: ``anchor`` (+ ``sample_idx``/``adata_link``).
+
+    No ``positive`` / ``negative_*`` columns are required or produced.
+    """
+    if modality == "bimodal":
+        ds = ds.map(
+            lambda row: {"anchor": f"{omics_prefix}{row[sample_id_col]}"},
+            desc="Building omics anchors",
+        )
+    else:  # text
+        if primary_cell_sentence != "anchor":
+            ds = ds.map(
+                lambda row: {"anchor": row[primary_cell_sentence]},
+                desc="Building text anchors",
+            )
+
+    # Retain identifiers needed downstream to align embeddings with AnnData.
+    keep = ["anchor"]
+    for col in (sample_id_col, "adata_link", "share_link"):
+        if col in ds.column_names and col not in keep:
+            keep.append(col)
+    ds = ds.select_columns(keep)
+    logger.info("Prepared inference split: %d rows, columns=%s", len(ds), ds.column_names)
+    return ds
+
+
+def prepare_dataset(
+    ds: Dataset | DatasetDict,
+    *,
+    purpose: Purpose = "train",
+    modality: Modality = "bimodal",
+    primary_cell_sentence: str = "cell_sentence_1",
+    sample_id_col: str = "sample_idx",
+    positive_col: str = "positive",
+    omics_prefix: str = "omics:",
+    use_hard_negatives: bool = True,
+    truncate: bool = False,
+    truncate_kwargs: dict | None = None,
+) -> Dataset | DatasetDict:
+    """Prepare a raw HuggingFace dataset for training or inference.
+
+    The first module of an mmcontext pipeline
+    (:class:`~mmcontext.modules.MMContextModule`) routes inputs by modality:
+    strings prefixed with *omics_prefix* are resolved through an attached
+    :class:`~mmcontext.io.VectorStore`, everything else is treated as text.
+    This function builds the ``anchor`` column accordingly and, for training,
+    the ``positive`` and hard-negative columns.
+
+    Parameters
+    ----------
+    ds : datasets.Dataset or datasets.DatasetDict
+        Raw dataset. A training dataset is expected to carry *positive_col* and
+        (optionally) ``negative_*_idx`` columns; an inference/test dataset only
+        needs *sample_id_col* / *primary_cell_sentence* (and an adata link for
+        the bimodal path).
+    purpose : {"train", "inference"}, default "train"
+        ``"train"`` produces ``anchor`` + ``positive`` (+ ``negative_*``).
+        ``"inference"`` produces only ``anchor`` plus retained identifier
+        columns; ``positive`` / ``negative_*`` are neither required nor created.
+    modality : {"bimodal", "text"}, default "bimodal"
+        ``"bimodal"`` builds ``anchor = f"{omics_prefix}{sample_idx}"`` so the
+        model resolves it via the VectorStore. ``"text"`` uses the
+        *primary_cell_sentence* text as the anchor.
+    primary_cell_sentence : str, default "cell_sentence_1"
+        Column holding the cell sentence (used as the text anchor and as the
+        primary column for negative resolution during training).
+    sample_id_col : str, default "sample_idx"
+        Column with sample identifiers (used to build omics anchors and to
+        resolve negatives). Required for ``modality="bimodal"``.
+    positive_col : str, default "positive"
+        Column with the positive text. Required for ``purpose="train"``.
+    omics_prefix : str, default "omics:"
+        Prefix marking an omics sample id. Must match the prefix configured on
+        the model's :class:`~mmcontext.modules.MMContextModule`.
+    use_hard_negatives : bool, default True
+        For ``purpose="train"``: resolve ``negative_*_idx`` columns into text
+        negatives when present. Ignored when no such columns exist.
+    truncate : bool, default False
+        For ``modality="text"``: truncate the cell sentence via
+        :func:`~mmcontext.utils.truncate_cell_sentences` before it becomes the
+        anchor.
+    truncate_kwargs : dict, optional
+        Extra keyword arguments forwarded to
+        :func:`~mmcontext.utils.truncate_cell_sentences` (e.g. ``max_length``,
+        ``filter_strings``). Defaults to ``max_length=64``.
+
+    Returns
+    -------
+    datasets.Dataset or datasets.DatasetDict
+        Same container type as *ds*, with prepared columns.
+
+    Examples
+    --------
+    >>> train_ds = prepare_dataset(raw, purpose="train", modality="bimodal")
+    >>> train_ds.column_names
+    ['anchor', 'positive', 'negative_1']
+    >>> test_ds = prepare_dataset(raw_test, purpose="inference", modality="bimodal")
+    >>> test_ds.column_names
+    ['anchor', 'sample_idx', 'adata_link']
+    """
+    kwargs = {
+        "purpose": purpose,
+        "modality": modality,
+        "primary_cell_sentence": primary_cell_sentence,
+        "sample_id_col": sample_id_col,
+        "positive_col": positive_col,
+        "omics_prefix": omics_prefix,
+        "use_hard_negatives": use_hard_negatives,
+        "truncate": truncate,
+        "truncate_kwargs": truncate_kwargs,
+    }
+
+    if isinstance(ds, DatasetDict):
+        return DatasetDict({name: _prepare_split(split, **kwargs) for name, split in ds.items()})
+    if isinstance(ds, Dataset):
+        return _prepare_split(ds, **kwargs)
+    raise TypeError(f"prepare_dataset expects a Dataset or DatasetDict, got {type(ds).__name__}")
+
+
+@dataclass
+class InferenceData:
+    """Bundle of everything needed to encode and evaluate one test chunk.
+
+    Attributes
+    ----------
+    dataset : datasets.Dataset
+        Anchor-ready dataset aligned to *adata* (encode ``dataset["anchor"]``).
+    adata : anndata.AnnData
+        AnnData chunk referenced by the dataset, subset to the dataset rows.
+    vector_store : VectorStore or None
+        Store attached to the model for the bimodal path; ``None`` for text.
+    local_path : pathlib.Path
+        Path to the downloaded/cached AnnData store.
+    """
+
+    dataset: Dataset
+    adata: ad.AnnData
+    vector_store: VectorStore | None
+    local_path: Path
+
+
+def prepare_inference(
+    model: SentenceTransformer,
+    ds: Dataset,
+    *,
+    modality: Modality = "bimodal",
+    obsm_key: str | None = None,
+    cache_dir: str | Path,
+    store_path: str | Path | None = None,
+    sample_id_col: str = "sample_idx",
+    adata_link_col: str = "adata_link",
+    omics_prefix: str = "omics:",
+    zenodo_token: str | None = None,
+    truncate: bool = False,
+    truncate_kwargs: dict | None = None,
+    overwrite_store: bool = False,
+    lean: bool = False,
+    obs_columns: list[str] | None = None,
+) -> InferenceData:
+    """Wire up a model + dataset + AnnData for inference on one test chunk.
+
+    This replaces the manual sequence of loading the AnnData chunk, subsetting
+    the dataset, building a VectorStore, and attaching it to the model. For the
+    bimodal path it builds the store from *obsm_key* and calls
+    ``model[0].set_vector_store(store)`` so ``model.encode(bundle.dataset["anchor"])``
+    works directly.
+
+    Parameters
+    ----------
+    model : sentence_transformers.SentenceTransformer
+        Model whose first module is an
+        :class:`~mmcontext.modules.MMContextModule` (only required for the
+        bimodal path, where the VectorStore is attached to it).
+    ds : datasets.Dataset
+        Test split referencing a single AnnData chunk via *adata_link_col*.
+    modality : {"bimodal", "text"}, default "bimodal"
+        ``"bimodal"`` builds and attaches a VectorStore; ``"text"`` skips it.
+    obsm_key : str, optional
+        ``adata.obsm`` key to extract for the VectorStore (e.g. ``"X_scvi_fm"``).
+        Required when ``modality="bimodal"``.
+    cache_dir : str or Path
+        Directory for caching the downloaded AnnData store (shared by the chunk
+        loader and the VectorStore builder, which use the same cache key).
+    store_path : str or Path, optional
+        Output ``.mmap`` path for the VectorStore. Defaults to
+        ``<cache_dir>/vector_store_inference.mmap``. Unused for text modality.
+    sample_id_col : str, default "sample_idx"
+        Column with sample identifiers.
+    adata_link_col : str, default "adata_link"
+        Column with the link to the AnnData chunk.
+    omics_prefix : str, default "omics:"
+        Prefix for omics anchors (must match the model's module).
+    zenodo_token : str, optional
+        Token for authenticating Zenodo draft downloads.
+    truncate : bool, default False
+        Forwarded to :func:`prepare_dataset` (text modality).
+    truncate_kwargs : dict, optional
+        Forwarded to :func:`prepare_dataset` (text modality).
+    overwrite_store : bool, default False
+        Rebuild the VectorStore even if *store_path* already exists.
+    lean : bool, default False
+        If True, load only the ``obs`` table (subset to *obs_columns*) and the
+        single ``obsm[obsm_key]`` layer for the returned AnnData — no ``X``. The
+        VectorStore is built independently from the store, so the returned
+        ``InferenceData.adata`` only needs its label columns downstream. Keeps
+        peak memory low on large chunks.
+    obs_columns : list[str] | None, default None
+        With *lean* True, restrict the returned ``obs`` to these columns (plus
+        the index). Ignored when *lean* is False.
+
+    Returns
+    -------
+    InferenceData
+        Bundle with the prepared dataset, the subset AnnData, the VectorStore
+        (or ``None``), and the local store path.
+
+    Examples
+    --------
+    >>> bundle = prepare_inference(model, test_ds, obsm_key="X_scvi_fm", cache_dir="./cache")
+    >>> emb = model.encode(bundle.dataset["anchor"])
+    >>> bundle.adata.obsm["mmcontext_emb"] = emb
+    """
+    # Imported here to keep module import light and avoid heavy/circular deps.
+    from mmcontext.file_utils import load_test_adata_from_hf_dataset, subset_dataset_by_chunk
+
+    if modality == "bimodal" and obsm_key is None:
+        raise ValueError("obsm_key is required when modality='bimodal'.")
+
+    # The text anchor source follows from the modality: bimodal anchors are
+    # omics ids (cell_sentence unused), text anchors use the text description.
+    primary_cell_sentence = "cell_sentence_1" if modality == "bimodal" else "cell_sentence_2"
+
+    cache_dir = Path(cache_dir)
+
+    # 1) Download + load the single AnnData chunk referenced by this split.
+    adata, local_path = load_test_adata_from_hf_dataset(
+        ds,
+        save_dir=cache_dir,
+        layer_key=obsm_key if modality == "bimodal" else None,
+        link_column=adata_link_col,
+        zenodo_token=zenodo_token,
+        lean=lean,
+        obs_columns=obs_columns,
+    )
+
+    # 2) Restrict the dataset to the rows present in this chunk.
+    adata, ds_sub = subset_dataset_by_chunk(adata, ds, sample_idx_col=sample_id_col)
+
+    # 3) Build the anchor-ready dataset (no positive/negative needed).
+    dataset = prepare_dataset(
+        ds_sub,
+        purpose="inference",
+        modality=modality,
+        primary_cell_sentence=primary_cell_sentence,
+        sample_id_col=sample_id_col,
+        omics_prefix=omics_prefix,
+        truncate=truncate,
+        truncate_kwargs=truncate_kwargs,
+    )
+
+    # 4) For bimodal, build + attach the VectorStore so omics anchors resolve.
+    vector_store: VectorStore | None = None
+    if modality == "bimodal":
+        from mmcontext.io import prepare_vector_store
+
+        if store_path is None:
+            store_path = cache_dir / "vector_store_inference.mmap"
+        vector_store = prepare_vector_store(
+            ds_sub,
+            obsm_key=obsm_key,
+            output_path=store_path,
+            cache_dir=cache_dir,
+            sample_id_column=sample_id_col,
+            adata_link_column=adata_link_col,
+            overwrite=overwrite_store,
+        )
+        model[0].set_vector_store(vector_store)
+
+    return InferenceData(dataset=dataset, adata=adata, vector_store=vector_store, local_path=local_path)
diff --git a/src/mmcontext/embed/dataset_utils.py b/src/mmcontext/embed/dataset_utils.py
index c66cf49..29d5c86 100644
--- a/src/mmcontext/embed/dataset_utils.py
+++ b/src/mmcontext/embed/dataset_utils.py
@@ -1,44 +1,27 @@
-# embedding_benchmark/model_utils.py
+"""Dataset loading helpers for the embed/inference path.
+
+Holds :func:`load_generic_dataset` (HF hub / on-disk / CSV loader) and
+re-exports :func:`~mmcontext.io.collect_adata_subset` (which now lives in
+:mod:`mmcontext.io`) so downstream consumers can keep importing both from this
+module, as ``mmcontext-benchmark`` does.
+"""
+
 from __future__ import annotations
 
 import logging
-from collections.abc import Sequence
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Literal
 
-import anndata as ad
 import pandas as pd
-import torch
-from datasets import Dataset, DatasetDict, load_dataset, load_from_disk
-from huggingface_hub import HfApi
+from datasets import Dataset, load_dataset, load_from_disk
 
-from mmcontext.file_utils import remove_corrupted_null_arrays
+# Re-export: collect_adata_subset moved to mmcontext.io during the v2.0 refactor;
+# keep the historical import path working for downstream consumers.
+from mmcontext.io import collect_adata_subset
 
 logger = logging.getLogger(__name__)
 
-
-@dataclass
-class SentenceDataset(Dataset):
-    """
-    A minimal Dataset that returns (index, text) tuples needed for batching.
-
-    Parameters
-    ----------
-    indices : List[int] | List[str]
-        Unique sample identifiers.
-    texts : List[str]
-        Sentences or tokenised strings.
-    """
-
-    indices: list[int]
-    texts: list[str]
-
-    def __len__(self):
-        return len(self.indices)
-
-    def __getitem__(self, idx):
-        return self.indices[idx], self.texts[idx]
+__all__ = ["collect_adata_subset", "load_generic_dataset"]
 
 
 def load_generic_dataset(
@@ -135,223 +118,3 @@ def load_generic_dataset(
         raise ValueError(f"Unknown format '{fmt}'")
 
     return ds
-
-
-def collect_adata_subset(
-    download_dir: str | Path | None = None,
-    sample_ids: Sequence[str] | None = None,
-    *,
-    file_paths: list[Path] | None = None,
-    join_vars: str = "outer",  # 'inner' if you need exact var overlap
-) -> ad.AnnData:
-    """
-    Load a subset of the data from the download directory or from specific file paths.
-
-    Gather a subset of observations that may live in **multiple** Zarr/H5AD
-    chunks into a single in-memory `AnnData`.
-
-    Parameters
-    ----------
-    download_dir : str | Path | None
-        Directory that contains files written by `get_initial_embeddings_from_adata_link`
-        (mix of `<dataset>.zarr/` folders and/or `<dataset>.h5ad` files).
-        If None, file_paths must be provided.
-    sample_ids : Sequence[str] | None
-        The observation indices to collect. Order will be preserved.
-        If None, all observations from the files will be collected.
-    file_paths : list[Path] | None
-        List of specific file paths to read from. If provided, download_dir is ignored.
-        Used when files are in their original locations (e.g., local datasets).
-    join_vars : {'outer', 'inner'}, default 'outer'
-        How to merge variables (genes) when concatenating chunks.
-
-    Returns
-    -------
-    AnnData
-        Combined object with `obs_names` matching the order of `sample_ids`.
-
-    Notes
-    -----
-    * For now, all chunks are loaded into memory.
-    * Either download_dir or file_paths must be provided.
-    """
-    if download_dir is None and file_paths is None:
-        raise ValueError("Either download_dir or file_paths must be provided")
-
-    if sample_ids is None:
-        raise ValueError("sample_ids must be provided")
-
-    if file_paths is not None:
-        # Use specific file paths
-        paths_to_scan = file_paths
-    else:
-        # Use directory scanning (existing behavior)
-        download_dir = Path(download_dir).expanduser()
-        if not download_dir.exists():
-            raise FileNotFoundError(download_dir)
-        paths_to_scan = list(download_dir.glob("*.zarr")) + list(download_dir.glob("*.h5ad"))
-
-    # keep a *mutable* set for O(1) look-up
-    remaining = set(sample_ids)
-    collected: list[ad.AnnData] = []
-
-    # scan zarr files/directories first (usually many small shards) -------------
-    for path in paths_to_scan:
-        if not remaining:
-            break
-
-        path = Path(path)
-        if path.suffix == ".zarr" or (path.is_dir() and path.name.endswith(".zarr")):
-            logger.debug("Scanning Zarr store %s", path)
-            # Fix corrupted null arrays before reading
-            remove_corrupted_null_arrays(path)
-            view = ad.read_zarr(path)
-            wanted = remaining.intersection(view.obs_names)
-            if wanted:
-                logger.info("  ↳ %d obs found in %s", len(wanted), path.name)
-                collected.append(view[list(wanted)].copy())  # loads only slice
-                remaining.difference_update(wanted)
-            if hasattr(view, "file") and view.file:
-                view.file.close()  # be nice to the OS
-
-    # scan *.h5ad files -----------------------------------------------
-    for path in paths_to_scan:
-        if not remaining:
-            break
-
-        path = Path(path)
-        if path.suffix == ".h5ad":
-            logger.debug("Scanning H5AD file %s", path.name)
-            view = ad.read_h5ad(path, backed="r")
-            wanted = remaining.intersection(view.obs_names)
-            if hasattr(view, "file") and view.file:
-                view.file.close()
-
-            if wanted:
-                logger.info("  ↳ %d obs found in %s", len(wanted), path.name)
-                full = ad.read_h5ad(path)  # loads entire file once
-                collected.append(full[list(wanted)].copy())
-                remaining.difference_update(wanted)
-
-    if remaining:
-        logger.warning(
-            "collect_adata_subset | %d sample_ids were *not* found in available files: %s",
-            len(remaining),
-            sorted(remaining)[:10],
-        )
-
-    if not collected:
-        raise ValueError("None of the requested sample_ids were found in any chunk.")
-
-    logger.info("Concatenating %d partial AnnData objects …", len(collected))
-    adata = ad.concat(
-        collected,
-        axis=0,
-        join=join_vars,
-        merge="same",
-        label=None,
-        index_unique=None,
-    )
-
-    # restore original order  -----------------------------------------
-    adata = adata[sample_ids].copy()
-    return adata
-
-
-def generate_revision_name(dataset_config: dict) -> str:
-    """
-    Generate a revision name for a dataset based on preprocessing parameters.
-
-    Parameters
-    ----------
-    dataset_config : dict
-        Dataset configuration containing gene_filter_strings, cs_length, etc.
-
-    Returns
-    -------
-    str
-        Revision name in format: [rps_rpl_mt][_cs{cs_length}]
-    """
-    parts = []
-
-    # Add gene filter strings if present
-    gene_filter_strings = dataset_config.get("gene_filter_strings", None)
-    if gene_filter_strings:
-        # Convert to lowercase and join with underscores
-        filter_str = "_".join([s.lower() for s in gene_filter_strings])
-        parts.append(filter_str)
-
-    # Add cs_length if specified
-    cs_length = dataset_config.get("cs_length", None)
-    if cs_length and cs_length > 0:
-        parts.append(f"cs{cs_length}")
-
-    return "_".join(parts) if parts else "processed"
-
-
-def check_revision_exists(dataset_source: str, revision: str) -> bool:
-    """
-    Check if a specific revision exists for a dataset on Hugging Face Hub.
-
-    Parameters
-    ----------
-    dataset_source : str
-        Full dataset source (e.g., "jo-mengr/dataset_name")
-    revision : str
-        Revision name to check
-
-    Returns
-    -------
-    bool
-        True if the revision exists, False otherwise
-    """
-    try:
-        api = HfApi()
-        # Try to get dataset info with specific revision
-        api.dataset_info(dataset_source, revision=revision)
-        return True
-    except Exception:
-        # If any error occurs (404, auth, etc.), assume revision doesn't exist
-        return False
-
-
-def push_dataset_revision(
-    dataset: Dataset,
-    dataset_source: str,
-    revision: str,
-    commit_message: str = None,
-) -> bool:
-    """
-    Push a processed dataset as a new revision to HuggingFace Hub.
-
-    Parameters
-    ----------
-    dataset : Dataset
-        The processed dataset to push
-    dataset_source : str
-        Full dataset source (e.g., "jo-mengr/dataset_name")
-    revision : str
-        Revision name for the processed dataset
-    commit_message : str, optional
-        Commit message for the revision
-
-    Returns
-    -------
-    bool
-        True if push was successful, False otherwise
-    """
-    try:
-        if commit_message is None:
-            commit_message = f"Add preprocessed dataset revision: {revision}"
-
-        logger.info(f"Pushing dataset as revision '{revision}' to {dataset_source}")
-        dataset.push_to_hub(
-            dataset_source,
-            revision=revision,
-            commit_message=commit_message,
-        )
-        logger.info(f"Successfully pushed revision '{revision}' for dataset '{dataset_source}'")
-        return True
-    except Exception as e:
-        logger.error(f"Failed to push revision '{revision}' for dataset '{dataset_source}': {e}")
-        return False
diff --git a/src/mmcontext/embed/embed_pipeline.py b/src/mmcontext/embed/embed_pipeline.py
deleted file mode 100644
index ccd9603..0000000
--- a/src/mmcontext/embed/embed_pipeline.py
+++ /dev/null
@@ -1,583 +0,0 @@
-import logging
-import os
-from collections.abc import Sequence
-from pathlib import Path
-
-import dotenv
-import pandas as pd
-import torch
-
-from mmcontext.embed.dataset_utils import (
-    check_revision_exists,
-    collect_adata_subset,
-    generate_revision_name,
-    load_generic_dataset,
-    push_dataset_revision,
-)
-from mmcontext.embed.model_utils import embed_labels, load_st_model, prepare_model_and_embed
-from mmcontext.file_utils import save_table
-from mmcontext.utils import truncate_cell_sentences
-
-dotenv.load_dotenv()
-logger = logging.getLogger(__name__)
-
-
-def check_required_files_exist(out_dir: Path, output_format: str) -> bool:
-    """
-    Check if required embedding files already exist in the output directory.
-
-    Parameters
-    ----------
-    out_dir : Path
-        Output directory path
-    output_format : str
-        Output format (e.g., 'parquet', 'csv')
-
-    Returns
-    -------
-    bool
-        True if required files exist, False otherwise
-    """
-    # Check for main embeddings file
-    embeddings_file = out_dir / f"embeddings.{output_format}"
-    meta_file = out_dir / "meta.yaml"
-
-    # Consider the output complete if both embeddings and meta files exist
-    return embeddings_file.exists() and meta_file.exists()
-
-
-def process_single_dataset_model(
-    ds_cfg, model_cfg, run_cfg, output_root, output_format, adata_cache, hf_cache=None
-) -> tuple[str, str, bool, str]:
-    """
-    Process a single dataset/model combination for embedding generation.
-
-    Parameters
-    ----------
-    ds_cfg : dict
-        Dataset configuration
-    model_cfg : dict
-        Model configuration
-    run_cfg : dict
-        Run configuration
-    output_root : str
-        Output root directory (resolved)
-    output_format : str
-        Output format (resolved)
-    adata_cache : str
-        AnnData cache directory (resolved)
-    hf_cache : str, optional
-        HuggingFace cache directory (resolved). If None, uses run_cfg.hf_cache for backward compatibility.
-
-    Returns
-    -------
-    tuple[str, str, bool, str]
-        Tuple of (dataset_name, model_id, success, error_msg)
-    """
-    dataset_name = ds_cfg.name
-    model_id = model_cfg.source
-    model_name = model_cfg.name
-    text_only = model_cfg.get("text_only", False)
-
-    # add a small string to the model indicating if it was used as text-only
-    if text_only:
-        model_name_for_path = model_name + "_text_only"
-    else:
-        model_name_for_path = model_name
-
-    logger.info(f"Processing {dataset_name} + {model_name_for_path}")
-
-    # Determine output directory
-    out_dir = (
-        Path(output_root)
-        / ds_cfg.name  # <— dataset-specific folder
-        / Path(model_name_for_path).name.replace("/", "_")
-    )
-
-    # Check if files already exist and skip if overwrite is False
-    if not run_cfg.overwrite and check_required_files_exist(out_dir, output_format):
-        logger.info(
-            f"Embeddings already exist for {ds_cfg.name} + {model_name_for_path}. Skipping. Use overwrite=true to force regeneration."
-        )
-        return dataset_name, model_name_for_path, True, "skipped_existing"
-
-    try:
-        # Load dataset
-        adata_download_dir = Path(adata_cache) / ds_cfg.name
-        # Use hf_cache parameter if provided, otherwise fall back to run_cfg.hf_cache for backward compatibility
-        cache_dir = hf_cache if hf_cache is not None else getattr(run_cfg, "hf_cache", None)
-
-        # Check if we should use revisions
-        use_revisions = getattr(run_cfg, "use_revisions", False)
-        revision_loaded = False
-
-        if use_revisions and ds_cfg.format == "hub" and text_only:
-            # Generate revision name based on preprocessing parameters
-            revision_name = generate_revision_name(ds_cfg)
-
-            # Check if revision exists and load it if available
-            if revision_name != "processed" and check_revision_exists(ds_cfg.source, revision_name):
-                logger.info(f"Found existing revision '{revision_name}' for dataset '{ds_cfg.name}', loading directly")
-                raw_ds = load_generic_dataset(
-                    source=ds_cfg.source,
-                    fmt=ds_cfg.format,
-                    split=ds_cfg.get("split", "test"),
-                    max_rows=run_cfg.n_rows,
-                    seed=run_cfg.seed,
-                    cache_dir=cache_dir,
-                    revision=revision_name,
-                )
-                revision_loaded = True
-                logger.info(f"Successfully loaded preprocessed dataset from revision '{revision_name}'")
-            else:
-                if revision_name != "processed":
-                    logger.info(
-                        f"Revision '{revision_name}' not found for dataset '{ds_cfg.name}', processing from scratch"
-                    )
-                # Load raw dataset
-                raw_ds = load_generic_dataset(
-                    source=ds_cfg.source,
-                    fmt=ds_cfg.format,
-                    split=ds_cfg.get("split", "test"),
-                    max_rows=run_cfg.n_rows,
-                    seed=run_cfg.seed,
-                    cache_dir=cache_dir,
-                )
-        else:
-            # Load raw dataset without revision
-            raw_ds = load_generic_dataset(
-                source=ds_cfg.source,
-                fmt=ds_cfg.format,
-                split=ds_cfg.get("split", "test"),
-                max_rows=run_cfg.n_rows,
-                seed=run_cfg.seed,
-                cache_dir=cache_dir,
-            )
-
-        numeric_data_available = "share_link" or "adata_link" in raw_ds.column_names
-
-        # Load model
-        st_model = load_st_model(model_id)
-        if text_only:
-            main_col = "cell_sentence_2"
-        else:
-            main_col = "cell_sentence_1"
-
-        # truncate cell sentences (only if revision wasn't loaded)
-        if not revision_loaded and ds_cfg.get("cs_length", None) is not None and main_col == "cell_sentence_2":
-            cs_length = ds_cfg.cs_length
-            filter_strings = ds_cfg.get("gene_filter_strings", None)
-            logger.info(f"Applying cell sentence truncation (cs_length={cs_length})")
-            raw_ds = truncate_cell_sentences(raw_ds, main_col, cs_length, filter_strings=filter_strings)
-
-            # Push processed dataset as new revision if enabled
-            if use_revisions and ds_cfg.format == "hub" and revision_name != "processed":
-                push_success = push_dataset_revision(
-                    raw_ds,
-                    ds_cfg.source,
-                    revision_name,
-                    commit_message=f"Processed dataset with {revision_name} settings",
-                )
-                if push_success:
-                    logger.info(f"Successfully pushed processed dataset as revision '{revision_name}'")
-                else:
-                    logger.warning(f"Failed to push processed dataset as revision '{revision_name}'")
-        elif revision_loaded:
-            logger.info("Skipping cell sentence truncation - using preprocessed revision")
-
-        # Generate embeddings
-        emb_df, path_map = prepare_model_and_embed(
-            st_model,
-            data=raw_ds,
-            main_col=main_col,
-            index_col=ds_cfg.index_col,
-            batch_size=run_cfg.batch_size,
-            num_workers=run_cfg.num_workers,
-            layer_key=getattr(model_cfg, "layer_key", None),  # For text only models, layer_key is not needed.
-            text_only=text_only,
-            adata_download_dir=adata_download_dir,
-            zenodo_token=os.getenv("ZENODO_TOKEN"),
-        )
-
-        if numeric_data_available:
-            # get sample_ids but without "sample_idx: prefix"
-            if "sample_idx:" in emb_df["sample_idx"][0]:
-                sample_ids = [sid.split(":")[1] for sid in emb_df["sample_idx"].tolist()]
-            elif "sample_idx" in emb_df.columns:
-                sample_ids = emb_df["sample_idx"].tolist()
-            else:
-                raise ValueError(f"sample_idx column not found in emb_df: {emb_df.columns}")
-
-            # Use path mapping if available (for local datasets), otherwise fall back to download_dir
-            if path_map is not None:
-                file_paths = list(path_map.values())
-                adata_subset = collect_adata_subset(
-                    sample_ids=sample_ids,
-                    file_paths=file_paths,
-                )
-            else:
-                adata_subset = collect_adata_subset(
-                    download_dir=adata_download_dir,
-                    sample_ids=sample_ids,
-                )
-        else:
-            adata_subset = None
-
-        # save the embeddings
-        save_table(
-            emb_df,
-            out_path=out_dir / "embeddings",
-            fmt=output_format,
-        )
-
-        # Save metadata
-        (out_dir / "meta.yaml").write_text(
-            f"model: {model_name_for_path}\ndataset: {ds_cfg.name}\nrows: {len(emb_df)}\n"
-        )
-
-        # Handle label embeddings if available
-        if adata_subset is not None:
-            subset_out = out_dir / "subset.h5ad"
-            adata_subset.write_h5ad(subset_out)
-            logger.info("Wrote subset AnnData → %s", subset_out)
-            # Define label types and their output prefixes
-            label_types = {"bio_label_list": "bio_label_embeddings", "batch_label_list": "batch_label_embeddings"}
-
-            # Loop over each label type
-            for label_list_attr, output_prefix in label_types.items():
-                if hasattr(ds_cfg, label_list_attr):
-                    label_cols = getattr(ds_cfg, label_list_attr)
-                    if label_cols is not None:
-                        for label_col in label_cols:
-                            if label_col in adata_subset.obs.columns:
-                                logger.info(f"Embedding {label_list_attr} from column: {label_col}")
-                                label_emb_df, label_to_index = embed_labels(
-                                    st_model,
-                                    adata_subset,
-                                    label_col,
-                                    batch_size=run_cfg.batch_size,
-                                    num_workers=run_cfg.num_workers,
-                                )
-                                # Save label embeddings
-                                save_table(
-                                    label_emb_df,
-                                    out_path=out_dir / f"{output_prefix}_{label_col}",
-                                    fmt=output_format,
-                                )
-
-                                # Save label mapping as JSON for downstream use
-                                import json
-
-                                mapping_path = out_dir / f"{output_prefix}_{label_col}_mapping.json"
-                                with open(mapping_path, "w") as f:
-                                    json.dump(label_to_index, f, indent=2)
-                                logger.info(f"Saved label mapping to {mapping_path}")
-                            else:
-                                logger.warning(f"Label column {label_col} not found in adata.obs")
-
-        logger.info(f"✓ Completed processing {dataset_name} + {model_name_for_path}")
-        return dataset_name, model_name_for_path, True, None
-
-    except Exception as e:
-        logger.error(f"✗ Error processing {dataset_name} + {model_name_for_path}: {e}")
-        return dataset_name, model_name_for_path, False, str(e)
-
-
-def _process_single_worker(args):
-    """
-    Worker function for parallel processing of embedding generation.
-
-    Parameters
-    ----------
-    args : tuple
-        Tuple of (ds_cfg, model_cfg, run_cfg, output_root, output_format, adata_cache)
-
-    Returns
-    -------
-    tuple[str, str, bool, str]
-        Tuple of (dataset_name, model_id, success, error_msg)
-    """
-    import os
-    import sys
-    import time
-    import traceback
-
-    import torch
-
-    ds_cfg, model_cfg, run_cfg, output_root, output_format, adata_cache = args
-    dataset_name = ds_cfg.name
-    model_id = model_cfg.source
-    text_only = model_cfg.get("text_only", False)
-    if text_only:
-        model_name_for_path = model_id + "_text_only"
-    else:
-        model_name_for_path = model_id
-
-    # Configure logging for this worker process
-    worker_logger = logging.getLogger(f"embed_worker_{os.getpid()}")
-    worker_logger.setLevel(logging.INFO)
-
-    # Create a handler that writes to stdout so we can see worker logs
-    if not worker_logger.handlers:
-        handler = logging.StreamHandler(sys.stdout)
-        handler.setFormatter(logging.Formatter(f"[Worker {os.getpid()}] %(levelname)s - %(message)s"))
-        worker_logger.addHandler(handler)
-
-    # Initialize CUDA context in worker process if available
-    if torch.cuda.is_available():
-        worker_logger.info(f"Initializing CUDA context - Device: {torch.cuda.get_device_name()}")
-        # Clear any cached memory from previous runs
-        torch.cuda.empty_cache()
-        # Set memory growth to avoid OOM in multi-process scenarios
-        torch.cuda.set_per_process_memory_fraction(0.8)  # Use 80% of GPU memory max
-    else:
-        worker_logger.warning("CUDA not available in worker process")
-
-    worker_logger.info(f"Starting processing: {dataset_name}/{model_name_for_path}")
-    start_time = time.time()
-
-    try:
-        dataset_name_result, model_id_result, success, error_msg = process_single_dataset_model(
-            ds_cfg, model_cfg, run_cfg, output_root, output_format, adata_cache
-        )
-
-        elapsed_time = time.time() - start_time
-        worker_logger.info(f"Completed processing: {dataset_name_result}/{model_id_result} in {elapsed_time:.1f}s")
-
-        # Clean up GPU memory after processing
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-        return dataset_name_result, model_id_result, success, error_msg
-
-    except Exception as e:
-        # Clean up GPU memory on error
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-        # Get full traceback for debugging
-        elapsed_time = time.time() - start_time
-        error_msg = f"{str(e)}\n{traceback.format_exc()}"
-        worker_logger.error(
-            f"Error processing {dataset_name}/{model_name_for_path} after {elapsed_time:.1f}s: {error_msg}"
-        )
-        return dataset_name, model_name_for_path, False, str(e)
-
-
-def embed_pipeline(cfg) -> None:
-    """
-    Orchestrate an end-to-end embedding generation run.
-
-    Parameters
-    ----------
-    cfg : omegaconf.DictConfig
-        Instantiated Hydra config. The schema matches ``conf/config.yaml``.
-    """
-    import concurrent.futures
-    import multiprocessing as mp
-    import time
-    from concurrent.futures import ProcessPoolExecutor, as_completed
-
-    print("\n=== Running Embedding Pipeline ===")
-
-    # Create list of all dataset/model combinations
-    # Resolve config values in main process to avoid Hydra interpolation issues in workers
-    tasks = []
-    for ds_cfg in cfg.datasets:
-        for model_cfg in cfg.models:
-            tasks.append((ds_cfg, model_cfg, cfg.run, cfg.output.root, cfg.output.format, cfg.output.adata_cache))
-
-    print(f"Total dataset/model combinations to process: {len(tasks)}")
-
-    # Check if parallel processing is enabled
-    enable_parallel = cfg.run.get("enable_parallel", False)
-    max_workers = cfg.run.get("max_workers", None)
-
-    if max_workers is None:
-        max_workers = mp.cpu_count()
-
-    start_time = time.time()
-
-    if not enable_parallel or len(tasks) == 1:
-        # Sequential processing
-        print("Running in sequential mode...")
-        successful = 0
-        failed = 0
-        skipped = 0
-
-        for i, (ds_cfg, model_cfg, run_cfg, output_root, output_format, adata_cache) in enumerate(tasks):
-            dataset_name = ds_cfg.name
-            model_id = model_cfg.source
-            model_name = model_cfg.name
-            text_only = model_cfg.get("text_only", False)
-            if text_only:
-                model_name_for_path = model_name + "_text_only"
-            else:
-                model_name_for_path = model_name
-
-            task_start = time.time()
-            print(f"\n=== [{i + 1}/{len(tasks)}] Processing dataset: {dataset_name}, model: {model_name_for_path} ===")
-
-            try:
-                dataset_name_result, model_id_result, success, error_msg = process_single_dataset_model(
-                    ds_cfg, model_cfg, run_cfg, output_root, output_format, adata_cache
-                )
-                task_time = time.time() - task_start
-
-                if success:
-                    if error_msg == "skipped_existing":
-                        skipped += 1
-                        print(f"↷ Skipped (already exists): {dataset_name_result}/{model_id_result} ({task_time:.1f}s)")
-                    else:
-                        successful += 1
-                        print(f"✓ Completed: {dataset_name_result}/{model_id_result} ({task_time:.1f}s)")
-                else:
-                    failed += 1
-                    print(f"✗ Failed: {dataset_name_result}/{model_id_result} ({task_time:.1f}s)")
-                    print(f"  Error: {error_msg}")
-
-                # Estimate remaining time
-                if i > 0:
-                    avg_time = (time.time() - start_time) / (i + 1)
-                    remaining_time = avg_time * (len(tasks) - i - 1)
-                    print(f"  → Estimated time remaining: {remaining_time:.1f}s")
-
-            except Exception as e:
-                failed += 1
-                print(f"✗ Error processing {dataset_name}/{model_name_for_path}: {e}")
-                print("  Continuing with next combination...")
-                continue
-
-        total_time = time.time() - start_time
-        print("\n=== Summary ===")
-        print(f"Total tasks: {len(tasks)}")
-        print(f"Successful: {successful}")
-        print(f"Skipped: {skipped}")
-        print(f"Failed: {failed}")
-        print(f"Success rate: {successful / len(tasks) * 100:.1f}%")
-        print(f"Total time: {total_time:.1f}s")
-        print(f"Average time per task: {total_time / len(tasks):.1f}s")
-
-    else:
-        # Parallel processing
-        print(f"Running in parallel mode with {max_workers} workers...")
-
-        # Get timeout from config (default: 2 hours per task for embedding)
-        task_timeout = cfg.run.get("task_timeout", 7200)
-        print(f"Task timeout: {task_timeout}s ({task_timeout / 60:.1f} minutes)")
-
-        # Choose executor based on GPU availability and configuration
-        use_threads_for_gpu = cfg.run.get("use_threads_for_gpu", True)  # Default to threads for GPU
-
-        if torch.cuda.is_available() and use_threads_for_gpu:
-            # Use ThreadPoolExecutor for GPU tasks to share GPU memory efficiently
-            from concurrent.futures import ThreadPoolExecutor
-
-            print("Using ThreadPoolExecutor for GPU processing (better memory sharing)")
-            executor_class = ThreadPoolExecutor
-            worker_func = process_single_dataset_model  # Use direct function for threads
-            task_args_func = lambda task: task  # Pass args directly for threads
-        else:
-            # Use ProcessPoolExecutor for CPU tasks or when explicitly requested
-            from concurrent.futures import ProcessPoolExecutor
-
-            print("Using ProcessPoolExecutor")
-            executor_class = ProcessPoolExecutor
-            worker_func = _process_single_worker  # Use wrapper for processes
-            task_args_func = lambda task: (task,)  # Wrap in tuple for processes
-
-        with executor_class(max_workers=max_workers) as executor:
-            # Submit all tasks
-            future_to_task = {}
-            for i, task in enumerate(tasks):
-                ds_cfg, model_cfg, run_cfg, output_root, output_format, adata_cache = task
-                dataset_name = ds_cfg.name
-                model_id = model_cfg.source
-                model_name = model_cfg.name
-                text_only = model_cfg.get("text_only", False)
-                if text_only:
-                    model_name_for_path = model_id + "_text_only"
-                else:
-                    model_name_for_path = model_id
-
-                future = executor.submit(worker_func, *task_args_func(task))
-                future_to_task[future] = (dataset_name, model_name_for_path, i, time.time())
-
-            # Process completed tasks with timeout
-            completed = 0
-            successful = 0
-            failed = 0
-            skipped = 0
-            timed_out = 0
-
-            try:
-                for future in as_completed(future_to_task, timeout=task_timeout + 60):  # Small buffer for cleanup
-                    dataset_name, model_name_for_path, task_idx, submit_time = future_to_task[future]
-                    completed += 1
-                    task_time = time.time() - submit_time
-
-                    try:
-                        # Get result with timeout
-                        dataset_name_result, model_id_result, success, error_msg = future.result(timeout=10)
-
-                        if success:
-                            if error_msg == "skipped_existing":
-                                skipped += 1
-                                print(
-                                    f"↷ [{completed}/{len(tasks)}] Skipped: {dataset_name_result}/{model_id_result} ({task_time:.1f}s)"
-                                )
-                            else:
-                                successful += 1
-                                print(
-                                    f"✓ [{completed}/{len(tasks)}] Completed: {dataset_name_result}/{model_id_result} ({task_time:.1f}s)"
-                                )
-                        else:
-                            failed += 1
-                            print(
-                                f"✗ [{completed}/{len(tasks)}] Failed: {dataset_name_result}/{model_id_result} ({task_time:.1f}s)"
-                            )
-                            print(f"  Error: {error_msg}")
-
-                        # Progress update
-                        if completed % max(1, len(tasks) // 10) == 0:  # Show progress every 10%
-                            elapsed = time.time() - start_time
-                            if completed > 0:
-                                eta = (elapsed / completed) * (len(tasks) - completed)
-                                print(f"  → Progress: {completed / len(tasks) * 100:.1f}% | ETA: {eta:.1f}s")
-
-                    except concurrent.futures.TimeoutError:
-                        timed_out += 1
-                        failed += 1
-                        print(
-                            f"✗ [{completed}/{len(tasks)}] Timeout: {dataset_name}/{model_name_for_path} (>{task_timeout}s)"
-                        )
-                        future.cancel()  # Try to cancel the future
-
-                    except Exception as e:
-                        failed += 1
-                        print(
-                            f"✗ [{completed}/{len(tasks)}] Exception in worker for {dataset_name}/{model_name_for_path}: {e}"
-                        )
-
-            except concurrent.futures.TimeoutError:
-                print("⚠ Warning: Overall timeout reached. Some tasks may not have completed.")
-                # Handle any remaining uncompleted tasks
-                remaining_tasks = len(tasks) - completed
-                if remaining_tasks > 0:
-                    print(f"⚠ Warning: {remaining_tasks} tasks did not complete")
-                    failed += remaining_tasks
-
-            total_time = time.time() - start_time
-            print("\n=== Summary ===")
-            print(f"Total tasks: {len(tasks)}")
-            print(f"Successful: {successful}")
-            print(f"Skipped: {skipped}")
-            print(f"Failed: {failed}")
-            if timed_out > 0:
-                print(f"Timed out: {timed_out}")
-            print(f"Success rate: {successful / len(tasks) * 100:.1f}%")
-            print(f"Total time: {total_time:.1f}s")
-            print(f"Average time per task: {total_time / len(tasks):.1f}s")
-            print(f"Speedup vs sequential: ~{max_workers}x (theoretical)")
-
-    print("\n=== Embedding pipeline completed ===")
diff --git a/src/mmcontext/embed/encode.py b/src/mmcontext/embed/encode.py
new file mode 100644
index 0000000..c0adc35
--- /dev/null
+++ b/src/mmcontext/embed/encode.py
@@ -0,0 +1,374 @@
+"""Embed helpers for running inference with mmcontext (and mmcontext-like) models.
+
+This module owns the *consumer-facing* glue for turning a loaded
+:class:`~sentence_transformers.SentenceTransformer` (which may be a legacy
+:class:`~mmcontext._legacy.mmcontextencoder.MMContextEncoder`) plus a dataset
+into an embedding table. It encodes mmcontext's own model contract — downloading
+numeric initial embeddings from a store link and registering them via
+``register_initial_embeddings`` — so downstream consumers do not have to reach
+into ``mmcontext._legacy`` themselves.
+
+Consumers:
+
+* ``mmcontext-benchmark`` — the ``MMContextAdapter`` legacy/v2 embed paths.
+* ``sc_language_foundation_models`` — the per-model ``predict.py`` scripts
+  (which still import from the back-compat :mod:`mmcontext.embed.model_utils`
+  shim).
+
+Only the helpers actually used by those consumers are kept here; the wider
+multi-model orchestration removed in the v2.0 cleanup is intentionally *not*
+restored.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+import torch
+from datasets import Dataset as HFDataset
+from sentence_transformers import SentenceTransformer
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from mmcontext._legacy.mmcontextencoder import MMContextEncoder as MMEnc
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SentenceDataset(Dataset):
+    """
+    A minimal Dataset that returns (index, text) tuples needed for batching.
+
+    Parameters
+    ----------
+    indices : list[int] | list[str]
+        Unique sample identifiers.
+    texts : list[str]
+        Sentences or tokenised strings.
+    """
+
+    indices: list[int]
+    texts: list[str]
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx):
+        return self.indices[idx], self.texts[idx]
+
+
+class HFIndexedDataset(torch.utils.data.Dataset):
+    """
+    Class to create a dataset that is used to embed the data.
+
+    A lightweight view on top of a 🤗 Dataset that returns
+    ``(sample_index, text)`` tuples.
+
+    Parameters
+    ----------
+    ds : datasets.Dataset
+        The underlying dataset.
+    main_col : str
+        Name of the column containing the raw sentence.
+    index_col : str
+        Name of the identifier column (kept for the final join).
+    """
+
+    def __init__(self, ds: Dataset, main_col: str, index_col: str):
+        self.ds = ds
+        self.main_col = main_col
+        self.index_col = index_col
+
+    def __len__(self):
+        return len(self.ds)
+
+    def __getitem__(self, idx):
+        row = self.ds[int(idx)]
+        return row[self.index_col], row[self.main_col]
+
+
+def load_st_model(model_id: str | Path) -> SentenceTransformer:
+    """Load a Sentence-Transformer model from HF Hub or local path with GPU support."""
+    import os
+
+    model_id = str(model_id)
+    logger.info("Loading model %s", model_id)
+
+    # Check device availability in worker process
+    if torch.cuda.is_available():
+        device = "cuda"
+        logger.info(f"CUDA available in worker process {os.getpid()}: {torch.cuda.get_device_name()}")
+    else:
+        device = "cpu"
+        logger.warning(f"CUDA not available in worker process {os.getpid()}, using CPU")
+
+    # Load model and explicitly move to device
+    model = SentenceTransformer(model_id, device=device, trust_remote_code=True)
+    logger.info(f"Model loaded on device: {model.device}")
+
+    return model
+
+
+def prepare_model_and_embed(
+    st_model: SentenceTransformer,
+    data: HFDataset,
+    *,
+    layer_key: str = None,
+    adata_download_dir: str | Path = None,
+    main_col: str = "cell_sentence_1",
+    index_col: str = "sample_index",
+    batch_size: int = 32,
+    num_workers: int = 0,
+    pin_memory: bool = torch.cuda.is_available(),
+    axis: Literal["var", "obs"] = "obs",
+    text_only: bool = False,
+    overwrite: bool = True,
+    zenodo_token: str | None = None,
+) -> tuple[pd.DataFrame, dict[str, Path] | None]:
+    """
+    Create embeddings using a Sentence-Transformer model, with some model-specific preparation.
+
+    Prepare a customised Sentence-Transformer model *if possible* and obtain
+    sentence embeddings in mini-batches.
+
+    Parameters
+    ----------
+    st_model : SentenceTransformer
+        The model wrapper.  The *first* module (``st_model[0]``) may expose
+        extra methods such as ``prefix_ds``.
+    df_subset : DataFrame
+        Test-set slice (≤ 5 k rows by default) **must** contain
+        ``main_col`` and ``index_col``.
+    main_col : str
+        Column containing the raw sentences.
+    index_col : str, default ``"sample_index"``
+        Column to keep as the identifier in the returned table.
+    batch_size : int, default 32
+        Number of sentences per encode() call.
+    num_workers : int, default 0
+        Worker processes for the PyTorch data loader.
+    pin_memory : bool, default ``True`` if CUDA is available
+        Forwarded to DataLoader for speed on GPUs.
+    layer_key : str
+        The name of the layer key to be used for the embedding.
+    axis : str, default ``"obs"``
+        The axis to be used for the embedding. Can be either ``"obs"`` or
+        ``"var"``.
+    text_only : bool, default ``False``
+        If ``True``, all input will be treated as text. If "False", the dataset has to contain
+        a column share_link, which points to a zarr store that can be used to get initial embeddings for each token.
+    overwrite : bool, default ``True``
+        If ``True``, the initial embeddings will be re-downloaded even if they already exist.
+    zenodo_token : str | None, optional
+        Zenodo access token for authenticating draft record downloads.
+        Required for draft records, optional for published records.
+
+    Returns
+    -------
+    tuple[DataFrame, dict[str, Path] | None]
+        A tuple containing:
+        - DataFrame with columns ``[index_col, "embedding"]`` where *embedding* is a list[float].
+        - Path mapping from original links to actual file locations (None if no numeric data)
+
+    Notes
+    -----
+    * Optional methods are only executed if present.
+    * ``token_df`` is **registered** with the model (if possible) but **not**
+      returned to the caller.
+    """
+    if not text_only and (layer_key is None or adata_download_dir is None):
+        raise ValueError("layer_key and adata_download_dir must be provided if text_only is False")
+    # If using a multiplets dataset, the column with the data representation is called "anchor". For the embedding workflow, only this data is embedded.
+    if "anchor" in data.column_names:
+        main_col = "anchor"
+    ###########################################################
+    # --- optional model-specific preparation ----------------
+    ###########################################################
+    impl = st_model[0] if len(st_model) > 0 else st_model
+    has_numeric = ("share_link" in data.column_names) or (
+        "adata_link" in data.column_names
+    )  # otherwise we cannot download the initial embeddings and register them with the model
+    path_map = None
+    if data is not None and has_numeric:
+        link_column = "share_link" if "share_link" in data.column_names else "adata_link"
+        # If the dataset is available, download get the token_df, even if the models doesnt use it. Just so the data is downloaded and the adata subset can be created downstream
+        logger.info("Extracting numeric intital embeddings from dataset via it's share_links …")
+        token_df, path_map = MMEnc.get_initial_embeddings_from_adata_link(
+            data,
+            layer_key=layer_key,
+            axis=axis,
+            download_dir=adata_download_dir,
+            overwrite=overwrite,
+            link_column=link_column,
+            zenodo_token=zenodo_token,
+        )  # type: ignore[arg-type]
+    else:
+        logger.info("""While the model supports initial embeddings, the dataset does not provide them.
+                    Therefore cell sentences will be analysed as texts.""")
+        token_df = None
+        # register only if the method exists
+    if hasattr(impl, "register_initial_embeddings") and token_df is not None and not text_only:
+        logger.info("Registering initial embeddings with MMContextEncoder module of SentenceTransformer Model …")
+        impl.register_initial_embeddings(token_df, data_origin=layer_key.split("X_")[1])  # type: ignore[arg-type]
+
+    if hasattr(impl, "prefix_ds") and not text_only:
+        logger.info("Calling prefix_ds on model-specific module …")
+        # Step 4: Apply prefixes using the new simplified prefix_ds method
+        logger.info(f"Applying prefixes to columns: {main_col}")
+        ds = impl.prefix_ds(
+            ds=data,
+            columns_to_prefix=main_col,
+        )  # type: ignore[arg-type]
+    else:
+        ds = data
+    ###########################################################
+    # --- sentence embeddings via encode() --------------------
+    ###########################################################
+    # -------- create a torch-compatible dataset ------------------------
+    if isinstance(ds, pd.DataFrame):
+        torch_ds = SentenceDataset(
+            indices=ds[index_col].tolist(),
+            texts=ds[main_col].astype(str).tolist(),
+        )
+    else:  # Hugging-Face Dataset
+        torch_ds = HFIndexedDataset(ds, main_col=main_col, index_col=index_col)
+
+    loader = DataLoader(
+        torch_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+    )
+
+    all_indices: list[int | str] = []
+    # Preallocate a single contiguous float32 buffer instead of accumulating a
+    # list[list[float]] (~1 GB of boxed Python floats for 100k × 768). The width
+    # D is only known after the first batch, so allocate lazily then fill by
+    # offset.
+    n_total = len(torch_ds)
+    out: np.ndarray | None = None
+    offset = 0
+
+    st_model.eval()
+    with torch.inference_mode():
+        for batch_indices, batch_texts in tqdm(loader, desc="Encoding"):
+            # DataLoader returns tuples → convert to list[str]
+            embeddings = st_model.encode(
+                list(batch_texts),
+                batch_size=len(batch_texts),
+                convert_to_numpy=True,
+                show_progress_bar=False,
+            )
+            embeddings = np.asarray(embeddings, dtype=np.float32)
+            if out is None:
+                out = np.empty((n_total, embeddings.shape[1]), dtype=np.float32)
+            out[offset : offset + embeddings.shape[0]] = embeddings
+            offset += embeddings.shape[0]
+            if isinstance(batch_indices, torch.Tensor):
+                batch_indices = batch_indices.tolist()
+            all_indices.extend(batch_indices)
+
+    if out is None:
+        out = np.empty((0, 0), dtype=np.float32)
+    else:
+        out = out[:offset]  # trim in case n_total over-counted (e.g. drop_last)
+
+    # Each "embedding" cell is a 1-D float32 view into the contiguous buffer.
+    emb_df = pd.DataFrame({"sample_idx": all_indices})
+    emb_df["embedding"] = list(out)
+    logger.info("Generated %d embeddings", len(emb_df))
+    return emb_df, path_map
+
+
+def create_label_dataset(adata, label_col: str) -> tuple[HFDataset, dict[str, int]]:
+    """
+    Create a small HF Dataset from unique labels in an AnnData object.
+
+    Parameters
+    ----------
+    adata : anndata.AnnData
+        AnnData object containing the labels
+    label_col : str
+        Name of the column in adata.obs containing the labels
+
+    Returns
+    -------
+    tuple[HFDataset, dict[str, int]]
+        A tuple containing:
+        - HuggingFace dataset with unique labels and their indices
+        - Mapping from label string to embedding index
+    """
+    from datasets import Dataset
+
+    # Get unique labels only
+    unique_labels = adata.obs[label_col].astype(str).unique().tolist()
+
+    # Create mapping from label to embedding index
+    label_to_index = {label: idx for idx, label in enumerate(unique_labels)}
+
+    # Create indices for the unique labels (just sequential)
+    indices = list(range(len(unique_labels)))
+
+    # Create dataset with unique labels only
+    ds = Dataset.from_dict({"sample_idx": indices, "label": unique_labels})
+
+    return ds, label_to_index
+
+
+def embed_labels(
+    st_model: SentenceTransformer,
+    adata,
+    label_col: str,
+    *,
+    batch_size: int = 32,
+    num_workers: int = 0,
+    pin_memory: bool = torch.cuda.is_available(),
+) -> tuple[pd.DataFrame, dict[str, int]]:
+    """
+    Embed unique labels from an AnnData object using a SentenceTransformer model.
+
+    Parameters
+    ----------
+    st_model : SentenceTransformer
+        The model to use for embedding
+    adata : anndata.AnnData
+        AnnData object containing the labels
+    label_col : str
+        Name of the column in adata.obs containing the labels
+    batch_size : int, default 32
+        Number of labels per encode() call
+    num_workers : int, default 0
+        Worker processes for the PyTorch data loader
+    pin_memory : bool, default True if CUDA is available
+        Forwarded to DataLoader for speed on GPUs
+
+    Returns
+    -------
+    tuple[pd.DataFrame, dict[str, int]]
+        A tuple containing:
+        - DataFrame with columns ['sample_idx', 'embedding'] containing unique label embeddings
+        - Mapping from label string to embedding index in the DataFrame
+    """
+    # Create dataset from unique labels
+    label_ds, label_to_index = create_label_dataset(adata, label_col)
+
+    # Use prepare_model_and_embed with text_only=True since we're just embedding text labels
+    emb_df, _ = prepare_model_and_embed(
+        st_model,
+        data=label_ds,
+        main_col="label",
+        index_col="sample_idx",
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        text_only=True,
+    )
+
+    return emb_df, label_to_index
diff --git a/src/mmcontext/embed/mining.py b/src/mmcontext/embed/mining.py
new file mode 100644
index 0000000..04b5b10
--- /dev/null
+++ b/src/mmcontext/embed/mining.py
@@ -0,0 +1,240 @@
+"""Hard-negative mining helpers for the MMContext pipeline.
+
+These wrap :func:`sentence_transformers.util.mine_hard_negatives` so the same
+core can back both an offline exploration script
+(``scripts/explore_hard_negatives.py``) and a future in-training re-mining
+callback for ``scripts/train_config.py``.
+
+Mining only makes sense against a model that already has a meaningful joint
+embedding space — mining with a freshly initialised bimodal model produces
+noise. The input is always an ``anchor`` / ``positive`` dataset as produced by
+:func:`~mmcontext.embed.prepare_dataset` (with ``use_hard_negatives=False`` for
+the omics path, so the dataset's own pre-baked negatives are dropped first):
+
+* For ``modality="bimodal"`` datasets the ``anchor`` column holds omics ids
+  (``"omics:..."``) that the model resolves through its attached VectorStore;
+  the ``positive`` texts form the corpus, so mined negatives are confusable
+  cell descriptions in the shared space.
+* For text datasets both columns are plain text.
+
+Because the sentence-transformers trainer samples each batch from a single
+dataset (in-batch negatives stay within-dataset), mining should likewise be run
+**per dataset** rather than over a merged corpus.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from datasets import Dataset
+    from sentence_transformers import SentenceTransformer
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MiningConfig:
+    """Parameters forwarded to :func:`sentence_transformers.util.mine_hard_negatives`.
+
+    The defaults mirror ``scripts/finetune_tiny.py`` and are a sensible starting
+    point for an already-aligned model. ``max_score`` is intentionally on by
+    default: as the model sharpens, the closest candidates increasingly become
+    *false* negatives (descriptions that are actually correct), and the ceiling
+    drops them.
+
+    Parameters
+    ----------
+    num_negatives : int, default 3
+        Negatives to mine per anchor.
+    range_min : int, default 1
+        Skip the N closest candidates (rank 0 is usually the true positive).
+    range_max : int or None, default None
+        Only consider candidates up to this rank (``None`` = no upper bound).
+    max_score : float or None, default 0.95
+        Drop candidates scoring above this ceiling (false-negative guard).
+    min_score : float or None, default None
+        Drop candidates scoring below this floor (too-easy guard).
+    absolute_margin : float or None, default None
+        Negative score must be at least this far below the positive score.
+    relative_margin : float or None, default None
+        Negative score must be below ``relative_margin * positive_score``
+        (e.g. ``0.95``).
+    sampling_strategy : {"top", "random"}, default "top"
+        Which qualifying candidates to keep.
+    use_faiss : bool, default False
+        Use FAISS for the similarity search (recommended for large corpora).
+    batch_size : int, default 32
+        Encoding batch size for the mining search.
+    """
+
+    num_negatives: int = 3
+    range_min: int = 1
+    range_max: int | None = None
+    max_score: float | None = 0.95
+    min_score: float | None = None
+    absolute_margin: float | None = None
+    relative_margin: float | None = None
+    sampling_strategy: Literal["top", "random"] = "top"
+    use_faiss: bool = False
+    batch_size: int = 32
+
+    def as_kwargs(self) -> dict[str, Any]:
+        """Return the config as ``mine_hard_negatives`` keyword arguments."""
+        return {
+            "num_negatives": self.num_negatives,
+            "range_min": self.range_min,
+            "range_max": self.range_max,
+            "max_score": self.max_score,
+            "min_score": self.min_score,
+            "absolute_margin": self.absolute_margin,
+            "relative_margin": self.relative_margin,
+            "sampling_strategy": self.sampling_strategy,
+            "use_faiss": self.use_faiss,
+            "batch_size": self.batch_size,
+        }
+
+
+def mine_negatives(
+    pairs: Dataset,
+    model: SentenceTransformer,
+    cfg: MiningConfig,
+    *,
+    anchor_column_name: str = "anchor",
+    positive_column_name: str = "positive",
+    output_format: Literal["triplet", "n-tuple", "labeled-pair", "labeled-list"] = "n-tuple",
+    output_scores: bool = False,
+    verbose: bool = True,
+) -> Dataset:
+    """Mine model-specific hard negatives for an ``anchor`` / ``positive`` dataset.
+
+    Thin wrapper around :func:`sentence_transformers.util.mine_hard_negatives`
+    that applies a :class:`MiningConfig`. The model is switched to ``eval`` mode
+    for the duration of mining and restored to its prior mode afterwards.
+
+    Parameters
+    ----------
+    pairs : datasets.Dataset
+        Dataset with at least *anchor_column_name* and *positive_column_name*.
+        For the bimodal path the anchors must be resolvable by the model's
+        attached VectorStore.
+    model : sentence_transformers.SentenceTransformer
+        An *already-trained* model (mining a fresh model yields noise).
+    cfg : MiningConfig
+        Mining parameters.
+    anchor_column_name, positive_column_name : str
+        Column names to use as query and corpus.
+    output_format : {"triplet", "n-tuple", "labeled-pair", "labeled-list"}
+        Forwarded to ``mine_hard_negatives``. ``"n-tuple"`` gives
+        ``anchor, positive, negative_1, ...`` which is what
+        :class:`~sentence_transformers.losses.MultipleNegativesRankingLoss`
+        expects.
+    output_scores : bool, default False
+        If True, ``mine_hard_negatives`` appends similarity-score columns.
+    verbose : bool, default True
+        Print mining statistics.
+
+    Returns
+    -------
+    datasets.Dataset
+        The mined dataset in the requested *output_format*.
+    """
+    from sentence_transformers.util import mine_hard_negatives
+
+    was_training = model.training
+    model.eval()
+    try:
+        mined = mine_hard_negatives(
+            pairs,
+            model,
+            anchor_column_name=anchor_column_name,
+            positive_column_name=positive_column_name,
+            output_format=output_format,
+            output_scores=output_scores,
+            verbose=verbose,
+            **cfg.as_kwargs(),
+        )
+    finally:
+        if was_training:
+            model.train()
+
+    logger.info(
+        "Mined %d rows from %d pairs (columns=%s)",
+        len(mined),
+        len(pairs),
+        mined.column_names,
+    )
+    return mined
+
+
+def mining_report(
+    mined: Dataset,
+    model: SentenceTransformer,
+    *,
+    anchor_column_name: str = "anchor",
+    positive_column_name: str = "positive",
+    max_rows: int | None = None,
+) -> pd.DataFrame:
+    """Build a human-readable table of what was mined, with cosine sims.
+
+    Re-encodes the anchor, positive and each ``negative_*`` column (cosine,
+    via normalised embeddings) so each row shows the positive similarity, each
+    negative similarity, and the positive-minus-negative margin. This is
+    version-independent (it does not rely on ``output_scores``) and gives the
+    positive similarity, which the mining call itself does not return.
+
+    Parameters
+    ----------
+    mined : datasets.Dataset
+        An ``n-tuple`` mined dataset (``anchor``, ``positive``, ``negative_*``).
+    model : sentence_transformers.SentenceTransformer
+        The model used for mining (used to recompute similarities).
+    anchor_column_name, positive_column_name : str
+        Column names in *mined*.
+    max_rows : int or None
+        Limit the report to the first N rows (``None`` = all).
+
+    Returns
+    -------
+    pandas.DataFrame
+        One row per anchor: anchor, positive, ``pos_sim``, and for each mined
+        negative its text, ``neg_<i>_sim`` and ``margin_<i>`` (pos − neg).
+    """
+    import numpy as np
+    import pandas as pd
+
+    ds = mined.select(range(min(max_rows, len(mined)))) if max_rows else mined
+    neg_cols = sorted(
+        (c for c in ds.column_names if c.startswith("negative") and not c.endswith("_idx")),
+        key=lambda c: (len(c), c),
+    )
+
+    def _enc(texts: list[str]) -> np.ndarray:
+        return model.encode(texts, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
+
+    anchor_emb = _enc(list(ds[anchor_column_name]))
+    pos_emb = _enc(list(ds[positive_column_name]))
+    pos_sim = np.sum(anchor_emb * pos_emb, axis=1)
+
+    rows: list[dict[str, Any]] = []
+    for i in range(len(ds)):
+        rows.append(
+            {
+                "anchor": ds[anchor_column_name][i],
+                "positive": ds[positive_column_name][i],
+                "pos_sim": round(float(pos_sim[i]), 4),
+            }
+        )
+
+    for j, col in enumerate(neg_cols, start=1):
+        neg_emb = _enc(list(ds[col]))
+        neg_sim = np.sum(anchor_emb * neg_emb, axis=1)
+        for i in range(len(ds)):
+            rows[i][f"negative_{j}"] = ds[col][i]
+            rows[i][f"neg_{j}_sim"] = round(float(neg_sim[i]), 4)
+            rows[i][f"margin_{j}"] = round(float(pos_sim[i] - neg_sim[i]), 4)
+
+    return pd.DataFrame(rows)
diff --git a/src/mmcontext/embed/model_utils.py b/src/mmcontext/embed/model_utils.py
index c90e9f7..a927868 100644
--- a/src/mmcontext/embed/model_utils.py
+++ b/src/mmcontext/embed/model_utils.py
@@ -1,656 +1,27 @@
-from __future__ import annotations
-
-import logging
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Literal, Union
-
-import anndata as ad
-import numpy as np
-import pandas as pd
-import torch
-from datasets import Dataset as HFDataset
-from sentence_transformers import SentenceTransformer
-from torch.utils.data import DataLoader, Dataset
-from tqdm import tqdm
-
-from mmcontext.mmcontextencoder import MMContextEncoder as MMEnc
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class SentenceDataset(Dataset):
-    """
-    A minimal Dataset that returns (index, text) tuples needed for batching.
-
-    Parameters
-    ----------
-    indices : list[int] | list[str]
-        Unique sample identifiers.
-    texts : list[str]
-        Sentences or tokenised strings.
-    """
-
-    indices: list[int]
-    texts: list[str]
-
-    def __len__(self):
-        return len(self.indices)
-
-    def __getitem__(self, idx):
-        return self.indices[idx], self.texts[idx]
-
-
-class HFIndexedDataset(torch.utils.data.Dataset):
-    """
-    Class to create a dataset that is used to embed the data.
-
-    A lightweight view on top of a 🤗 Dataset that returns
-    ``(sample_index, text)`` tuples.
-
-    Parameters
-    ----------
-    ds : datasets.Dataset
-        The underlying dataset.
-    main_col : str
-        Name of the column containing the raw sentence.
-    index_col : str
-        Name of the identifier column (kept for the final join).
-    """
-
-    def __init__(self, ds: Dataset, main_col: str, index_col: str):
-        self.ds = ds
-        self.main_col = main_col
-        self.index_col = index_col
-
-    def __len__(self):
-        return len(self.ds)
-
-    def __getitem__(self, idx):
-        row = self.ds[int(idx)]
-        return row[self.index_col], row[self.main_col]
-
-
-def load_st_model(model_id: str | Path) -> SentenceTransformer:
-    """Load a Sentence-Transformer model from HF Hub or local path with GPU support."""
-    import os
-
-    model_id = str(model_id)
-    logger.info("Loading model %s", model_id)
-
-    # Check device availability in worker process
-    if torch.cuda.is_available():
-        device = "cuda"
-        logger.info(f"CUDA available in worker process {os.getpid()}: {torch.cuda.get_device_name()}")
-    else:
-        device = "cpu"
-        logger.warning(f"CUDA not available in worker process {os.getpid()}, using CPU")
-
-    # Load model and explicitly move to device
-    model = SentenceTransformer(model_id, device=device, trust_remote_code=True)
-    logger.info(f"Model loaded on device: {model.device}")
-
-    return model
-
-
-def prepare_model_and_embed(
-    st_model: SentenceTransformer,
-    data: HFDataset,
-    *,
-    layer_key: str = None,
-    adata_download_dir: str | Path = None,
-    main_col: str = "cell_sentence_1",
-    index_col: str = "sample_index",
-    batch_size: int = 32,
-    num_workers: int = 0,
-    pin_memory: bool = torch.cuda.is_available(),
-    axis: Literal["var", "obs"] = "obs",
-    text_only: bool = False,
-    overwrite: bool = True,
-    zenodo_token: str | None = None,
-) -> tuple[pd.DataFrame, dict[str, Path] | None]:
-    """
-    Create embeddings using a Sentence-Transformer model, with some model-specific preparation.
-
-    Prepare a customised Sentence-Transformer model *if possible* and obtain
-    sentence embeddings in mini-batches.
-
-    Parameters
-    ----------
-    st_model : SentenceTransformer
-        The model wrapper.  The *first* module (``st_model[0]``) may expose
-        extra methods such as ``prefix_ds``.
-    df_subset : DataFrame
-        Test-set slice (≤ 5 k rows by default) **must** contain
-        ``main_col`` and ``index_col``.
-    main_col : str
-        Column containing the raw sentences.
-    index_col : str, default ``"sample_index"``
-        Column to keep as the identifier in the returned table.
-    batch_size : int, default 32
-        Number of sentences per encode() call.
-    num_workers : int, default 0
-        Worker processes for the PyTorch data loader.
-    pin_memory : bool, default ``True`` if CUDA is available
-        Forwarded to DataLoader for speed on GPUs.
-    layer_key : str
-        The name of the layer key to be used for the embedding.
-    axis : str, default ``"obs"``
-        The axis to be used for the embedding. Can be either ``"obs"`` or
-        ``"var"``.
-    text_only : bool, default ``False``
-        If ``True``, all input will be treated as text. If "False", the dataset has to contain
-        a column share_link, which points to a zarr store that can be used to get initial embeddings for each token.
-    overwrite : bool, default ``True``
-        If ``True``, the initial embeddings will be re-downloaded even if they already exist.
-    zenodo_token : str | None, optional
-        Zenodo access token for authenticating draft record downloads.
-        Required for draft records, optional for published records.
-
-    Returns
-    -------
-    tuple[DataFrame, dict[str, Path] | None]
-        A tuple containing:
-        - DataFrame with columns ``[index_col, "embedding"]`` where *embedding* is a list[float].
-        - Path mapping from original links to actual file locations (None if no numeric data)
-
-    Notes
-    -----
-    * Optional methods are only executed if present.
-    * ``token_df`` is **registered** with the model (if possible) but **not**
-      returned to the caller.
-    """
-    if not text_only and (layer_key is None or adata_download_dir is None):
-        raise ValueError("layer_key and adata_download_dir must be provided if text_only is False")
-    # If using a multiplets dataset, the column with the data representation is called "anchor". For the embedding workflow, only this data is embedded.
-    if "anchor" in data.column_names:
-        main_col = "anchor"
-    ###########################################################
-    # --- optional model-specific preparation ----------------
-    ###########################################################
-    impl = st_model[0] if len(st_model) > 0 else st_model
-    has_numeric = ("share_link" in data.column_names) or (
-        "adata_link" in data.column_names
-    )  # otherwise we cannot download the initial embeddings and register them with the model
-    path_map = None
-    if data is not None and has_numeric:
-        link_column = "share_link" if "share_link" in data.column_names else "adata_link"
-        # If the dataset is available, download get the token_df, even if the models doesnt use it. Just so the data is downloaded and the adata subset can be created downstream
-        logger.info("Extracting numeric intital embeddings from dataset via it's share_links …")
-        token_df, path_map = MMEnc.get_initial_embeddings_from_adata_link(
-            data,
-            layer_key=layer_key,
-            axis=axis,
-            download_dir=adata_download_dir,
-            overwrite=overwrite,
-            link_column=link_column,
-            zenodo_token=zenodo_token,
-        )  # type: ignore[arg-type]
-    else:
-        logger.info("""While the model supports initial embeddings, the dataset does not provide them.
-                    Therefore cell sentences will be analysed as texts.""")
-        token_df = None
-        # register only if the method exists
-    if hasattr(impl, "register_initial_embeddings") and token_df is not None and not text_only:
-        logger.info("Registering initial embeddings with MMContextEncoder module of SentenceTransformer Model …")
-        impl.register_initial_embeddings(token_df, data_origin=layer_key.split("X_")[1])  # type: ignore[arg-type]
-
-    if hasattr(impl, "prefix_ds") and not text_only:
-        logger.info("Calling prefix_ds on model-specific module …")
-        # Step 4: Apply prefixes using the new simplified prefix_ds method
-        logger.info(f"Applying prefixes to columns: {main_col}")
-        ds = impl.prefix_ds(
-            ds=data,
-            columns_to_prefix=main_col,
-        )  # type: ignore[arg-type]
-    else:
-        ds = data
-    ###########################################################
-    # --- sentence embeddings via encode() --------------------
-    ###########################################################
-    # -------- create a torch-compatible dataset ------------------------
-    if isinstance(ds, pd.DataFrame):
-        torch_ds = SentenceDataset(
-            indices=ds[index_col].tolist(),
-            texts=ds[main_col].astype(str).tolist(),
-        )
-    else:  # Hugging-Face Dataset
-        torch_ds = HFIndexedDataset(ds, main_col=main_col, index_col=index_col)
-
-    loader = DataLoader(
-        torch_ds,
-        batch_size=batch_size,
-        shuffle=False,
-        num_workers=num_workers,
-        pin_memory=pin_memory,
-    )
-
-    all_indices: list[int | str] = []
-    all_embeddings: list[list[float]] = []
-
-    st_model.eval()
-    with torch.inference_mode():
-        for batch_indices, batch_texts in tqdm(loader, desc="Encoding"):
-            # DataLoader returns tuples → convert to list[str]
-            embeddings = st_model.encode(
-                list(batch_texts),
-                batch_size=len(batch_texts),
-                convert_to_numpy=True,
-                show_progress_bar=False,
-            )
-            if isinstance(batch_indices, torch.Tensor):
-                batch_indices = batch_indices.tolist()
-            all_indices.extend(batch_indices)
-            all_embeddings.extend(embeddings.tolist())
-
-    # Create DataFrame using from_dict to avoid deprecation warning
-    emb_df = pd.DataFrame.from_dict({"sample_idx": all_indices, "embedding": all_embeddings})
-    logger.info("Generated %d embeddings", len(emb_df))
-    return emb_df, path_map
-
-
-def prepare_model_and_ds(
-    st_model: SentenceTransformer,
-    data: HFDataset | pd.DataFrame,
-    *,
-    layer_key: str = None,
-    adata_download_dir: str | Path = None,
-    main_col: str = "cell_sentence_1",
-    index_col: str = "sample_index",
-    axis: Literal["var", "obs"] = "obs",
-    text_only: bool = False,
-    overwrite: bool = True,
-) -> tuple[SentenceTransformer, HFDataset | pd.DataFrame]:
-    """
-    Prepare a Sentence-Transformer model and dataset for encoding.
-
-    This function handles model-specific preparation including:
-    - Downloading and registering initial numeric embeddings (if available)
-    - Applying dataset prefixes (if the model supports it)
-    - Storing path_map as an attribute on the model
-    - Returning the prepared model and dataset ready for encoding
-
-    Parameters
-    ----------
-    st_model : SentenceTransformer
-        The model wrapper. The *first* module (``st_model[0]``) may expose
-        extra methods such as ``prefix_ds`` and ``register_initial_embeddings``.
-        The ``path_map`` (if available) will be stored as ``st_model.path_map``.
-    data : HFDataset | pd.DataFrame
-        Dataset containing the data to encode. Must contain ``main_col`` and
-        ``index_col`` columns.
-    layer_key : str, optional
-        The name of the layer key to be used for the embedding (required if
-        ``text_only=False``).
-    adata_download_dir : str | Path, optional
-        Directory to download AnnData files to (required if ``text_only=False``).
-    main_col : str, default "cell_sentence_1"
-        Column containing the raw sentences or data to encode. This value is
-        preserved and not modified.
-    index_col : str, default "sample_index"
-        Column containing unique sample identifiers.
-    axis : Literal["var", "obs"], default "obs"
-        The axis to be used for the embedding. Can be either "obs" or "var".
-    text_only : bool, default False
-        If ``True``, all input will be treated as text. If ``False``, the dataset
-        must contain a column ``share_link`` or ``adata_link``, which points to
-        a zarr store that can be used to get initial embeddings for each token.
-    overwrite : bool, default True
-        If ``True``, the initial embeddings will be re-downloaded even if they
-        already exist.
-
-    Returns
-    -------
-    tuple[SentenceTransformer, HFDataset | pd.DataFrame]
-        A tuple containing:
-        - The prepared SentenceTransformer model (same object, may be modified in-place).
-          The ``path_map`` (if available) is stored as ``st_model.path_map``.
-        - The prepared dataset (with prefixes applied if applicable)
-
-    Notes
-    -----
-    * Optional methods are only executed if present on the model.
-    * The returned dataset can be used directly with ``model.encode(ds[main_col])``.
-    * The ``path_map`` (mapping from original links to actual file locations) is
-      stored as an attribute on the model: ``st_model.path_map``.
-
-    Examples
-    --------
-    >>> model, ds = prepare_model_and_ds(st_model, data, layer_key="X_pca", adata_download_dir="./cache")
-    >>> embeddings = model.encode(ds[main_col])
-    >>> # Access path_map if needed
-    >>> path_map = getattr(model, "path_map", None)
-    """
-    if not text_only and (layer_key is None or adata_download_dir is None):
-        raise ValueError("layer_key and adata_download_dir must be provided if text_only is False")
+"""Back-compat shim — the embed helpers now live in :mod:`mmcontext.embed.encode`.
 
-    ###########################################################
-    # --- optional model-specific preparation ----------------
-    ###########################################################
-    impl = st_model[0] if len(st_model) > 0 else st_model
-    has_numeric = (
-        isinstance(data, pd.DataFrame) and ("share_link" in data.columns or "adata_link" in data.columns)
-    ) or (
-        not isinstance(data, pd.DataFrame) and ("share_link" in data.column_names or "adata_link" in data.column_names)
-    )  # otherwise we cannot download the initial embeddings and register them with the model
+This module re-exports the public embed helpers under their historical
+``mmcontext.embed.model_utils`` import path so existing consumers (notably the
+``sc_language_foundation_models`` ``predict.py`` scripts) keep working without
+changes. New code should import from :mod:`mmcontext.embed.encode`.
+"""
 
-    path_map = None
-    if data is not None and has_numeric and not text_only:
-        if isinstance(data, pd.DataFrame):
-            link_column = "share_link" if "share_link" in data.columns else "adata_link"
-        else:
-            link_column = "share_link" if "share_link" in data.column_names else "adata_link"
-        # If the dataset is available, download get the token_df, even if the models doesnt use it.
-        # Just so the data is downloaded and the adata subset can be created downstream
-        logger.info("Extracting numeric initial embeddings from dataset via it's share_links …")
-        token_df, path_map = MMEnc.get_initial_embeddings_from_adata_link(
-            data,
-            layer_key=layer_key,
-            axis=axis,
-            download_dir=adata_download_dir,
-            overwrite=overwrite,
-            link_column=link_column,
-        )  # type: ignore[arg-type]
-    else:
-        logger.info("""All input will be treated as text.""")
-        token_df = None
-
-    # Store path_map as an attribute on the model
-    st_model.path_map = path_map  # type: ignore[attr-defined]
-
-    # register only if the method exists
-    if hasattr(impl, "register_initial_embeddings") and token_df is not None and not text_only:
-        logger.info("Registering initial embeddings with MMContextEncoder module of SentenceTransformer Model …")
-        impl.register_initial_embeddings(token_df, data_origin=layer_key.split("X_")[1])  # type: ignore[arg-type]
-
-    if hasattr(impl, "prefix_ds") and not text_only:
-        logger.info("Calling prefix_ds on model-specific module …")
-        logger.info(f"Applying prefixes to columns: {main_col}")
-        ds = impl.prefix_ds(
-            ds=data,
-            columns_to_prefix=main_col,
-        )  # type: ignore[arg-type]
-    else:
-        ds = data
-
-    return st_model, ds
-
-
-def encode_adata(
-    st_model: SentenceTransformer,
-    data: HFDataset | pd.DataFrame,
-    adata: ad.AnnData,
-    *,
-    layer_key: str = None,
-    adata_download_dir: str | Path = None,
-    main_col: str = "cell_sentence_1",
-    index_col: str = "sample_index",
-    axis: Literal["var", "obs"] = "obs",
-    text_only: bool = False,
-    overwrite: bool = True,
-    batch_size: int = 32,
-    show_progress_bar: bool = True,
-    obsm_key: str = "mmcontext_emb",
-) -> ad.AnnData:
-    """
-    Encode dataset using a Sentence-Transformer model and store results in AnnData object.
-
-    This function prepares the model and dataset, generates embeddings, validates
-    that indices match between the dataset and AnnData object, and stores the
-    embeddings in ``adata.obsm[obsm_key]`` with rows sorted to match
-    ``adata.obs.index``.
-
-    Parameters
-    ----------
-    st_model : SentenceTransformer
-        The model wrapper. The *first* module (``st_model[0]``) may expose
-        extra methods such as ``prefix_ds`` and ``register_initial_embeddings``.
-    data : HFDataset | pd.DataFrame
-        Dataset containing the data to encode. Must contain ``main_col`` and
-        ``index_col`` columns. The indices in ``index_col`` must match
-        ``adata.obs.index``.
-    adata : anndata.AnnData
-        AnnData object where embeddings will be stored. The ``adata.obs.index``
-        must match the indices in the dataset's ``index_col``.
-    layer_key : str, optional
-        The name of the layer key to be used for the embedding (required if
-        ``text_only=False``).
-    adata_download_dir : str | Path, optional
-        Directory to download AnnData files to (required if ``text_only=False``).
-    main_col : str, default "cell_sentence_1"
-        Column containing the raw sentences or data to encode.
-    index_col : str, default "sample_index"
-        Column containing unique sample identifiers that match ``adata.obs.index``.
-    axis : Literal["var", "obs"], default "obs"
-        The axis to be used for the embedding. Can be either "obs" or "var".
-    text_only : bool, default False
-        If ``True``, all input will be treated as text. If ``False``, the dataset
-        must contain a column ``share_link`` or ``adata_link``, which points to
-        a zarr store that can be used to get initial embeddings for each token.
-    overwrite : bool, default True
-        If ``True``, the initial embeddings will be re-downloaded even if they
-        already exist.
-    batch_size : int, default 32
-        Batch size for encoding. Passed to ``model.encode()``.
-    show_progress_bar : bool, default True
-        Whether to show a progress bar during encoding.
-    obsm_key : str, default "mmcontext_emb"
-        Key to store embeddings in ``adata.obsm``.
-
-    Returns
-    -------
-    anndata.AnnData
-        The input AnnData object with embeddings stored in ``adata.obsm[obsm_key]``.
-        The embeddings are sorted to match the order of ``adata.obs.index``.
-
-    Raises
-    ------
-    ValueError
-        If indices in the dataset do not match ``adata.obs.index``, or if there
-        are missing indices in either direction.
-
-    Notes
-    -----
-    * This function validates that all indices in the dataset are present in
-      ``adata.obs.index`` and vice versa. If there's a mismatch, an error is
-      raised suggesting to verify that the same AnnData object was used to build
-      the dataset as was passed to this function.
-    * The embeddings are automatically sorted to match the order of
-      ``adata.obs.index`` before being stored.
-
-    Examples
-    --------
-    >>> adata_with_emb = encode_adata(st_model, dataset, adata, layer_key="X_pca", adata_download_dir="./cache")
-    >>> # Embeddings are now available in adata_with_emb.obsm["mmcontext_emb"]
-    """
-    # Prepare model and dataset
-    prepared_model, prepared_ds = prepare_model_and_ds(
-        st_model=st_model,
-        data=data,
-        layer_key=layer_key,
-        adata_download_dir=adata_download_dir,
-        main_col=main_col,
-        index_col=index_col,
-        axis=axis,
-        text_only=text_only,
-        overwrite=overwrite,
-    )
-
-    # Extract indices from the prepared dataset
-    if isinstance(prepared_ds, pd.DataFrame):
-        dataset_indices = prepared_ds[index_col].tolist()
-    else:  # Hugging-Face Dataset
-        dataset_indices = prepared_ds[index_col]  # type: ignore[assignment]
-        if not isinstance(dataset_indices, list):
-            dataset_indices = list(dataset_indices)
-
-    # Validate indices match
-    adata_indices = list(adata.obs.index)
-    dataset_indices_set = set(dataset_indices)
-    adata_indices_set = set(adata_indices)
-
-    missing_in_adata = dataset_indices_set - adata_indices_set
-    missing_in_dataset = adata_indices_set - dataset_indices_set
-
-    if missing_in_adata or missing_in_dataset:
-        error_msg = "Index mismatch between dataset and AnnData object.\n"
-        if missing_in_adata:
-            error_msg += f"  - {len(missing_in_adata)} indices found in dataset but not in adata.obs.index: {list(missing_in_adata)[:10]}{'...' if len(missing_in_adata) > 10 else ''}\n"
-        if missing_in_dataset:
-            error_msg += f"  - {len(missing_in_dataset)} indices found in adata.obs.index but not in dataset: {list(missing_in_dataset)[:10]}{'...' if len(missing_in_dataset) > 10 else ''}\n"
-        error_msg += "\nPlease double-check that you used the same AnnData object to build the dataset as you passed to this method."
-        raise ValueError(error_msg)
-
-    # Extract texts for encoding
-    if isinstance(prepared_ds, pd.DataFrame):
-        texts = prepared_ds[main_col].astype(str).tolist()
-    else:  # Hugging-Face Dataset
-        texts = prepared_ds[main_col]  # type: ignore[assignment]
-        if not isinstance(texts, list):
-            texts = list(texts)
-        texts = [str(t) for t in texts]
-
-    # Encode
-    logger.info(f"Encoding {len(texts)} samples...")
-    embeddings = prepared_model.encode(
-        texts,
-        batch_size=batch_size,
-        convert_to_numpy=True,
-        show_progress_bar=show_progress_bar,
-    )
-
-    # Create DataFrame with embeddings
-    emb_df = pd.DataFrame.from_dict(
-        {
-            "sample_idx": dataset_indices,
-            "embedding": embeddings.tolist(),
-        }
-    )
-
-    # Validate that embedding indices match dataset indices (sanity check)
-    if len(emb_df) != len(dataset_indices):
-        raise ValueError(
-            f"Number of embeddings ({len(emb_df)}) does not match number of dataset indices ({len(dataset_indices)}). "
-            "This should not happen - please report this as a bug."
-        )
-
-    # Validate that all embedding indices are in adata.obs.index before reindexing
-    emb_indices_set = set(emb_df["sample_idx"])
-    adata_indices_set = set(adata.obs.index)
-
-    missing_emb_in_adata = emb_indices_set - adata_indices_set
-    if missing_emb_in_adata:
-        raise ValueError(
-            f"Found {len(missing_emb_in_adata)} embedding indices that are not in adata.obs.index: "
-            f"{list(missing_emb_in_adata)[:10]}{'...' if len(missing_emb_in_adata) > 10 else ''}. "
-            "Please double-check that you used the same AnnData object to build the dataset as you passed to this method."
-        )
-
-    missing_adata_in_emb = adata_indices_set - emb_indices_set
-    if missing_adata_in_emb:
-        raise ValueError(
-            f"Found {len(missing_adata_in_emb)} indices in adata.obs.index that do not have embeddings: "
-            f"{list(missing_adata_in_emb)[:10]}{'...' if len(missing_adata_in_emb) > 10 else ''}. "
-            "Please double-check that you used the same AnnData object to build the dataset as you passed to this method."
-        )
-
-    # Sort emb_df so that the sample_idx (index column) matches the order of adata.obs.index
-    # Convert adata.obs.index to a list to avoid pandas interpreting it as a single key
-    emb_df_sorted = emb_df.set_index("sample_idx").reindex(list(adata.obs.index)).reset_index()
-
-    # Convert to numpy array
-    embedding_matrix = np.vstack(emb_df_sorted["embedding"].to_numpy())
-
-    # Store in adata
-    adata.obsm[obsm_key] = embedding_matrix
-    logger.info(f"Stored embeddings in adata.obsm['{obsm_key}'] with shape {embedding_matrix.shape}")
-
-    return adata
-
-
-def create_label_dataset(adata, label_col: str) -> tuple[HFDataset, dict[str, int]]:
-    """
-    Create a small HF Dataset from unique labels in an AnnData object.
-
-    Parameters
-    ----------
-    adata : anndata.AnnData
-        AnnData object containing the labels
-    label_col : str
-        Name of the column in adata.obs containing the labels
-
-    Returns
-    -------
-    tuple[HFDataset, dict[str, int]]
-        A tuple containing:
-        - HuggingFace dataset with unique labels and their indices
-        - Mapping from label string to embedding index
-    """
-    from datasets import Dataset
-
-    # Get unique labels only
-    unique_labels = adata.obs[label_col].astype(str).unique().tolist()
-
-    # Create mapping from label to embedding index
-    label_to_index = {label: idx for idx, label in enumerate(unique_labels)}
-
-    # Create indices for the unique labels (just sequential)
-    indices = list(range(len(unique_labels)))
-
-    # Create dataset with unique labels only
-    ds = Dataset.from_dict({"sample_idx": indices, "label": unique_labels})
-
-    return ds, label_to_index
-
-
-def embed_labels(
-    st_model: SentenceTransformer,
-    adata: ad.AnnData,
-    label_col: str,
-    *,
-    batch_size: int = 32,
-    num_workers: int = 0,
-    pin_memory: bool = torch.cuda.is_available(),
-) -> tuple[pd.DataFrame, dict[str, int]]:
-    """
-    Embed unique labels from an AnnData object using a SentenceTransformer model.
-
-    Parameters
-    ----------
-    st_model : SentenceTransformer
-        The model to use for embedding
-    adata : anndata.AnnData
-        AnnData object containing the labels
-    label_col : str
-        Name of the column in adata.obs containing the labels
-    batch_size : int, default 32
-        Number of labels per encode() call
-    num_workers : int, default 0
-        Worker processes for the PyTorch data loader
-    pin_memory : bool, default True if CUDA is available
-        Forwarded to DataLoader for speed on GPUs
-
-    Returns
-    -------
-    tuple[pd.DataFrame, dict[str, int]]
-        A tuple containing:
-        - DataFrame with columns ['sample_idx', 'embedding'] containing unique label embeddings
-        - Mapping from label string to embedding index in the DataFrame
-    """
-    # Create dataset from unique labels
-    label_ds, label_to_index = create_label_dataset(adata, label_col)
-
-    # Use prepare_model_and_embed with text_only=True since we're just embedding text labels
-    emb_df, _ = prepare_model_and_embed(
-        st_model,
-        data=label_ds,
-        main_col="label",
-        index_col="sample_idx",
-        batch_size=batch_size,
-        num_workers=num_workers,
-        pin_memory=pin_memory,
-        text_only=True,
-    )
+from __future__ import annotations
 
-    return emb_df, label_to_index
+from mmcontext.embed.encode import (
+    HFIndexedDataset,
+    SentenceDataset,
+    create_label_dataset,
+    embed_labels,
+    load_st_model,
+    prepare_model_and_embed,
+)
+
+__all__ = [
+    "HFIndexedDataset",
+    "SentenceDataset",
+    "create_label_dataset",
+    "embed_labels",
+    "load_st_model",
+    "prepare_model_and_embed",
+]
diff --git a/src/mmcontext/embed/pipeline.py b/src/mmcontext/embed/pipeline.py
new file mode 100644
index 0000000..9343ddb
--- /dev/null
+++ b/src/mmcontext/embed/pipeline.py
@@ -0,0 +1,77 @@
+"""Builder for the mmcontext sentence-transformers pipeline.
+
+The standard obs-level pipeline is::
+
+    [MMContextModule, AdapterModule, Pooling, Normalize]
+
+:func:`build_pipeline` assembles it with a single call so training scripts
+(``scripts/train_tiny.py``, ``scripts/train_config.py``) share one definition.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.sentence_transformer.modules import Normalize, Pooling
+
+from mmcontext.modules import AdapterModule, MMContextModule
+
+logger = logging.getLogger(__name__)
+
+
+def build_pipeline(
+    text_model: str,
+    *,
+    omics_dim: int | None = None,
+    shared_dim: int = 256,
+    adapter_hidden_dim: int | None = 512,
+    max_seq_length: int = 512,
+) -> SentenceTransformer:
+    """Build the MMContext sentence-transformer pipeline.
+
+    Parameters
+    ----------
+    text_model : str
+        HuggingFace model name/path for the text encoder.
+    omics_dim : int, optional
+        Dimension of omics vectors. If ``None``, the adapter's omics head is
+        sized to match the text dimension (text-only / placeholder use).
+    shared_dim : int, default 256
+        Output dimension of the shared embedding space.
+    adapter_hidden_dim : int or None, default 512
+        Hidden layer size of each adapter projection. ``None``/``0`` uses a
+        single ``Linear -> LayerNorm`` head instead of an MLP.
+    max_seq_length : int, default 512
+        Maximum sequence length for text tokenization.
+
+    Returns
+    -------
+    sentence_transformers.SentenceTransformer
+        The assembled pipeline.
+    """
+    mmcontext = MMContextModule(model_name_or_path=text_model, max_seq_length=max_seq_length)
+    text_dim = mmcontext.get_word_embedding_dimension()
+
+    if omics_dim is None:
+        omics_dim = text_dim
+
+    adapter = AdapterModule(
+        text_input_dim=text_dim,
+        omics_input_dim=omics_dim,
+        shared_dim=shared_dim,
+        hidden_dim=adapter_hidden_dim,
+    )
+    pooling = Pooling(embedding_dimension=shared_dim, pooling_mode="mean")
+    normalize = Normalize()
+
+    pipeline = SentenceTransformer(modules=[mmcontext, adapter, pooling, normalize])
+    logger.info(
+        "Pipeline: text_dim=%d, omics_dim=%d, shared_dim=%d, adapter_hidden_dim=%s, params=%d",
+        text_dim,
+        omics_dim,
+        shared_dim,
+        adapter_hidden_dim,
+        sum(p.numel() for p in pipeline.parameters()),
+    )
+    return pipeline
diff --git a/src/mmcontext/embed/scsa_utils.py b/src/mmcontext/embed/scsa_utils.py
deleted file mode 100644
index 70a6000..0000000
--- a/src/mmcontext/embed/scsa_utils.py
+++ /dev/null
@@ -1,999 +0,0 @@
-"""
-SCSA (Single-Cell RNA-seq Annotation) utilities for integration with mmcontext pipeline.
-
-SCSA is a marker-gene-based cell type annotation tool that works at the cluster level.
-Unlike embedding models, it produces direct predictions by matching differentially
-expressed genes against curated cell marker databases.
-
-This module provides functions to run the full SCSA pipeline and compute evaluation
-metrics compatible with the mmcontext metrics collection system.
-"""
-
-import logging
-import os
-import subprocess
-import sys
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any
-
-import anndata as ad
-import numpy as np
-import pandas as pd
-import scanpy as sc
-
-from mmcontext.eval.utils import LabelKind
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class HarmonizeResult:
-    """Outcome of the calmate label-harmonisation step."""
-
-    success: bool
-    all_reviewed: bool
-    mapped_csv: Path | None = None
-    messages: list[str] = field(default_factory=list)
-
-
-def _reorder_harmonized_columns(df: pd.DataFrame) -> pd.DataFrame:
-    """Reorder columns for the harmonised predictions CSV.
-
-    Desired order (when present):
-    ``cell_id``, ``true_labels``, ``true_labels_mapped``,
-    ``predicted_labels``, ``predicted_labels_mapped``, then any remaining
-    columns in their original order.
-    """
-    preferred: list[str] = []
-    if "cell_id" in df.columns:
-        preferred.append("cell_id")
-    for col in [
-        "true_labels",
-        "true_labels_mapped",
-        "predicted_labels",
-        "predicted_labels_mapped",
-    ]:
-        if col in df.columns and col not in preferred:
-            preferred.append(col)
-    remaining = [c for c in df.columns if c not in preferred]
-    return df[preferred + remaining]
-
-
-def ensure_scsa_setup(modules_dir: Path) -> bool:
-    """
-    Ensure SCSA is properly set up by checking for the repository and database.
-
-    Runs prepare_scsa.sh if the SCSA directory is missing.
-
-    Parameters
-    ----------
-    modules_dir : Path
-        Directory containing (or to contain) the SCSA installation.
-
-    Returns
-    -------
-    bool
-        True if setup is complete, False otherwise.
-    """
-    scsa_dir = modules_dir / "SCSA"
-    scsa_db = scsa_dir / "whole_v2.db"
-
-    if scsa_dir.exists() and scsa_db.exists():
-        logger.info("SCSA already set up at %s", scsa_dir)
-        return True
-
-    prepare_script = modules_dir / "prepare_scsa.sh"
-    if not prepare_script.exists():
-        logger.error("SCSA prepare script not found at %s", prepare_script)
-        return False
-
-    try:
-        logger.info("Running SCSA prepare script...")
-        result = subprocess.run(
-            ["bash", str(prepare_script)],
-            cwd=str(modules_dir),
-            capture_output=True,
-            text=True,
-            timeout=300,
-        )
-
-        if result.returncode != 0:
-            logger.error("SCSA prepare script failed (rc=%d): %s", result.returncode, result.stderr)
-            return False
-
-        logger.info("SCSA prepare script completed successfully")
-        if result.stdout:
-            logger.info("Prepare output: %s", result.stdout.strip())
-
-        return scsa_dir.exists() and scsa_db.exists()
-
-    except subprocess.TimeoutExpired:
-        logger.error("SCSA prepare script timed out")
-        return False
-    except Exception as e:
-        logger.error("Error running SCSA prepare script: %s", e)
-        return False
-
-
-def create_marker_csv(
-    adata: ad.AnnData,
-    output_csv: Path,
-    cluster_col: str = "louvain",
-    resolution: float = 1.0,
-    n_top_genes: int = 2000,
-    n_neighbors: int = 15,
-    n_pcs: int = 50,
-    random_state: int = 0,
-) -> tuple[ad.AnnData, str]:
-    """
-    Cluster cells (if needed) and create a marker-gene CSV in SCSA's expected format.
-
-    The marker detection is performed on normalized (not scaled) data so that
-    log fold changes are valid for SCSA's filtering.
-
-    Parameters
-    ----------
-    adata : AnnData
-        Input AnnData object with expression data.
-    output_csv : Path
-        Path to write the marker CSV.
-    cluster_col : str
-        Name of the cluster column to use or create.
-    resolution : float
-        Resolution for Louvain clustering if clustering is needed.
-    n_top_genes : int
-        Number of highly variable genes for preprocessing.
-    n_neighbors : int
-        Number of neighbors for the neighborhood graph.
-    n_pcs : int
-        Number of principal components.
-    random_state : int
-        Random seed for reproducibility.
-
-    Returns
-    -------
-    tuple[AnnData, str]
-        The (possibly modified) AnnData and the cluster column name used.
-    """
-    needs_clustering = cluster_col not in adata.obs.columns
-
-    if "gene_names" in adata.var.columns:
-        adata.var.index = adata.var["gene_names"]
-
-    if needs_clustering:
-        logger.info("No cluster column '%s' found — running preprocessing + Louvain clustering", cluster_col)
-
-        sc.pp.highly_variable_genes(adata, n_top_genes=n_top_genes, flavor="seurat_v3")
-        adata = adata[:, adata.var.highly_variable].copy()
-
-        adata_for_markers = adata.copy()
-
-        sc.pp.scale(adata, max_value=10)
-        sc.tl.pca(adata, svd_solver="arpack", n_comps=n_pcs)
-        sc.pp.neighbors(adata, n_neighbors=n_neighbors, n_pcs=n_pcs, random_state=random_state)
-        sc.tl.louvain(adata, resolution=resolution, key_added=cluster_col, random_state=random_state)
-
-        adata_for_markers.obs[cluster_col] = adata.obs[cluster_col]
-    else:
-        logger.info("Using existing cluster column '%s'", cluster_col)
-        adata_for_markers = adata.copy()
-        if "log1p" in adata.layers:
-            adata_for_markers.X = adata.layers["log1p"].copy()
-        elif adata.raw is not None:
-            adata_for_markers = adata.raw.to_adata()
-            adata_for_markers.obs[cluster_col] = adata.obs[cluster_col]
-
-    n_clusters = len(adata.obs[cluster_col].unique())
-    logger.info("Found %d clusters", n_clusters)
-
-    logger.info("Computing marker genes per cluster (on normalized data)...")
-    sc.tl.rank_genes_groups(adata_for_markers, groupby=cluster_col, method="t-test")
-
-    result = adata_for_markers.uns["rank_genes_groups"]
-    groups = result["names"].dtype.names
-
-    dat = pd.DataFrame(
-        {
-            group + "_" + key[:1]: result[key][group]
-            for group in groups
-            for key in ["names", "logfoldchanges", "scores", "pvals"]
-        }
-    )
-
-    dat.to_csv(output_csv, index=True)
-    logger.info("Wrote marker CSV with %d genes and %d columns to %s", len(dat), len(dat.columns), output_csv)
-
-    return adata, cluster_col
-
-
-def run_scsa(
-    csv_path: Path,
-    db_path: Path,
-    output_path: Path,
-    scsa_dir: Path,
-    *,
-    scsa_python: str | Path | None = None,
-    foldchange: float = 1.5,
-    pvalue: float = 0.01,
-    species: str = "Human",
-    tissue: str = "All",
-    gensymbol: bool = True,
-) -> Path:
-    """
-    Run SCSA.py on a marker-gene CSV.
-
-    SCSA requires numpy<2.0 and must be run with a dedicated Python
-    interpreter from a separate virtual environment.
-
-    Parameters
-    ----------
-    csv_path : Path
-        Path to marker CSV produced by :func:`create_marker_csv`.
-    db_path : Path
-        Path to the SCSA database file.
-    output_path : Path
-        Path to write SCSA results.
-    scsa_dir : Path
-        Directory containing SCSA.py.
-    scsa_python : str or Path or None
-        Python executable from the SCSA venv.  When *None* the function
-        looks for ``modules/scsa_venv/bin/python`` relative to *scsa_dir*'s
-        parent, and falls back to ``sys.executable`` with a warning.
-    foldchange : float
-        Minimum log fold-change threshold.
-    pvalue : float
-        Maximum p-value threshold.
-    species : str
-        Species for database lookup.
-    tissue : str
-        Tissue type for database lookup.
-    gensymbol : bool
-        If True, use gene symbols instead of Ensembl IDs.
-
-    Returns
-    -------
-    Path
-        Path to the SCSA output file.
-    """
-    if scsa_python is None:
-        default_venv_python = scsa_dir.parent / "scsa_venv" / "bin" / "python"
-        if default_venv_python.exists():
-            scsa_python = str(default_venv_python)
-        else:
-            logger.warning(
-                "SCSA venv not found at %s — falling back to current interpreter. "
-                "This will fail if numpy>=2.0 is installed. "
-                "Create the venv with: python -m venv modules/scsa_venv && "
-                "modules/scsa_venv/bin/pip install -r modules/requirements_scsa.txt",
-                default_venv_python,
-            )
-            scsa_python = sys.executable
-    else:
-        scsa_python = str(scsa_python)
-
-    cmd = [
-        scsa_python,
-        str(scsa_dir / "SCSA.py"),
-        "-d",
-        str(db_path),
-        "-i",
-        str(csv_path),
-        "-s",
-        "scanpy",
-        "-o",
-        str(output_path),
-        "-f",
-        str(foldchange),
-        "-p",
-        str(pvalue),
-        "-g",
-        species,
-        "-k",
-        tissue,
-        "-T",
-        "normal",
-        "-t",
-        "cellmarker",
-        "-m",
-        "txt",
-    ]
-
-    if gensymbol:
-        cmd.append("-E")
-
-    logger.info("Running SCSA: %s", " ".join(cmd))
-    result = subprocess.run(cmd, capture_output=True, text=True, cwd=str(scsa_dir))
-
-    if result.returncode != 0:
-        logger.error("SCSA failed (rc=%d)\nstdout: %s\nstderr: %s", result.returncode, result.stdout, result.stderr)
-        raise RuntimeError(f"SCSA failed with return code {result.returncode}: {result.stderr}")
-
-    if not output_path.exists() or output_path.stat().st_size == 0:
-        logger.warning("SCSA output is empty or missing at %s", output_path)
-        if result.stdout:
-            logger.info("SCSA stdout: %s", result.stdout[:1000])
-
-    logger.info("SCSA completed, output at %s", output_path)
-    return output_path
-
-
-def map_predictions_to_cells(
-    scsa_result_file: Path,
-    adata: ad.AnnData,
-    cluster_col: str,
-    annotation_key: str,
-) -> pd.DataFrame:
-    """
-    Map SCSA cluster-level predictions to individual cells.
-
-    For each cluster, the cell type with the highest Z-score is selected.
-    Clusters without a prediction are labelled "Unknown".
-
-    Parameters
-    ----------
-    scsa_result_file : Path
-        Path to SCSA tab-separated output file.
-    adata : AnnData
-        AnnData with cluster assignments and true labels.
-    cluster_col : str
-        Column in ``adata.obs`` with cluster IDs.
-    annotation_key : str
-        Column in ``adata.obs`` with true cell-type labels.
-
-    Returns
-    -------
-    DataFrame
-        Columns: ``cell_id``, ``true_labels``, ``predicted_labels``.
-    """
-    if not scsa_result_file.exists() or scsa_result_file.stat().st_size == 0:
-        logger.warning("SCSA result file empty or missing: %s — all cells will be 'Unknown'", scsa_result_file)
-        return pd.DataFrame(
-            {
-                "cell_id": adata.obs.index,
-                "true_labels": adata.obs[annotation_key].values,
-                "predicted_labels": "Unknown",
-            }
-        )
-
-    scsa_results = pd.read_csv(scsa_result_file, sep="\t")
-    expected_cols = {"Cell Type", "Z-score", "Cluster"}
-    if not expected_cols.issubset(scsa_results.columns):
-        raise ValueError(
-            f"SCSA output missing expected columns. Found {list(scsa_results.columns)}, need {expected_cols}"
-        )
-
-    # Pick the top prediction (highest Z-score) per cluster
-    cluster_predictions: dict[str, str] = {}
-    for cluster in scsa_results["Cluster"].unique():
-        cluster_data = scsa_results[scsa_results["Cluster"] == cluster].copy()
-        cluster_data["Z-score"] = pd.to_numeric(cluster_data["Z-score"], errors="coerce")
-        cluster_data = cluster_data.dropna(subset=["Z-score"])
-        if len(cluster_data) > 0:
-            best = cluster_data.loc[cluster_data["Z-score"].idxmax()]
-            cluster_predictions[str(cluster)] = best["Cell Type"]
-
-    logger.info("SCSA predictions for %d clusters", len(cluster_predictions))
-
-    if annotation_key not in adata.obs.columns:
-        raise ValueError(f"Annotation key '{annotation_key}' not found. Available: {list(adata.obs.columns)}")
-
-    cell_data = []
-    unmapped = set()
-    for idx, cluster in enumerate(adata.obs[cluster_col]):
-        cstr = str(cluster)
-        pred = cluster_predictions.get(cstr)
-        if pred is None:
-            unmapped.add(cstr)
-            pred = "Unknown"
-        cell_data.append(
-            {
-                "cell_id": adata.obs.index[idx],
-                "true_labels": adata.obs[annotation_key].iloc[idx],
-                "predicted_labels": pred,
-            }
-        )
-
-    if unmapped:
-        logger.warning("%d clusters had no SCSA prediction: %s", len(unmapped), sorted(unmapped))
-
-    return pd.DataFrame(cell_data)
-
-
-def compute_metrics(
-    true_labels: np.ndarray | pd.Series,
-    predicted_labels: np.ndarray | pd.Series,
-    dataset_name: str,
-    annotation_key: str,
-) -> list[dict[str, Any]]:
-    """
-    Compute evaluation metrics directly from predictions.
-
-    Produces rows in the same schema as the LabelSimilarity evaluator so that
-    ``collect_metrics.py`` can fuse them with other model results.
-
-    Parameters
-    ----------
-    true_labels : array-like
-        Ground truth labels.
-    predicted_labels : array-like
-        Predicted labels.
-    dataset_name : str
-        Name of the dataset.
-    annotation_key : str
-        Name of the label column.
-
-    Returns
-    -------
-    list[dict]
-        One dict per metric, each with keys:
-        ``dataset``, ``model``, ``model_name``, ``label``, ``label_kind``,
-        ``metric``, ``value``.
-    """
-    true_labels = np.asarray(true_labels)
-    predicted_labels = np.asarray(predicted_labels)
-
-    unique_labels = np.unique(true_labels)
-    n_labels = len(unique_labels)
-
-    # Accuracy
-    accuracy = float(np.mean(true_labels == predicted_labels))
-
-    # Balanced accuracy (mean per-class recall)
-    per_class_recall = []
-    for label in unique_labels:
-        mask = true_labels == label
-        if mask.sum() > 0:
-            per_class_recall.append(float(np.mean(predicted_labels[mask] == label)))
-    balanced_accuracy = float(np.mean(per_class_recall)) if per_class_recall else 0.0
-
-    # Macro F1
-    f1_scores = []
-    for label in unique_labels:
-        mask_true = true_labels == label
-        mask_pred = predicted_labels == label
-        tp = np.sum(mask_true & mask_pred)
-        fp = np.sum(~mask_true & mask_pred)
-        fn = np.sum(mask_true & ~mask_pred)
-        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
-        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
-        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
-        f1_scores.append(f1)
-    macro_f1 = float(np.mean(f1_scores)) if f1_scores else 0.0
-
-    # Random baseline
-    random_baseline = 1.0 / n_labels if n_labels > 0 else 0.0
-    accuracy_over_random = accuracy / random_baseline if random_baseline > 0 else 0.0
-
-    metrics = {
-        "LabelSimilarity/accuracy": accuracy,
-        "LabelSimilarity/balanced_accuracy": balanced_accuracy,
-        "LabelSimilarity/macro_f1": macro_f1,
-        "LabelSimilarity/random_baseline_accuracy": random_baseline,
-        "LabelSimilarity/accuracy_over_random": accuracy_over_random,
-        "LabelSimilarity/n_labels": float(n_labels),
-        # Embedding-only metrics — not computable for SCSA
-        "LabelSimilarity/mean_auc": float("nan"),
-        "LabelSimilarity/std_auc": float("nan"),
-        "LabelSimilarity/topk_accuracy@1": float("nan"),
-        "LabelSimilarity/topk_accuracy@3": float("nan"),
-        "LabelSimilarity/topk_accuracy@5": float("nan"),
-        "LabelSimilarity/mrr": float("nan"),
-    }
-
-    rows = []
-    for metric_name, value in metrics.items():
-        rows.append(
-            {
-                "dataset": dataset_name,
-                "model": "scsa",
-                "model_name": "scsa",
-                "label": annotation_key,
-                "label_kind": LabelKind.BIO,
-                "metric": metric_name,
-                "value": value,
-            }
-        )
-
-    return rows
-
-
-def _find_calmate_cli(calmate_python: str | Path | None, modules_dir: Path) -> str | None:
-    """Locate the ``calmate`` console-script in a venv.
-
-    Resolution order:
-    1. Derive from *calmate_python* (``…/bin/python`` → ``…/bin/calmate``)
-    2. Auto-detect ``modules/calmate_venv/bin/calmate``
-    """
-    if calmate_python is not None:
-        candidate = Path(calmate_python).parent / "calmate"
-        if candidate.exists():
-            return str(candidate)
-
-    default = modules_dir / "calmate_venv" / "bin" / "calmate"
-    if default.exists():
-        return str(default)
-
-    return None
-
-
-def _run_calmate(calmate_cli: str, args: list[str], label: str) -> tuple[bool, str]:
-    """Run a single calmate CLI command and return (ok, combined_output).
-
-    Both stdout and stderr are always logged so that diagnostic messages
-    from calmate (e.g. backend availability warnings) are visible.
-    """
-    cmd = [calmate_cli, *args]
-    logger.info("Running calmate (%s): %s", label, " ".join(cmd))
-    result = subprocess.run(cmd, capture_output=True, text=True)
-    output = (result.stdout or "") + (result.stderr or "")
-    if output.strip():
-        for line in output.strip().splitlines():
-            logger.info("calmate [%s]: %s", label, line)
-    if result.returncode != 0:
-        logger.error("calmate %s failed (rc=%d)", label, result.returncode)
-        return False, output.strip()
-    return True, output.strip()
-
-
-def harmonize_labels(
-    predictions_csv: Path,
-    store_path: Path,
-    modules_dir: Path,
-    *,
-    dataset_name: str,
-    calmate_python: str | Path | None = None,
-    cache_dir: str | None = None,
-) -> HarmonizeResult:
-    """Map SCSA predicted and ground-truth labels to Cell Ontology terms via calmate.
-
-    Calls the calmate CLI twice for ``map`` (predicted + true labels) then twice
-    for ``apply`` (predicted + true labels), yielding a final CSV with
-    ``predicted_labels_mapped`` and ``true_labels_mapped`` columns.
-
-    Parameters
-    ----------
-    predictions_csv:
-        CSV produced by :func:`map_predictions_to_cells` with at least
-        ``predicted_labels`` and ``true_labels`` columns.
-    store_path:
-        Path to the shared calmate mapping-store CSV.
-    modules_dir:
-        ``modules/`` directory, used to auto-detect the calmate venv.
-    calmate_python:
-        Explicit path to the Python interpreter inside the calmate venv.
-        When *None*, ``modules/calmate_venv/bin/calmate`` is tried.
-    cache_dir:
-        Ontology / model cache directory forwarded to ``calmate map``.
-
-    Returns
-    -------
-    HarmonizeResult
-    """
-    calmate_cli = _find_calmate_cli(calmate_python, modules_dir)
-    if calmate_cli is None:
-        return HarmonizeResult(
-            success=False,
-            all_reviewed=False,
-            messages=[
-                "Calmate venv not found. Create it with:",
-                "  python -m venv modules/calmate_venv",
-                "  modules/calmate_venv/bin/pip install -r modules/requirements_calmate.txt",
-            ],
-        )
-
-    store_path.parent.mkdir(parents=True, exist_ok=True)
-    store_arg = ["--store", str(store_path)]
-    diagnostics: list[str] = []
-
-    # --- 1. Map predicted labels ---
-    map_pred_args = [
-        *store_arg,
-        "map",
-        str(predictions_csv),
-        "-c",
-        "predicted_labels",
-        "-o",
-        f"scsa:{dataset_name}",
-    ]
-    if cache_dir:
-        map_pred_args += ["--cache-dir", cache_dir]
-    ok, out = _run_calmate(calmate_cli, map_pred_args, "map predicted_labels")
-    if out:
-        diagnostics.append(f"[map predicted_labels] {out}")
-    if not ok:
-        return HarmonizeResult(
-            success=False, all_reviewed=False, messages=[f"calmate map predicted_labels failed: {out}"]
-        )
-
-    # --- 2. Map true labels ---
-    map_true_args = [
-        *store_arg,
-        "map",
-        str(predictions_csv),
-        "-c",
-        "true_labels",
-        "-o",
-        f"ground_truth:{dataset_name}",
-    ]
-    if cache_dir:
-        map_true_args += ["--cache-dir", cache_dir]
-    ok, out = _run_calmate(calmate_cli, map_true_args, "map true_labels")
-    if out:
-        diagnostics.append(f"[map true_labels] {out}")
-    if not ok:
-        return HarmonizeResult(success=False, all_reviewed=False, messages=[f"calmate map true_labels failed: {out}"])
-
-    # --- 3. Check for unreviewed mappings ---
-    has_unreviewed = False
-    if store_path.exists():
-        try:
-            store_df = pd.read_csv(store_path)
-            if "reviewed" in store_df.columns:
-                reviewed_mask = store_df["reviewed"].astype(str).str.strip().str.lower().isin(["true", "1", "yes"])
-                n_unreviewed = int((~reviewed_mask).sum())
-                has_unreviewed = n_unreviewed > 0
-                logger.info("Calmate store: %d total, %d unreviewed", len(store_df), n_unreviewed)
-        except Exception as exc:
-            logger.warning("Could not read calmate store at %s: %s", store_path, exc)
-            has_unreviewed = True
-
-    if has_unreviewed:
-        msgs = diagnostics + [
-            "",
-            f"{n_unreviewed} label mapping(s) need manual review before metrics can be computed.",
-            "To review interactively:",
-            f"  {calmate_cli} --store {store_path} review",
-            "",
-            "Then re-run with:",
-            "  python scripts/run_scsa.py settings.skip_existing=false",
-        ]
-        return HarmonizeResult(success=True, all_reviewed=False, messages=msgs)
-
-    # --- 4. Apply mappings (predicted) ---
-    intermediate_csv = predictions_csv.with_name(predictions_csv.stem + "_mapped_pred.csv")
-    apply_pred_args = [
-        *store_arg,
-        "apply",
-        str(predictions_csv),
-        "-c",
-        "predicted_labels",
-        "-o",
-        str(intermediate_csv),
-        "--all",
-    ]
-    ok, out = _run_calmate(calmate_cli, apply_pred_args, "apply predicted_labels")
-    if not ok:
-        return HarmonizeResult(
-            success=False, all_reviewed=True, messages=[f"calmate apply predicted_labels failed: {out}"]
-        )
-
-    # --- 5. Apply mappings (true) ---
-    final_csv = predictions_csv.with_name(predictions_csv.stem + "_harmonized.csv")
-    apply_true_args = [*store_arg, "apply", str(intermediate_csv), "-c", "true_labels", "-o", str(final_csv), "--all"]
-    ok, out = _run_calmate(calmate_cli, apply_true_args, "apply true_labels")
-    if not ok:
-        return HarmonizeResult(success=False, all_reviewed=True, messages=[f"calmate apply true_labels failed: {out}"])
-
-    # Clean up intermediate file
-    if intermediate_csv.exists():
-        intermediate_csv.unlink()
-
-    return HarmonizeResult(success=True, all_reviewed=True, mapped_csv=final_csv)
-
-
-def process_scsa_dataset(
-    ds_cfg: Any,
-    run_cfg: Any,
-    output_root: str,
-    modules_dir: Path,
-    scsa_cfg: Any,
-    adata_cache: str,
-    hf_cache: str | None = None,
-    skip_existing: bool = True,
-    temp_dir: str | None = None,
-) -> tuple[str, bool, str | None, list[str]]:
-    """
-    Run the full SCSA pipeline for a single dataset.
-
-    Loads the dataset, runs clustering + SCSA annotation, harmonises labels
-    via calmate (if available), computes metrics, and writes results to
-    ``{output_root}/{dataset_name}/scsa/eval/metrics.csv``.
-
-    Parameters
-    ----------
-    ds_cfg : DictConfig
-        Dataset configuration (from dataset_list.yaml).
-    run_cfg : DictConfig
-        Run configuration (seed, n_rows).
-    output_root : str
-        Root directory for evaluation results.
-    modules_dir : Path
-        Directory containing the SCSA installation.
-    scsa_cfg : DictConfig
-        SCSA-specific parameters (foldchange, pvalue, species, etc.).
-    adata_cache : str
-        AnnData cache directory.
-    hf_cache : str or None
-        HuggingFace cache directory.
-    skip_existing : bool
-        If True, skip datasets that already have SCSA results.
-    temp_dir : str or None
-        Directory for intermediate files. Defaults to a subdir of output_root.
-
-    Returns
-    -------
-    tuple[str, bool, str | None, list[str]]
-        ``(dataset_name, success, error_message, calmate_messages)``
-    """
-    from mmcontext.embed.dataset_utils import collect_adata_subset, load_generic_dataset
-    from mmcontext.file_utils import save_table
-
-    dataset_name = ds_cfg.name
-    calmate_messages: list[str] = []
-
-    out_dir = Path(output_root) / dataset_name / "scsa"
-    eval_dir = out_dir / "eval"
-
-    remap_only = bool(scsa_cfg.get("remap_only", False))
-
-    # For full runs we can skip datasets that already have metrics.
-    if not remap_only and skip_existing and (eval_dir / "metrics.csv").exists():
-        logger.info("SCSA results already exist for %s — skipping", dataset_name)
-        return dataset_name, True, "skipped_existing", calmate_messages
-
-    if temp_dir is None:
-        temp_dir_path = out_dir / "temp"
-    else:
-        temp_dir_path = Path(temp_dir) / dataset_name
-    temp_dir_path.mkdir(parents=True, exist_ok=True)
-
-    scsa_dir = modules_dir / "SCSA"
-    scsa_db = scsa_dir / "whole_v2.db"
-
-    try:
-        # ------------------------------------------------------------------
-        # Fast path: re-run only the calmate mapping / metrics step
-        # ------------------------------------------------------------------
-        if remap_only:
-            logger.info("remap_only=True — reusing existing SCSA predictions for %s", dataset_name)
-
-            bio_labels = ds_cfg.get("bio_label_list", []) or []
-            all_metric_rows: list[dict] = []
-
-            calmate_store = Path(scsa_cfg.get("calmate_store", "modules/.calmate/mappings.csv"))
-            calmate_cache = scsa_cfg.get("calmate_cache", None)
-            calmate_python_path = scsa_cfg.get("calmate_python", None)
-
-            # Load existing adata (if available) so we can update ontology labels.
-            adata_path = eval_dir / "subset.h5ad"
-            adata: ad.AnnData | None = None
-            if adata_path.exists():
-                adata = ad.read_h5ad(adata_path)
-                logger.info("Loaded existing adata for remap_only from %s", adata_path)
-
-            for annotation_key in bio_labels:
-                pred_path = eval_dir / f"scsa_predictions_{annotation_key}.csv"
-                if not pred_path.exists():
-                    logger.warning(
-                        "Prediction file %s not found for %s / %s — skipping",
-                        pred_path,
-                        dataset_name,
-                        annotation_key,
-                    )
-                    continue
-
-                logger.info("Re-running calmate label harmonisation for '%s'", annotation_key)
-                harm = harmonize_labels(
-                    predictions_csv=pred_path,
-                    store_path=calmate_store,
-                    modules_dir=modules_dir,
-                    dataset_name=dataset_name,
-                    calmate_python=calmate_python_path,
-                    cache_dir=calmate_cache,
-                )
-
-                if not harm.success or not harm.all_reviewed:
-                    calmate_messages.extend(harm.messages)
-                    logger.warning(
-                        "Skipping metrics for %s / %s — label harmonisation incomplete",
-                        dataset_name,
-                        annotation_key,
-                    )
-                    continue
-
-                mapped_df = pd.read_csv(harm.mapped_csv)  # type: ignore[arg-type]
-                mapped_df = _reorder_harmonized_columns(mapped_df)
-                mapped_df.to_csv(harm.mapped_csv, index=False)  # keep on disk in the same order
-                true_col = "true_labels_mapped" if "true_labels_mapped" in mapped_df.columns else "true_labels"
-                pred_col = (
-                    "predicted_labels_mapped" if "predicted_labels_mapped" in mapped_df.columns else "predicted_labels"
-                )
-                logger.info("Harmonised CSV at %s (%d rows)", harm.mapped_csv, len(mapped_df))
-
-                # Store mapped true labels back into adata for downstream use
-                if adata is not None and true_col in mapped_df.columns:
-                    ontology_col = f"{annotation_key}_ontology"
-                    label_map = dict(zip(mapped_df["true_labels"], mapped_df[true_col], strict=False))
-                    adata.obs[ontology_col] = adata.obs[annotation_key].map(label_map).fillna(adata.obs[annotation_key])
-                    logger.info("Updated adata.obs['%s'] with ontology-mapped labels", ontology_col)
-
-                metric_rows = compute_metrics(
-                    mapped_df[true_col],
-                    mapped_df[pred_col],
-                    dataset_name,
-                    annotation_key,
-                )
-                all_metric_rows.extend(metric_rows)
-
-            if all_metric_rows:
-                metrics_df = pd.DataFrame(all_metric_rows)
-                save_table(metrics_df, eval_dir / "metrics", fmt="csv")
-                logger.info("Saved %d metrics for %s to %s (remap_only)", len(all_metric_rows), dataset_name, eval_dir)
-
-            if adata is not None:
-                adata.write_h5ad(adata_path)
-                logger.info("Re-saved adata with updated ontology columns to %s", adata_path)
-
-            return dataset_name, True, None, calmate_messages
-
-        # ------------------------------------------------------------------
-        # Full pipeline: download data, run SCSA, run calmate, compute metrics
-        # ------------------------------------------------------------------
-        # Load HF dataset to get share links and sample IDs
-        adata_download_dir = Path(adata_cache) / ds_cfg.name
-        cache_dir = hf_cache
-
-        raw_ds = load_generic_dataset(
-            source=ds_cfg.source,
-            fmt=ds_cfg.format,
-            split=ds_cfg.get("split", "test"),
-            max_rows=run_cfg.n_rows,
-            seed=run_cfg.seed,
-            cache_dir=cache_dir,
-        )
-
-        link_column = None
-        if "share_link" in raw_ds.column_names:
-            link_column = "share_link"
-        elif "adata_link" in raw_ds.column_names:
-            link_column = "adata_link"
-
-        if link_column is None:
-            return dataset_name, False, "No numeric data (adata_link/share_link) available for SCSA", calmate_messages
-
-        from mmcontext.file_utils import collect_unique_links, download_and_extract_links
-
-        links, _ = collect_unique_links(raw_ds, link_column=link_column)
-        adata_download_dir.mkdir(parents=True, exist_ok=True)
-        logger.info("Downloading %d adata chunks to %s ...", len(links), adata_download_dir)
-        download_and_extract_links(
-            links,
-            target_dir=adata_download_dir,
-            overwrite=False,
-        )
-
-        sample_ids = raw_ds[ds_cfg.index_col]
-        if isinstance(sample_ids[0], str) and "sample_idx:" in sample_ids[0]:
-            sample_ids = [sid.split(":")[1] for sid in sample_ids]
-
-        adata = collect_adata_subset(
-            download_dir=adata_download_dir,
-            sample_ids=sample_ids,
-        )
-
-        logger.info("Loaded AnnData for %s: %d obs x %d vars", dataset_name, adata.n_obs, adata.n_vars)
-
-        # Run SCSA pipeline: cluster -> markers -> SCSA -> map
-        marker_csv = temp_dir_path / f"{dataset_name}_markers.csv"
-        scsa_output = temp_dir_path / f"{dataset_name}_scsa_result.txt"
-
-        adata, cluster_col = create_marker_csv(
-            adata,
-            marker_csv,
-            cluster_col=scsa_cfg.get("cluster_col", "louvain"),
-            resolution=scsa_cfg.get("resolution", 1.0),
-            n_top_genes=scsa_cfg.get("n_top_genes", 2000),
-            random_state=run_cfg.seed,
-        )
-
-        scsa_python = scsa_cfg.get("scsa_python", None)
-
-        run_scsa(
-            csv_path=marker_csv,
-            db_path=scsa_db,
-            output_path=scsa_output,
-            scsa_dir=scsa_dir,
-            scsa_python=scsa_python,
-            foldchange=scsa_cfg.get("foldchange", 1.5),
-            pvalue=scsa_cfg.get("pvalue", 0.01),
-            species=scsa_cfg.get("species", "Human"),
-            tissue=scsa_cfg.get("tissue", "All"),
-            gensymbol=scsa_cfg.get("gensymbol", True),
-        )
-
-        # Evaluate for each bio label
-        bio_labels = ds_cfg.get("bio_label_list", []) or []
-        all_metric_rows: list[dict] = []
-
-        calmate_store = Path(scsa_cfg.get("calmate_store", "modules/.calmate/mappings.csv"))
-        calmate_cache = scsa_cfg.get("calmate_cache", None)
-        calmate_python_path = scsa_cfg.get("calmate_python", None)
-
-        for annotation_key in bio_labels:
-            if annotation_key not in adata.obs.columns:
-                logger.warning(
-                    "Annotation key '%s' not found in adata.obs for %s — skipping", annotation_key, dataset_name
-                )
-                continue
-
-            logger.info("Mapping SCSA predictions for label '%s'", annotation_key)
-            predictions_df = map_predictions_to_cells(scsa_output, adata, cluster_col, annotation_key)
-
-            eval_dir.mkdir(parents=True, exist_ok=True)
-            pred_path = eval_dir / f"scsa_predictions_{annotation_key}.csv"
-            predictions_df.to_csv(pred_path, index=False)
-            logger.info("Saved raw predictions to %s", pred_path)
-
-            # --- Label harmonisation via calmate ---
-            harm = harmonize_labels(
-                predictions_csv=pred_path,
-                store_path=calmate_store,
-                modules_dir=modules_dir,
-                dataset_name=dataset_name,
-                calmate_python=calmate_python_path,
-                cache_dir=calmate_cache,
-            )
-
-            if not harm.success or not harm.all_reviewed:
-                calmate_messages.extend(harm.messages)
-                logger.warning(
-                    "Skipping metrics for %s / %s — label harmonisation incomplete",
-                    dataset_name,
-                    annotation_key,
-                )
-                continue
-
-            # Read harmonised CSV and compute metrics on mapped columns
-            mapped_df = pd.read_csv(harm.mapped_csv)  # type: ignore[arg-type]
-            mapped_df = _reorder_harmonized_columns(mapped_df)
-            mapped_df.to_csv(harm.mapped_csv, index=False)
-            true_col = "true_labels_mapped" if "true_labels_mapped" in mapped_df.columns else "true_labels"
-            pred_col = (
-                "predicted_labels_mapped" if "predicted_labels_mapped" in mapped_df.columns else "predicted_labels"
-            )
-            logger.info("Harmonised CSV at %s (%d rows)", harm.mapped_csv, len(mapped_df))
-
-            # Store mapped true labels back into adata for downstream use
-            ontology_col = f"{annotation_key}_ontology"
-            if true_col in mapped_df.columns:
-                label_map = dict(zip(mapped_df["true_labels"], mapped_df[true_col], strict=False))
-                adata.obs[ontology_col] = adata.obs[annotation_key].map(label_map).fillna(adata.obs[annotation_key])
-                logger.info("Stored mapped true labels in adata.obs['%s']", ontology_col)
-
-            metric_rows = compute_metrics(
-                mapped_df[true_col],
-                mapped_df[pred_col],
-                dataset_name,
-                annotation_key,
-            )
-            all_metric_rows.extend(metric_rows)
-
-        if all_metric_rows:
-            metrics_df = pd.DataFrame(all_metric_rows)
-            save_table(metrics_df, eval_dir / "metrics", fmt="csv")
-            logger.info("Saved %d metrics for %s to %s", len(all_metric_rows), dataset_name, eval_dir)
-
-        # Re-save adata if ontology columns were added
-        ontology_cols = [c for c in adata.obs.columns if c.endswith("_ontology")]
-        if ontology_cols:
-            adata_out = eval_dir / "subset.h5ad"
-            adata.write_h5ad(adata_out)
-            logger.info("Saved adata with ontology columns to %s", adata_out)
-
-        logger.info("SCSA pipeline completed for %s", dataset_name)
-        return dataset_name, True, None, calmate_messages
-
-    except Exception as e:
-        logger.error("SCSA pipeline failed for %s: %s", dataset_name, e, exc_info=True)
-        return dataset_name, False, str(e), calmate_messages
diff --git a/src/mmcontext/eval/__init__.py b/src/mmcontext/eval/__init__.py
index 7772b56..4a962f6 100644
--- a/src/mmcontext/eval/__init__.py
+++ b/src/mmcontext/eval/__init__.py
@@ -1,18 +1,20 @@
-import importlib
-import inspect
-import pkgutil
-from pathlib import Path
+"""mmcontext.eval — Evaluation tools for mmcontext models.
 
-# Automatically import all classes from all submodules in pp
-__all__ = []
+Evaluators follow a registry pattern: import a module to register its class,
+then retrieve it by name via :func:`get`.  The evaluate-model-2.0 notebook
+uses ``LabelSimilarity`` and ``OmicsQueryAnnotator`` as its primary tools.
+"""
 
-package_dir = Path(__file__).resolve().parent
-for module_info in pkgutil.iter_modules([str(package_dir)]):
-    module = importlib.import_module(f".{module_info.name}", package=__name__)
+# Import evaluator modules so their @register decorators run.
+from . import ari, label_similarity
+from .base import BaseEvaluator, EvalResult
+from .query_annotate import OmicsQueryAnnotator
+from .registry import get, register
 
-    # Dynamically get all classes in the module and add them to __all__
-    for name, obj in inspect.getmembers(module):
-        if inspect.isclass(obj) or inspect.isfunction(obj):
-            if obj.__module__ == module.__name__:
-                globals()[name] = obj  # Add class to the global namespace
-                __all__.append(name)  # Add class to __all__ to make them importable
+__all__ = [
+    "BaseEvaluator",
+    "EvalResult",
+    "OmicsQueryAnnotator",
+    "get",
+    "register",
+]
diff --git a/src/mmcontext/eval/embedding_alignment.py b/src/mmcontext/eval/embedding_alignment.py
deleted file mode 100644
index 267576f..0000000
--- a/src/mmcontext/eval/embedding_alignment.py
+++ /dev/null
@@ -1,153 +0,0 @@
-import logging
-
-import numpy as np
-import pandas as pd
-from numpy.linalg import norm
-
-logger = logging.getLogger(__name__)
-
-
-def cosine_similarity(x: np.ndarray, y: np.ndarray) -> float:
-    """
-    Compute the cosine similarity between two vectors.
-
-    Parameters
-    ----------
-    x : np.ndarray
-        First embedding vector.
-    y : np.ndarray
-        Second embedding vector.
-
-    Returns
-    -------
-    float
-        Cosine similarity between x and y.
-    """
-    denom = norm(x) * norm(y)
-    return float(np.dot(x, y) / denom) if denom else 0.0
-
-
-def evaluate_modality_alignment(df: pd.DataFrame) -> (float, float):
-    """
-    Evaluate two cross-modal alignment scores for omics/text embeddings.
-
-    How close are pairs of cross-modal embeddings in the embedding space, compared to false pairs?
-    The input DataFrame is assumed to have exactly one 'omics' and one 'text'
-    embedding for each unique ID (e.g., cell ID) in order. Each row contains:
-    - A vector of float embeddings in one column (e.g. 'embedding'),
-    - A modality column with values in {'omics', 'text'}.
-
-    The function calculates two metrics:
-
-    1) Modality-gap-irrelevant score: For each (omics, text) pair i,
-       compare the similarity of that true pair only to other cross-modal
-       similarities (ignoring any intra-modal comparisons).
-
-       The score is the average fraction of cross-modal pairs whose similarity
-       is less than the true pair's similarity.
-
-    2) Full-comparison score: For each (omics, text) pair i,
-       compare the similarity of that true pair to *both* other cross-modal
-       pairs and all intra-modal pairs.
-
-       The score is again the average fraction of *all* other pairs (cross- or
-       intra-modal) whose similarity is less than the true pair's similarity.
-
-    Parameters
-    ----------
-    df : pd.DataFrame
-        DataFrame with columns:
-            - 'embedding_type': (e.g. 'omics' or 'text')
-            - 'embedding': vector of floats (list or np.ndarray)
-            - some ID columns to match pairs (e.g. 'sample_index').
-        It is assumed each ID has exactly two rows: one 'omics' and one 'text'.
-
-    Returns
-    -------
-    (float, float)
-        A tuple (score_modality_gap_irrelevant, score_full_comparison).
-    """
-    logger.info("Starting modality alignment evaluation.")
-
-    # Convert embedding column to numpy arrays (if stored as lists).
-    df["embedding"] = df["embedding"].apply(np.array)
-
-    # Separate data for each modality, indexing by some ID.
-    # Here we assume there's a column 'sample_index' that identifies the pairs.
-    # Adapt if your ID column is different.
-    omics_df = df[df["embedding_type"] == "omics"].copy()
-    text_df = df[df["embedding_type"] == "text"].copy()
-
-    # Merge omics and text on the ID to get direct pairs
-    # (Assuming 'sample_index' is the matching column).
-    merged = pd.merge(
-        omics_df[["sample_index", "embedding"]],
-        text_df[["sample_index", "embedding"]],
-        on="sample_index",
-        suffixes=("_omics", "_text"),
-    )
-
-    # Build lists for all cross-modal pairs
-    omics_embeddings = omics_df["embedding"].values
-    text_embeddings = text_df["embedding"].values
-
-    # For easy iteration over cross-modal combos
-    all_cross_pairs = []
-    for _i, emb_o in enumerate(omics_embeddings):
-        for _j, emb_t in enumerate(text_embeddings):
-            all_cross_pairs.append(cosine_similarity(emb_o, emb_t))
-    all_cross_pairs = np.array(all_cross_pairs)
-
-    # Build lists for all intra-modal pairs
-    # We do omics-omics and text-text separately.
-    all_omics_pairs = []
-    omics_embeddings_list = omics_embeddings.tolist()
-    for i in range(len(omics_embeddings_list)):
-        for j in range(i + 1, len(omics_embeddings_list)):
-            all_omics_pairs.append(cosine_similarity(omics_embeddings_list[i], omics_embeddings_list[j]))
-    all_omics_pairs = np.array(all_omics_pairs)
-
-    all_text_pairs = []
-    text_embeddings_list = text_embeddings.tolist()
-    for i in range(len(text_embeddings_list)):
-        for j in range(i + 1, len(text_embeddings_list)):
-            all_text_pairs.append(cosine_similarity(text_embeddings_list[i], text_embeddings_list[j]))
-    all_text_pairs = np.array(all_text_pairs)
-
-    # Prepare accumulators for each metric
-    ranks_modality_gap_irrelevant = []
-    ranks_full_comparison = []
-
-    # For each 'true' pair (row in merged), get its sim and rank it vs others
-    for _idx, row in merged.iterrows():
-        emb_o = row["embedding_omics"]
-        emb_t = row["embedding_text"]
-        true_sim = cosine_similarity(emb_o, emb_t)
-
-        # 1) Compare to cross-modal only
-        # Count how many cross-modal similarities are below 'true_sim'
-        # Subtract 1 if needed to exclude the pair i,i from the distribution
-        # but that’s typically negligible if each ID is unique.
-        cross_count_below = np.sum(all_cross_pairs < true_sim)
-        # Turn into a fraction
-        rank_cross = cross_count_below / len(all_cross_pairs)
-        ranks_modality_gap_irrelevant.append(rank_cross)
-
-        # 2) Compare to cross-modal + intra-modal
-        # Combine cross-modal array with omics-omics & text-text
-        # For a single “true” ID i, you might also exclude the sim of
-        # (emb_o, emb_o) or (emb_t, emb_t) if it appears, but typically
-        # we’re ignoring self-sim because it’s not in the data.
-        all_pairs_combined = np.concatenate([all_cross_pairs, all_omics_pairs, all_text_pairs])
-        count_below = np.sum(all_pairs_combined < true_sim)
-        rank_full = count_below / len(all_pairs_combined)
-        ranks_full_comparison.append(rank_full)
-
-    # Average rank across all pairs
-    modality_gap_irrelevant_score = float(np.mean(ranks_modality_gap_irrelevant))
-    full_comparison_score = float(np.mean(ranks_full_comparison))
-
-    logger.info("Modality-gap-irrelevant score: %f", modality_gap_irrelevant_score)
-    logger.info("Full-comparison score: %f", full_comparison_score)
-
-    return modality_gap_irrelevant_score, full_comparison_score
diff --git a/src/mmcontext/eval/eval_pipeline.py b/src/mmcontext/eval/eval_pipeline.py
deleted file mode 100644
index 983102f..0000000
--- a/src/mmcontext/eval/eval_pipeline.py
+++ /dev/null
@@ -1,946 +0,0 @@
-# embedding_benchmark/eval_pipeline.py
-import importlib
-import logging
-import pkgutil
-from pathlib import Path
-from typing import Any
-
-import anndata as ad
-import numpy as np
-import pandas as pd
-
-import mmcontext.eval as ev_pkg
-from mmcontext.eval.registry import get as get_evaluator
-
-# from mmcontext.adata_utils import collect_adata_subset
-from mmcontext.eval.utils import LabelKind, LabelSpec
-from mmcontext.file_utils import save_table
-
-# ---- discover all evaluator modules so decorators run ---------------
-for m in pkgutil.walk_packages(ev_pkg.__path__, prefix=ev_pkg.__name__ + "."):
-    importlib.import_module(m.name)
-
-logger = logging.getLogger(__name__)
-
-
-def _find_project_root(marker_name: str = "mmcontext") -> Path | None:
-    """
-    Best-effort project root discovery.
-
-    This is used to locate default cache directories when Hydra changes CWD.
-    """
-    try:
-        cur = Path.cwd().resolve()
-        while True:
-            if cur.name == marker_name:
-                return cur
-            if cur.parent == cur:
-                return None
-            cur = cur.parent
-    except Exception:
-        return None
-
-
-def _default_hf_cache_dir() -> Path | None:
-    """Default HF cache path used by this repo: <project_root>/cache/huggingface."""
-    root = _find_project_root()
-    if root is None:
-        return None
-    return root / "cache" / "huggingface"
-
-
-def _read_cellwhisperer_logit_scale_from_emb_dir(emb_dir: Path) -> float | None:
-    """
-    Read logit_scale persisted during CellWhisperer embedding generation.
-
-    Expected file: emb_dir/cellwhisperer_meta.json with {"logit_scale": <float>}.
-    """
-    meta_path = emb_dir / "cellwhisperer_meta.json"
-    if not meta_path.exists():
-        return None
-
-    try:
-        import json
-
-        payload = json.loads(meta_path.read_text())
-        val = payload.get("logit_scale")
-        if val is None:
-            return None
-        return float(val)
-    except Exception as e:
-        logger.warning(f"Failed reading {meta_path} for logit_scale: {e}")
-        return None
-
-
-def _extract_cellwhisperer_logit_scale(modules_dir: Path, hf_cache: str | None = None) -> float | None:
-    """
-    Extract logit_scale from CellWhisperer checkpoint **without** importing CellWhisperer.
-
-    Parameters
-    ----------
-    modules_dir : Path
-        Directory containing CellWhisperer installation (kept for backward compatibility)
-    hf_cache : str, optional
-        HuggingFace cache directory
-
-    Returns
-    -------
-    float or None
-        Extracted logit_scale value, or None if extraction fails
-    """
-    try:
-        import torch
-
-        # Preferred location: HF cache (matches CellWhisperer standalone runner behavior)
-        cache_dir = Path(hf_cache) if hf_cache else _default_hf_cache_dir()
-
-        ckpt_candidates: list[Path] = []
-        if cache_dir is not None:
-            ckpt_candidates.extend(
-                [
-                    cache_dir / "cellwhisperer_emb" / "cellwhisperer_jointemb_v1.ckpt",
-                    cache_dir / "cellwhisperer_emb" / "cellwhisperer_clip_v1.ckpt",
-                ]
-            )
-
-        # Backward-compatible fallbacks (older layouts)
-        ckpt_candidates.extend(
-            [
-                modules_dir / "cellwhisperer_emb" / "cellwhisperer_jointemb_v1.ckpt",
-                modules_dir / "cellwhisperer_emb" / "cellwhisperer_clip_v1.ckpt",
-            ]
-        )
-
-        ckpt_path = next((p for p in ckpt_candidates if p.exists()), None)
-        if ckpt_path is None:
-            logger.warning(
-                "CellWhisperer checkpoint not found in any expected location, cannot extract logit_scale. "
-                f"Tried: {[str(p) for p in ckpt_candidates]}"
-            )
-            return None
-
-        logger.info(f"Loading CellWhisperer checkpoint to extract logit_scale from {ckpt_path}")
-
-        try:
-            ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
-        except TypeError:
-            ckpt = torch.load(ckpt_path, map_location="cpu")
-
-        state_dict = ckpt.get("state_dict", ckpt) if isinstance(ckpt, dict) else None
-        if not isinstance(state_dict, dict):
-            logger.warning(
-                "CellWhisperer checkpoint format unexpected (no dict state_dict), cannot extract logit_scale"
-            )
-            return None
-
-        # Try to find the discriminator temperature parameter
-        temp_key = None
-        for k in state_dict.keys():
-            if "discriminator.temperature" in k:
-                temp_key = k
-                break
-        if temp_key is None:
-            logger.warning("Could not find discriminator.temperature in checkpoint state_dict")
-            return None
-
-        temperature = state_dict[temp_key]
-        if not isinstance(temperature, torch.Tensor):
-            logger.warning(f"Checkpoint key {temp_key} is not a tensor, cannot extract logit_scale")
-            return None
-
-        logit_scale = float(torch.exp(temperature).item())
-        logger.info(f"✓ Extracted CellWhisperer logit_scale: {logit_scale}")
-
-        return logit_scale
-
-    except Exception as e:
-        logger.warning(f"Failed to extract CellWhisperer logit_scale: {e}")
-        return None
-
-
-def run_scib_evaluation(cfg) -> None:
-    """
-    Run ScIB evaluation separately for all dataset/model combinations.
-
-    ScIB has its own interface and internal parallelization, so it's handled
-    separately from the main evaluation pipeline.
-
-    Args:
-        cfg: Hydra configuration object
-    """
-    print("\n=== Running ScIB Evaluation ===")
-
-    # Create list of all dataset/model combinations
-    tasks = []
-    for ds_cfg in cfg.datasets:
-        for model_cfg in cfg.models:
-            tasks.append((ds_cfg, model_cfg))
-
-    print(f"Total dataset/model combinations for ScIB: {len(tasks)}")
-
-    successful = 0
-    failed = 0
-    skipped = 0
-
-    for i, (ds_cfg, model_cfg) in enumerate(tasks):
-        dataset_name = ds_cfg.name
-        model_id = model_cfg.source
-        model_name = model_cfg.get("name", model_cfg.source)  # Use name if available, fallback to source
-        text_only = model_cfg.get("text_only", False)
-        model_dir_name = model_name
-        if text_only:
-            model_id = model_id + "_text_only"
-            model_dir_name = model_dir_name + "_text_only"
-
-        print(f"\n=== [{i + 1}/{len(tasks)}] ScIB for: {dataset_name}/{model_id} (name: {model_name}) ===")
-
-        emb_dir = Path(cfg.output.root) / dataset_name / model_dir_name
-
-        # ── Check if eval directory exists and skip if requested ─────────────
-        skip_if_eval_exists = cfg.eval.get("skip_if_eval_exists", False)
-        eval_dir = emb_dir / "eval"
-
-        if skip_if_eval_exists and eval_dir.exists():
-            print(f"↷ Skipping {dataset_name}/{model_id} - eval directory already exists: {eval_dir}")
-            skipped += 1
-            continue
-
-        # Check if required files exist
-        embeddings_file = emb_dir / "embeddings.parquet"
-        adata_file = emb_dir / "subset.h5ad"
-
-        if not embeddings_file.exists() or not adata_file.exists():
-            print(f"✗ Missing required files for {dataset_name}/{model_id}")
-            failed += 1
-            continue
-
-        try:
-            # Load data
-            print(f"Loading data from: {emb_dir}")
-            emb_df = pd.read_parquet(embeddings_file)
-            E1 = np.vstack(emb_df["embedding"].to_numpy())
-            adata = ad.read_h5ad(adata_file)
-
-            print(f"Loaded embeddings: {E1.shape}, AnnData: {adata.n_obs} obs × {adata.n_vars} vars")
-
-            # Pre-validate data for ScIB
-            if E1.size == 0:
-                print(f"✗ Empty embedding array for {dataset_name}/{model_id}")
-                failed += 1
-                continue
-
-            if adata.n_obs == 0:
-                print(f"✗ Empty AnnData for {dataset_name}/{model_id}")
-                failed += 1
-                continue
-
-            # Run ScIB evaluation
-            ScibClass = get_evaluator("scib")
-            scib_evaluator = ScibClass()
-
-            # Handle cases where label lists might be None or empty
-            bio_labels = ds_cfg.bio_label_list or []
-            batch_labels = ds_cfg.batch_label_list or []
-
-            scib_results = scib_evaluator.compute_dataset_model(
-                emb1=E1,
-                adata=adata,
-                dataset_name=dataset_name,
-                model_id=model_id,
-                model_name=model_name,
-                bio_labels=bio_labels,
-                batch_labels=batch_labels,
-                **cfg.eval,
-            )
-
-            # Save results
-            if scib_results:
-                scib_df = pd.DataFrame(scib_results)
-                eval_dir = emb_dir / "eval"
-                eval_dir.mkdir(parents=True, exist_ok=True)
-                save_table(scib_df, eval_dir / "scib_metrics", fmt="csv")
-                print(f"✓ Saved {len(scib_results)} ScIB metrics")
-                successful += 1
-            else:
-                print(f"✗ No ScIB results generated for {dataset_name}/{model_id}")
-                failed += 1
-
-        except Exception as e:
-            print(f"✗ ScIB evaluation failed for {dataset_name}/{model_id}: {e}")
-
-            # Save error to file
-            try:
-                eval_dir = emb_dir / "eval"
-                eval_dir.mkdir(parents=True, exist_ok=True)
-                error_results = [
-                    {
-                        "dataset": dataset_name,
-                        "model": model_id,
-                        "model_name": model_name,
-                        "bio_label": "unknown",
-                        "batch_label": "unknown",
-                        "metric": "scib/error",
-                        "value": str(e),
-                        "data_id": "",
-                        "hvg": "",
-                        "type": "",
-                    }
-                ]
-                error_df = pd.DataFrame(error_results)
-                save_table(error_df, eval_dir / "scib_metrics", fmt="csv")
-            except Exception as save_error:
-                print(f"  Also failed to save error: {save_error}")
-
-            failed += 1
-
-    print("\n=== ScIB Evaluation Summary ===")
-    print(f"Total tasks: {len(tasks)}")
-    print(f"Successful: {successful}")
-    print(f"Skipped: {skipped}")
-    print(f"Failed: {failed}")
-    print(f"Success rate: {successful / len(tasks) * 100:.1f}%")
-
-
-def process_single_dataset_model(
-    ds_cfg: Any, model_cfg: Any, eval_cfg: dict[str, Any], output_root: str, modules_dir: Path | None = None
-) -> list[dict]:
-    """
-    Process a single dataset/model combination for standard evaluators only.
-
-    ScIB evaluation is handled separately in run_scib_evaluation().
-
-    Args:
-        ds_cfg: Dataset configuration
-        model_cfg: Model configuration
-        eval_cfg: Evaluation configuration
-        output_root: Output root directory
-        modules_dir: Path to modules directory (for CellWhisperer logit_scale extraction)
-
-    Returns
-    -------
-        List of evaluation results as dictionaries
-    """
-    dataset_name = ds_cfg.name
-    model_id = model_cfg.source
-    model_name = model_cfg.get("name", model_cfg.source)  # Use name if available, fallback to source
-    text_only = model_cfg.get("text_only", False)
-    if text_only:
-        model_name = model_name + "_text_only"
-
-    plot_only = eval_cfg.get("plot_only", False)
-    skip_missing_cache = eval_cfg.get("skip_missing_cache", True)
-
-    if plot_only:
-        logger.info(f"Plot-only mode: {model_name} for dataset: {dataset_name}")
-    else:
-        logger.info(f"Processing model: {model_name} for dataset: {dataset_name}")
-
-    # Handle cases where label lists might be None or empty
-    bio_labels = ds_cfg.bio_label_list or []
-    batch_labels = ds_cfg.batch_label_list or []
-
-    label_specs = [LabelSpec(n, LabelKind.BIO) for n in bio_labels] + [
-        LabelSpec(n, LabelKind.BATCH) for n in batch_labels
-    ]
-
-    rows = []
-
-    emb_dir = Path(output_root) / dataset_name / Path(model_name).name.replace("/", "_")
-    logger.info(f"Looking for embeddings in: {emb_dir}")
-
-    # ── Check if eval directory exists and skip if requested ─────────────
-    skip_if_eval_exists = eval_cfg.get("skip_if_eval_exists", False)
-    eval_dir = emb_dir / "eval"
-
-    if skip_if_eval_exists and eval_dir.exists():
-        logger.info(f"Skipping {dataset_name}/{model_name} - eval directory already exists: {eval_dir}")
-        return []  # Return empty list to indicate skipped
-
-    # ── Handle plot-only mode ───────────────────────────────────────
-    if plot_only:
-        return process_plot_only_mode(
-            ds_cfg,
-            model_cfg,
-            eval_cfg,
-            emb_dir,
-            label_specs,
-            dataset_name,
-            model_id,
-            model_name,
-            skip_missing_cache,
-            modules_dir,
-        )
-
-    # ── load shared artefacts *once* per model ───────────────────────
-    try:
-        logger.info(f"Loading embeddings from: {emb_dir / 'embeddings.parquet'}")
-        emb_df = pd.read_parquet(emb_dir / "embeddings.parquet")
-        logger.info(f"✓ Loaded embeddings: {len(emb_df)} rows")
-
-        E1 = np.vstack(emb_df["embedding"].to_numpy())
-        logger.info(f"✓ Stacked embeddings: {E1.shape}")
-
-        logger.info(f"Loading AnnData from: {emb_dir / 'subset.h5ad'}")
-        adata = ad.read_h5ad(emb_dir / "subset.h5ad")
-        logger.info(f"✓ Loaded AnnData: {adata.n_obs} obs × {adata.n_vars} vars")
-
-    except FileNotFoundError as e:
-        logger.error(f"✗ Missing file for {dataset_name}/{model_name}: {e}")
-        raise
-    except Exception as e:
-        logger.error(f"✗ Error loading data for {dataset_name}/{model_name}: {e}")
-        raise
-
-    # ──── Handle CellWhisperer vs other models ────────────────────
-    # Check if we're dealing with a CellWhisperer model
-    is_cellwhisperer = (
-        model_id == "cellwhisperer"
-        or model_name.lower() == "cellwhisperer"
-        or str(model_cfg.get("source", "")).lower() == "cellwhisperer"
-    )
-
-    # Make a copy of eval_cfg to modify
-    eval_cfg = eval_cfg.copy()
-
-    # Set parameters based on model type
-    if is_cellwhisperer:
-        # For CellWhisperer: use dot similarity + softmax + logit_scale
-        logger.info("Detected CellWhisperer model - using dot similarity + softmax normalization")
-        eval_cfg["similarity"] = "dot"
-        eval_cfg["score_norm_method"] = "softmax"
-
-        # Get logit_scale from eval_cfg, or extract it if None
-        logit_scale = eval_cfg.get("logit_scale")
-        if logit_scale is None:
-            # Prefer persisted metadata produced during CellWhisperer embedding generation
-            persisted_scale = _read_cellwhisperer_logit_scale_from_emb_dir(emb_dir)
-            if persisted_scale is not None:
-                logit_scale = persisted_scale
-                logger.info(f"Using persisted CellWhisperer logit_scale from embeddings dir: {logit_scale}")
-                eval_cfg["logit_scale"] = logit_scale
-            else:
-                # Backward-compatible fallback for older runs: read from checkpoint without importing cellwhisperer
-                if modules_dir is not None and modules_dir.exists():
-                    extracted_scale = _extract_cellwhisperer_logit_scale(modules_dir, hf_cache=eval_cfg.get("hf_cache"))
-                    if extracted_scale is not None:
-                        logit_scale = extracted_scale
-                        logger.info(f"Using extracted CellWhisperer logit_scale from checkpoint: {logit_scale}")
-                        eval_cfg["logit_scale"] = logit_scale
-                    else:
-                        logger.warning("Could not extract CellWhisperer logit_scale, using default (1.0)")
-                        eval_cfg["logit_scale"] = 1.0
-                else:
-                    logger.warning(
-                        f"modules_dir not provided or doesn't exist: {modules_dir}, using default logit_scale (1.0)"
-                    )
-                    eval_cfg["logit_scale"] = 1.0
-    else:
-        # For all other models: use cosine similarity + no normalization + no logit_scale
-        logger.info("Detected standard model - using cosine similarity + no normalization")
-        eval_cfg["similarity"] = "cosine"
-        eval_cfg["score_norm_method"] = None
-        eval_cfg["logit_scale"] = None  # Will default to 1.0 in LabelSimilarity.__init__
-
-    # ──── Run standard evaluators ────────────────────
-    evaluators = eval_cfg.get("suite", [])
-    logger.info(f"Running evaluators: {evaluators}")
-
-    # try to load label embeddings only once
-    label_emb_cache = {}  # (kind, name) → tuple[ndarray, dict] | None
-    # ------------- ITERATE over labels AND evaluators --------------
-    for label_spec in label_specs:
-        if label_spec.name not in adata.obs.columns:
-            logger.info(f"  Skipping label {label_spec.name} - not found in adata.obs")
-            continue  # skip silently
-        logger.info(f"  Processing label: {label_spec.name} ({label_spec.kind})")
-        y = adata.obs[label_spec.name].unique().tolist()
-        if (label_spec.kind, label_spec.name) not in label_emb_cache:
-            prefix = "bio_label_embeddings" if label_spec.kind == LabelKind.BIO else "batch_label_embeddings"
-            path = emb_dir / f"{prefix}_{label_spec.name}.parquet"
-            mapping_path = emb_dir / f"{prefix}_{label_spec.name}_mapping.json"
-            if path.exists():
-                try:
-                    df2 = pd.read_parquet(path)
-                    embeddings = np.vstack(df2["embedding"].to_numpy())
-
-                    # Load label mapping if available
-                    label_to_index = None
-                    if mapping_path.exists():
-                        import json
-
-                        with open(mapping_path) as f:
-                            label_to_index = json.load(f)
-                        logger.info(f"    ✓ Loaded label mapping for {label_spec.name}")
-                    else:
-                        logger.warning(f"    - No label mapping found for {label_spec.name}, using legacy format")
-
-                    label_emb_cache[(label_spec.kind, label_spec.name)] = (embeddings, label_to_index)
-                    logger.info(f"    ✓ Loaded label embeddings for {label_spec.name}")
-                except Exception as e:
-                    logger.error(f"    ✗ Error loading label embeddings for {label_spec.name}: {e}")
-                    label_emb_cache[(label_spec.kind, label_spec.name)] = None
-            else:
-                logger.info(f"    - No label embeddings found for {label_spec.name}")
-                label_emb_cache[(label_spec.kind, label_spec.name)] = None
-
-        # Extract embeddings and mapping
-        cache_entry = label_emb_cache[(label_spec.kind, label_spec.name)]
-        if cache_entry is not None:
-            E2, label_to_index = cache_entry
-        else:
-            E2, label_to_index = None, None
-
-        for ev_name in evaluators:
-            logger.info(f"    Running evaluator: {ev_name}")
-            try:
-                EvClass = get_evaluator(ev_name)
-
-                # Extract LabelSimilarity-specific parameters for constructor
-                # These are constructor parameters, not method parameters
-                evaluator_kwargs = {}
-                if ev_name == "LabelSimilarity":
-                    if "similarity" in eval_cfg:
-                        evaluator_kwargs["similarity"] = eval_cfg["similarity"]
-                    if "logit_scale" in eval_cfg:
-                        evaluator_kwargs["logit_scale"] = eval_cfg["logit_scale"]
-                    if "score_norm_method" in eval_cfg:
-                        evaluator_kwargs["score_norm_method"] = eval_cfg["score_norm_method"]
-
-                ev = EvClass(**evaluator_kwargs)
-
-                # Remove constructor parameters from eval_cfg before passing to compute/plot
-                method_eval_cfg = {
-                    k: v for k, v in eval_cfg.items() if k not in ["similarity", "logit_scale", "score_norm_method"]
-                }
-
-                result = ev.compute(
-                    omics_embeddings=E1,
-                    label_embeddings=E2,
-                    query_labels=y,
-                    true_labels=adata.obs[label_spec.name],
-                    label_key=label_spec.name,
-                    out_dir=emb_dir,
-                    **method_eval_cfg,
-                )
-
-                for key, val in result.items():
-                    rows.append(
-                        {
-                            "dataset": dataset_name,
-                            "model": model_id,
-                            "model_name": model_name,
-                            "label": label_spec.name,
-                            "label_kind": label_spec.kind,
-                            "metric": f"{ev_name}/{key}",
-                            "value": val,
-                        }
-                    )
-                logger.info(f"      ✓ {ev_name} completed")
-
-                if ev.produces_plot:
-                    logger.info(f"      Generating plots for {ev_name}")
-                    # …/eval/<Evaluator>/<label-name>/figure.png
-                    plot_dir = (
-                        emb_dir / "eval" / ev_name / label_spec.name  # <— NEW: sub-folder per label value
-                    )
-                    plot_dir.mkdir(parents=True, exist_ok=True)
-
-                    ev.plot(
-                        omics_embeddings=E1,
-                        out_dir=plot_dir,
-                        label_embeddings=E2,
-                        query_labels=y,
-                        true_labels=adata.obs[label_spec.name],
-                        label_key=label_spec.name,
-                        **method_eval_cfg,
-                    )
-                    logger.info(f"      ✓ Plots saved to {plot_dir}")
-            except Exception as e:
-                logger.error(f"      ✗ Error in evaluator {ev_name}: {e}")
-                # Continue with next evaluator instead of crashing
-                continue
-
-    # ──── Save results ────────────────────────────────────────────────
-    if rows:
-        out_df = pd.DataFrame(rows)
-        save_table(out_df, emb_dir / "eval/metrics", fmt="csv")
-        logger.info(f"✓ Saved {len(rows)} evaluation results")
-
-    logger.info(f"✓ Completed processing model {model_name} for dataset {dataset_name}")
-
-    return rows
-
-
-def process_plot_only_mode(
-    ds_cfg: Any,
-    model_cfg: Any,
-    eval_cfg: dict[str, Any],
-    emb_dir: Path,
-    label_specs: list,
-    dataset_name: str,
-    model_id: str,
-    model_name: str,
-    skip_missing_cache: bool,
-    modules_dir: Path | None = None,
-) -> list[dict]:
-    """
-    Process a single dataset/model combination in plot-only mode.
-
-    This function regenerates plots using cached similarity matrices without recomputing embeddings.
-
-    Args:
-        ds_cfg: Dataset configuration
-        model_cfg: Model configuration
-        eval_cfg: Evaluation configuration
-        emb_dir: Embedding directory path
-        label_specs: List of label specifications
-        dataset_name: Dataset name
-        model_id: Model ID
-        model_name: Model name
-        skip_missing_cache: Whether to skip if cache is missing
-
-    Returns
-    -------
-        List of evaluation results (empty in plot-only mode since we're not recomputing)
-    """
-    evaluators = eval_cfg.get("suite", [])
-    logger.info(f"Plot-only mode: Regenerating plots for evaluators: {evaluators}")
-
-    generated_any_plots = False
-
-    for label_spec in label_specs:
-        logger.info(f"  Processing label: {label_spec.name} ({label_spec.kind})")
-
-        for ev_name in evaluators:
-            try:
-                EvClass = get_evaluator(ev_name)
-                ev = EvClass()
-
-                # Only process evaluators that produce plots and support plot-only mode
-                if not ev.produces_plot:
-                    logger.info(f"      Skipping {ev_name} - doesn't produce plots")
-                    continue
-
-                # Check if evaluator supports plot_only method (currently only LabelSimilarity)
-                if not hasattr(ev, "plot_only"):
-                    logger.info(f"      Skipping {ev_name} - doesn't support plot-only mode")
-                    continue
-
-                plot_dir = emb_dir / "eval" / ev_name / label_spec.name
-                cache_file = plot_dir / "label_similarity_cache.pkl"
-
-                if not cache_file.exists():
-                    if skip_missing_cache:
-                        logger.info(f"      Skipping {ev_name}/{label_spec.name} - no cache found")
-                        continue
-                    else:
-                        logger.error(f"      ✗ Cache not found for {ev_name}/{label_spec.name}: {cache_file}")
-                        continue
-
-                logger.info(f"      Regenerating plots for {ev_name} using cached data")
-                plot_dir.mkdir(parents=True, exist_ok=True)
-
-                # Use plot_only method for supported evaluators
-                ev.plot_only(
-                    out_dir=plot_dir,
-                    label_kind=label_spec.kind,
-                    label_key=label_spec.name,
-                    **eval_cfg,  # forward any extra Hydra knobs
-                )
-
-                logger.info(f"      ✓ Plots regenerated for {ev_name}")
-                generated_any_plots = True
-
-            except Exception as e:
-                logger.error(f"      ✗ Error regenerating plots for {ev_name}: {e}")
-                continue
-
-    if generated_any_plots:
-        logger.info(f"✓ Plot-only mode completed for {model_id} for dataset {dataset_name}")
-    else:
-        logger.warning(f"⚠ No plots were regenerated for {model_id} for dataset {dataset_name}")
-
-    # Return empty list since we're not computing new metrics in plot-only mode
-    return []
-
-
-def _process_single_worker(args):
-    """
-    Worker function for parallel processing of standard evaluators.
-
-    Args:
-        args: Tuple of (ds_cfg, model_cfg, eval_cfg, output_root, modules_dir)
-
-    Returns
-    -------
-        Tuple of (dataset_name, model_id, success, result_count, error_msg, was_skipped)
-    """
-    import os
-    import sys
-    import time
-    import traceback
-
-    if len(args) == 5:
-        ds_cfg, model_cfg, eval_cfg, output_root, modules_dir = args
-    else:
-        # Backwards compatibility
-        ds_cfg, model_cfg, eval_cfg, output_root = args
-        modules_dir = None
-    dataset_name = ds_cfg.name
-    model_id = model_cfg.source
-    model_name = model_cfg.get("name", model_cfg.source)  # Use name if available, fallback to source
-    text_only = model_cfg.get("text_only", False)
-    if text_only:
-        model_id = model_id + "_text_only"
-
-    # Configure logging for this worker process
-    worker_logger = logging.getLogger(f"eval_worker_{os.getpid()}")
-    worker_logger.setLevel(logging.INFO)
-
-    # Create a handler that writes to stdout so we can see worker logs
-    if not worker_logger.handlers:
-        handler = logging.StreamHandler(sys.stdout)
-        handler.setFormatter(logging.Formatter(f"[Worker {os.getpid()}] %(levelname)s - %(message)s"))
-        worker_logger.addHandler(handler)
-
-        worker_logger.info(f"Starting processing: {dataset_name}/{model_id} (name: {model_name})")
-    start_time = time.time()
-
-    try:
-        results = process_single_dataset_model(ds_cfg, model_cfg, eval_cfg, output_root, modules_dir)
-
-        # Check if this was skipped (empty results + skip flag enabled)
-        was_skipped = False
-        if len(results) == 0 and eval_cfg.get("skip_if_eval_exists", False):
-            eval_dir = Path(output_root) / dataset_name / Path(model_name).name.replace("/", "_") / "eval"
-            if eval_dir.exists():
-                was_skipped = True
-
-        elapsed_time = time.time() - start_time
-        worker_logger.info(
-            f"Completed processing: {dataset_name}/{model_id} (name: {model_name}) in {elapsed_time:.1f}s"
-        )
-        return dataset_name, model_id, True, len(results), None, was_skipped
-
-    except Exception as e:
-        # Get full traceback for debugging
-        elapsed_time = time.time() - start_time
-        error_msg = f"{str(e)}\n{traceback.format_exc()}"
-        worker_logger.error(
-            f"Error processing {dataset_name}/{model_id} (name: {model_name}) after {elapsed_time:.1f}s: {error_msg}"
-        )
-        return dataset_name, model_id, False, 0, str(e), False
-
-
-def run_eval_suite(cfg) -> None:
-    """
-    Run the evaluation suite.
-
-    Is called from the eval.py script.
-    """
-    import concurrent.futures
-    import multiprocessing as mp
-    import time
-    from concurrent.futures import ProcessPoolExecutor, as_completed
-
-    # Check if we're in plot-only mode
-    plot_only = cfg.eval.get("plot_only", False)
-
-    if plot_only:
-        print("\n=== Running Evaluation Suite in Plot-Only Mode ===")
-        print("This will regenerate plots using cached similarity matrices.")
-        print("Skipping ScIB evaluation (not supported in plot-only mode).")
-    else:
-        print("\n=== Running Full Evaluation Suite ===")
-
-    # Create list of all dataset/model combinations
-    modules_dir = Path(cfg.eval.get("modules_dir", "")) if cfg.eval.get("modules_dir") else None
-    tasks = []
-    for ds_cfg in cfg.datasets:
-        for model_cfg in cfg.models:
-            tasks.append((ds_cfg, model_cfg, cfg.eval, cfg.output.root, modules_dir))
-
-    print(f"Total dataset/model combinations to process: {len(tasks)}")
-
-    # Check if parallel processing is enabled
-    enable_parallel = cfg.eval.get("enable_parallel", True)
-    max_workers = cfg.eval.get("max_workers", None)
-
-    if max_workers is None:
-        max_workers = mp.cpu_count()
-
-    start_time = time.time()
-
-    if not enable_parallel or len(tasks) == 1:
-        # Sequential processing
-        print("Running in sequential mode...")
-        successful = 0
-        failed = 0
-        skipped = 0
-
-        for i, task in enumerate(tasks):
-            if len(task) == 5:
-                ds_cfg, model_cfg, eval_cfg, output_root, modules_dir = task
-            else:
-                ds_cfg, model_cfg, eval_cfg, output_root = task
-                modules_dir = None
-            dataset_name = ds_cfg.name
-            model_id = model_cfg.source
-            model_name = model_cfg.get("name", model_cfg.source)  # Use name if available, fallback to source
-            text_only = model_cfg.get("text_only", False)
-            if text_only:
-                model_id = model_id + "_text_only"
-
-            task_start = time.time()
-            print(
-                f"\n=== [{i + 1}/{len(tasks)}] Processing dataset: {dataset_name}, model: {model_id} (name: {model_name}) ==="
-            )
-            try:
-                results = process_single_dataset_model(ds_cfg, model_cfg, eval_cfg, output_root, modules_dir)
-                task_time = time.time() - task_start
-
-                if len(results) == 0 and eval_cfg.get("skip_if_eval_exists", False):
-                    # Check if this was actually skipped (empty results + skip flag enabled)
-                    eval_dir = Path(output_root) / dataset_name / Path(model_name).name.replace("/", "_") / "eval"
-                    if eval_dir.exists():
-                        skipped += 1
-                        print(
-                            f"↷ Skipped {model_id} (name: {model_name}) for dataset {dataset_name} - eval directory exists ({task_time:.1f}s)"
-                        )
-                        continue
-
-                successful += 1
-                print(
-                    f"✓ Completed processing model {model_id} (name: {model_name}) for dataset {dataset_name} ({task_time:.1f}s)"
-                )
-                print(f"  → {len(results)} evaluation metrics")
-
-                # Estimate remaining time
-                if i > 0:
-                    avg_time = (time.time() - start_time) / (i + 1)
-                    remaining_time = avg_time * (len(tasks) - i - 1)
-                    print(f"  → Estimated time remaining: {remaining_time:.1f}s")
-
-            except Exception as e:
-                failed += 1
-                print(f"✗ Error processing model {model_id} (name: {model_name}) for dataset {dataset_name}: {e}")
-                print("  Continuing with next combination...")
-                continue
-
-        total_time = time.time() - start_time
-        print("\n=== Summary ===")
-        print(f"Total tasks: {len(tasks)}")
-        print(f"Successful: {successful}")
-        print(f"Skipped: {skipped}")
-        print(f"Failed: {failed}")
-        print(f"Success rate: {successful / len(tasks) * 100:.1f}%")
-        print(f"Total time: {total_time:.1f}s")
-        print(f"Average time per task: {total_time / len(tasks):.1f}s")
-    else:
-        # Parallel processing
-        print(f"Running in parallel mode with {max_workers} workers...")
-
-        # Get timeout from config (default: 1 hour per task)
-        task_timeout = cfg.eval.get("task_timeout", 3600)
-        print(f"Task timeout: {task_timeout}s ({task_timeout / 60:.1f} minutes)")
-
-        with ProcessPoolExecutor(max_workers=max_workers) as executor:
-            # Submit all tasks
-            future_to_task = {}
-            for i, task in enumerate(tasks):
-                if len(task) == 5:
-                    ds_cfg, model_cfg, eval_cfg, output_root, modules_dir = task
-                else:
-                    ds_cfg, model_cfg, eval_cfg, output_root = task
-                    modules_dir = None
-                dataset_name = ds_cfg.name
-                model_id = model_cfg.source
-                model_name = model_cfg.get("name", model_cfg.source)  # Use name if available, fallback to source
-                text_only = model_cfg.get("text_only", False)
-                if text_only:
-                    model_id = model_id + "_text_only"
-
-                future = executor.submit(_process_single_worker, task)
-                future_to_task[future] = (dataset_name, model_id, model_name, i, time.time())
-
-            # Process completed tasks with timeout
-            completed = 0
-            successful = 0
-            failed = 0
-            skipped = 0
-            timed_out = 0
-
-            try:
-                for future in as_completed(future_to_task, timeout=task_timeout + 60):  # Small buffer for cleanup
-                    dataset_name, model_id, model_name, task_idx, submit_time = future_to_task[future]
-                    completed += 1
-                    task_time = time.time() - submit_time
-
-                    try:
-                        # Get result with timeout
-                        dataset_name_result, model_id_result, success, result_count, error_msg, was_skipped = (
-                            future.result(timeout=10)
-                        )
-
-                        if success:
-                            if was_skipped:
-                                skipped += 1
-                                print(
-                                    f"↷ [{completed}/{len(tasks)}] Skipped: {dataset_name_result}/{model_id_result} (name: {model_name}) - eval exists ({task_time:.1f}s)"
-                                )
-                            else:
-                                successful += 1
-                                print(
-                                    f"✓ [{completed}/{len(tasks)}] Completed: {dataset_name_result}/{model_id_result} (name: {model_name}) ({task_time:.1f}s)"
-                                )
-                                print(f"  → {result_count} evaluation metrics")
-                        else:
-                            failed += 1
-                            print(
-                                f"✗ [{completed}/{len(tasks)}] Failed: {dataset_name_result}/{model_id_result} (name: {model_name}) ({task_time:.1f}s)"
-                            )
-                            print(f"  Error: {error_msg}")
-
-                        # Progress update
-                        if completed % max(1, len(tasks) // 10) == 0:  # Show progress every 10%
-                            elapsed = time.time() - start_time
-                            if completed > 0:
-                                eta = (elapsed / completed) * (len(tasks) - completed)
-                                print(f"  → Progress: {completed / len(tasks) * 100:.1f}% | ETA: {eta:.1f}s")
-
-                    except concurrent.futures.TimeoutError:
-                        timed_out += 1
-                        failed += 1
-                        print(
-                            f"✗ [{completed}/{len(tasks)}] Timeout: {dataset_name}/{model_id} (name: {model_name}) (>{task_timeout}s)"
-                        )
-                        future.cancel()  # Try to cancel the future
-
-                    except Exception as e:
-                        failed += 1
-                        print(
-                            f"✗ [{completed}/{len(tasks)}] Exception in worker for {dataset_name}/{model_id} (name: {model_name}): {e}"
-                        )
-
-            except concurrent.futures.TimeoutError:
-                print("⚠ Warning: Overall timeout reached. Some tasks may not have completed.")
-                # Handle any remaining uncompleted tasks
-                remaining_tasks = len(tasks) - completed
-                if remaining_tasks > 0:
-                    print(f"⚠ Warning: {remaining_tasks} tasks did not complete")
-                    failed += remaining_tasks
-
-            total_time = time.time() - start_time
-            print("\n=== Summary ===")
-            print(f"Total tasks: {len(tasks)}")
-            print(f"Successful: {successful}")
-            print(f"Skipped: {skipped}")
-            print(f"Failed: {failed}")
-            if timed_out > 0:
-                print(f"Timed out: {timed_out}")
-            print(f"Success rate: {successful / len(tasks) * 100:.1f}%")
-            print(f"Total time: {total_time:.1f}s")
-            print(f"Average time per task: {total_time / len(tasks):.1f}s")
-            print(f"Speedup vs sequential: ~{max_workers}x (theoretical)")
-
-    print("\n=== Evaluation suite completed ===")
diff --git a/src/mmcontext/eval/label_similarity.py b/src/mmcontext/eval/label_similarity.py
index 5689b79..757c697 100644
--- a/src/mmcontext/eval/label_similarity.py
+++ b/src/mmcontext/eval/label_similarity.py
@@ -850,7 +850,7 @@ def _compute_topk_accuracies(
 
         n_cells, n_labels = similarity_matrix.shape
         if n_labels == 0 or n_cells == 0:
-            return {k: 0.0 for k in ks}
+            return dict.fromkeys(ks, 0.0)
 
         # Sanitize ks: positive, <= n_labels, unique and sorted
         valid_ks = sorted({int(k) for k in ks if int(k) > 0})
diff --git a/src/mmcontext/eval/umap_plotter.py b/src/mmcontext/eval/umap_plotter.py
deleted file mode 100644
index 81d0940..0000000
--- a/src/mmcontext/eval/umap_plotter.py
+++ /dev/null
@@ -1,172 +0,0 @@
-# mmcontext/eval/umap_plotter.py
-from __future__ import annotations
-
-import logging
-from pathlib import Path
-
-import anndata as ad
-import numpy as np
-
-from mmcontext.eval.base import BaseEvaluator, EvalResult
-from mmcontext.eval.registry import register
-from mmcontext.eval.utils import LabelKind
-from mmcontext.pl.plotting import plot_umap
-from mmcontext.utils import consolidate_low_frequency_categories
-
-logger = logging.getLogger(__name__)
-
-
-@register
-class UmapPlotter(BaseEvaluator):
-    """
-    Evaluator that generates UMAP plots for all embeddings in adata.obsm.
-
-    This evaluator doesn't compute any metrics but generates UMAP visualizations
-    for all embeddings found in the AnnData object, colored by the specified labels.
-    """
-
-    name = "UmapPlotter"
-    requires_pair = False  # Can work with just one embedding
-    produces_plot = True  # This evaluator is all about plotting
-
-    def compute(
-        self,
-        emb1: np.ndarray,
-        *,
-        labels: np.ndarray,
-        adata: ad.AnnData,
-        label_kind: LabelKind,
-        label_key: str,
-        **kw,
-    ) -> EvalResult:
-        """Compute method - just returns empty results since this evaluator is for plotting only."""
-        return EvalResult()
-
-    def plot(
-        self,
-        emb1: np.ndarray,
-        out_dir: Path,
-        *,
-        emb2: np.ndarray = None,
-        labels: np.ndarray,
-        adata: ad.AnnData,
-        label_kind: LabelKind,
-        label_key: str,
-        save_format: str = "pdf",
-        figsize: tuple = (8, 8),
-        dpi: int = 300,
-        point_size: int = None,
-        font_size: int = 12,
-        font_style: str = "normal",
-        font_weight: str = "normal",
-        legend_fontsize: int = 10,
-        legend_loc: str = "right margin",
-        axis_label_size: int = 12,
-        axis_tick_size: int = 10,
-        frameon: bool = False,
-        **kw,
-    ) -> None:
-        """
-        Generate UMAP plots for all embeddings in adata.obsm.
-
-        Parameters
-        ----------
-        emb1 : np.ndarray
-            Primary embedding (will be stored as 'data_rep_embedding' in adata.obsm)
-        out_dir : Path
-            Directory to save plots
-        emb2 : np.ndarray, optional
-            Secondary embedding (will be stored as 'caption_embedding' in adata.obsm)
-        labels : np.ndarray
-            Labels for coloring the UMAP
-        adata : ad.AnnData
-            AnnData object to store embeddings and use for plotting
-        label_kind : LabelKind
-            Type of label (bio or batch)
-        label_key : str
-            Key for the labels in adata.obs
-        save_format : str, optional
-            Format to save plots (default: "pdf")
-        figsize : tuple, optional
-            Figure size (default: (8, 8))
-        dpi : int, optional
-            DPI for saved plots (default: 300)
-        point_size : int, optional
-            Size of points in UMAP (default: None, uses scanpy default)
-        font_size : int, optional
-            General font size (default: 12)
-        font_style : str, optional
-            Font style (default: "normal")
-        font_weight : str, optional
-            Font weight (default: "normal")
-        legend_fontsize : int, optional
-            Font size for legends (default: 10)
-        legend_loc : str, optional
-            Legend location (default: "right margin")
-        axis_label_size : int, optional
-            Size of axis labels (default: 12)
-        axis_tick_size : int, optional
-            Size of axis tick labels (default: 10)
-        frameon : bool, optional
-            Whether to show frame around plots (default: False)
-        **kw
-            Additional keyword arguments
-        """
-        # Make a copy of adata to avoid modifying the original
-        adata_full = adata.copy()
-        adata_copy = consolidate_low_frequency_categories(adata_full, label_key, 10)
-        # Store the embeddings in obsm
-        adata_copy.obsm["data_rep_embedding"] = emb1
-        if emb2 is not None:
-            adata_copy.obsm["caption_embedding"] = emb2
-
-        # Create output directory
-        out_dir.mkdir(parents=True, exist_ok=True)
-
-        # Get all embedding keys from obsm
-        embedding_keys = list(adata_copy.obsm.keys())
-
-        logger.info(f"Generating UMAP plots for {len(embedding_keys)} embeddings: {embedding_keys}")
-
-        # Loop over all embeddings and create UMAP plots
-        for emb_key in embedding_keys:
-            try:
-                logger.info(f"Generating UMAP for embedding: {emb_key}")
-
-                # Generate a clean title
-                title = f"UMAP - {emb_key} ({label_key})"
-
-                # Create nametag for filename
-                nametag = f"{emb_key}_{label_key}_{label_kind.value}"
-
-                # Call the plot_umap function with all parameters
-                plot_umap(
-                    adata=adata_copy,
-                    embedding_key=emb_key,
-                    color_key=label_key,
-                    title=title,
-                    save_plot=True,
-                    save_dir=str(out_dir),
-                    nametag=nametag,
-                    figsize=figsize,
-                    dpi=dpi,
-                    point_size=point_size,
-                    save_format=save_format,
-                    frameon=frameon,
-                    legend_fontsize=legend_fontsize,
-                    font_weight=font_weight,
-                    legend_loc=legend_loc,
-                    # Additional parameters that we'll pass through
-                    font_size=font_size,
-                    font_style=font_style,
-                    axis_label_size=axis_label_size,
-                    axis_tick_size=axis_tick_size,
-                )
-
-                logger.info(f"✓ UMAP plot saved for {emb_key}")
-
-            except Exception as e:
-                logger.error(f"✗ Failed to generate UMAP for {emb_key}: {e}")
-                continue
-
-        logger.info(f"✓ All UMAP plots completed for {label_key} ({label_kind.value})")
diff --git a/src/mmcontext/eval/utils.py b/src/mmcontext/eval/utils.py
index 52f6499..c02842e 100644
--- a/src/mmcontext/eval/utils.py
+++ b/src/mmcontext/eval/utils.py
@@ -2,7 +2,7 @@
 import os
 import random
 from dataclasses import dataclass
-from enum import Enum
+from enum import StrEnum
 
 import numpy as np
 import pandas as pd
@@ -11,7 +11,7 @@
 logger = logging.getLogger(__name__)
 
 
-class LabelKind(str, Enum):
+class LabelKind(StrEnum):
     """Name for the kind of label."""
 
     BIO = "bio"
diff --git a/src/mmcontext/file_utils.py b/src/mmcontext/file_utils.py
index b7a66a8..a9b40f1 100644
--- a/src/mmcontext/file_utils.py
+++ b/src/mmcontext/file_utils.py
@@ -1,4 +1,5 @@
 # tests/utils.py
+import hashlib
 import json
 import logging
 import os
@@ -27,6 +28,16 @@
 logger = logging.getLogger(__name__)
 
 
+def url_to_cache_name(url: str) -> str:
+    """Deterministic short name for a URL, used as a cache file/directory name.
+
+    Mirrors the scheme used by :func:`mmcontext.io.prepare_store.prepare_vector_store`
+    so that the same remote store maps to a stable, collision-free cache key
+    regardless of its position in a links list.
+    """
+    return hashlib.sha256(url.encode()).hexdigest()[:16]
+
+
 def remove_corrupted_null_arrays(zarr_path: str | Path) -> list[str]:
     """Remove corrupted null arrays from zarr store.
 
@@ -64,7 +75,7 @@ def find_and_remove(group, path=""):
         for key in list(group.keys()):
             current_path = f"{path}/{key}" if path else key
             item = group[key]
-            if isinstance(item, zarr.core.Array):
+            if isinstance(item, zarr.Array):
                 attrs = dict(item.attrs)
                 if attrs.get("encoding-type") == "null":
                     # Remove the corrupted array from filesystem
@@ -75,7 +86,7 @@ def find_and_remove(group, path=""):
                         else:
                             full_path.unlink()
                         removed.append(current_path)
-            elif isinstance(item, zarr.hierarchy.Group):
+            elif isinstance(item, zarr.Group):
                 find_and_remove(item, current_path)
 
     find_and_remove(z)
@@ -98,6 +109,8 @@ def load_test_adata_from_hf_dataset(
     link_column: str = "adata_link",
     zenodo_token: str | None = None,
     force_drafts: bool = False,
+    lean: bool = False,
+    obs_columns: list[str] | None = None,
 ) -> tuple[ad.AnnData, Path]:
     """
     Download the unique AnnData chunk referenced in *test_split* and load it.
@@ -119,11 +132,26 @@ def load_test_adata_from_hf_dataset(
         If False, automatically converts Zenodo draft links to published links.
         This is a quick fix for datasets created with draft links that were
         published remotely afterwards. Set to True to use actual draft links.
+    lean : bool, default False
+        If True, do **not** materialise the full AnnData. Instead read only the
+        ``obs`` table (subset to *obs_columns* if given) and the single
+        ``obsm[layer_key]`` layer, returning a minimal AnnData with no ``X`` and
+        ``n_vars == 0``. This keeps memory low for large zarr chunks (50k–100k
+        cells) where only the obs index, label columns, and one embedding layer
+        are needed downstream. Requires ``axis == "obs"``; for ``axis == "var"``
+        the function falls back to a full read with a warning. For ``.h5ad``
+        files the lean path reads via ``backed="r"``; a full read remains the
+        accepted fallback.
+    obs_columns : list[str] | None, default None
+        When *lean* is True, restrict the returned ``obs`` to these columns
+        (plus the index). Columns absent from the store are ignored. Ignored
+        when *lean* is False.
 
     Returns
     -------
     adata : AnnData
-        The loaded AnnData (in memory, not backed).
+        The loaded AnnData. Fully in-memory unless *lean* is True, in which case
+        it contains only ``obs`` and ``obsm[layer_key]``.
     local_path : pathlib.Path
         Path to the downloaded store (`*.zarr`, `*.h5ad`, or `*.zip`).
 
@@ -131,6 +159,15 @@ def load_test_adata_from_hf_dataset(
     ------
     ValueError
         If the split references multiple different files.
+
+    Notes
+    -----
+    The downloaded store is cached under *save_dir* using a deterministic hash
+    of its URL (``<sha256[:16]>``), matching
+    :func:`mmcontext.io.prepare_store.prepare_vector_store`. This makes the
+    cache key stable and collision-free across datasets sharing the same
+    *save_dir* (the legacy ``chunk_0`` naming caused different datasets to
+    overwrite/false-hit each other).
     """
     # 1) ensure there is exactly ONE unique link
     links, _ = collect_unique_links({"test": test_split}, split="test", link_column=link_column)
@@ -146,10 +183,33 @@ def load_test_adata_from_hf_dataset(
         overwrite=False,
         zenodo_token=zenodo_token,
         force_drafts=force_drafts,
+        name_by_url_hash=True,
     )
     local_path = next(iter(local_map.values()))
 
     # 3) open store
+    if lean and axis != "obs":
+        logger.warning("lean=True is only supported for axis='obs'; falling back to a full read.")
+        lean = False
+
+    if lean:
+        # Lean path: read only obs (optionally column-subset) and the single
+        # obsm[layer_key] layer. No X, no other obsm layers — keeps memory low.
+        from mmcontext.io._zarr_read import read_obs_and_obsm
+
+        if local_path.suffix not in (".h5ad", ".zip") and local_path.is_dir():
+            # Fix corrupted null arrays before reading (zarr directory only)
+            remove_corrupted_null_arrays(local_path)
+
+        obs_df, obsm = read_obs_and_obsm(local_path, obsm_key=layer_key, obs_columns=obs_columns)
+        adata = ad.AnnData(obs=obs_df)
+        if layer_key is not None:
+            if obsm is None:
+                raise KeyError(f"Layer '{layer_key}' could not be read from obsm of file {local_path}")
+            adata.obsm[layer_key] = obsm
+        logger.info("Loaded lean test AnnData with %d cells (obs + obsm['%s'] only)", adata.n_obs, layer_key)
+        return adata, local_path
+
     if local_path.suffix == ".h5ad":
         adata = ad.read_h5ad(local_path)
     elif local_path.suffix == ".zip":
@@ -470,6 +530,7 @@ def download_and_extract_links(
     zenodo_token: str | None = None,
     chunk_size: int = 8 * (1 << 20),  # 8MB chunks for better performance
     force_drafts: bool = False,
+    name_by_url_hash: bool = False,
 ) -> dict[str, Path]:
     """
     Download every share-link or handle local paths.  If it is a ZIP (Nextcloud folder-download) either
@@ -502,6 +563,14 @@ def download_and_extract_links(
         by removing ``/draft`` from URLs. This is a quick fix to allow using
         datasets created with draft links that were published remotely afterwards.
         Set to True if you actually need to download from a draft record.
+    name_by_url_hash : bool, default False
+        If False, downloaded files are named ``chunk_<idx>.<ext>`` using the
+        link's position in *links*. This is **not** stable across calls — two
+        different links each passed as the sole item collide on ``chunk_0``.
+        If True, files are named by a deterministic hash of the URL
+        (``<sha256[:16]>.<ext>``), matching the caching scheme of
+        :func:`mmcontext.io.prepare_store.prepare_vector_store` and giving a
+        stable, collision-free cache key per URL.
 
     Returns
     -------
@@ -546,13 +615,17 @@ def _is_local_path(link: str) -> bool:
             continue
 
         # For URLs, proceed with download logic
+        # Base name for cached outputs: a stable per-URL hash, or the legacy
+        # position-based ``chunk_<idx>`` name.
+        base = url_to_cache_name(link) if name_by_url_hash else f"chunk_{idx}"
+
         # ------------------------- check if present----
-        already = target_dir / f"chunk_{idx}.zip"
+        already = target_dir / f"{base}.zip"
         if already.exists() and not overwrite:
             out_map[link] = already
             continue  # ← skip download altogether
 
-        already = target_dir / f"chunk_{idx}.zarr"
+        already = target_dir / f"{base}.zarr"
         if already.exists() and not overwrite:
             out_map[link] = already
             continue
@@ -689,20 +762,20 @@ def _is_local_path(link: str) -> bool:
         m4 = tmp.read_bytes()[:4]
         if m4 == PK_MAGIC:  # ⇒ ZIP archive
             if extract:
-                out_path = target_dir / f"chunk_{idx}.zarr"
+                out_path = target_dir / f"{base}.zarr"
                 if overwrite and out_path.exists():
                     shutil.rmtree(out_path)
                 with zipfile.ZipFile(tmp) as zf:
                     zf.extractall(out_path)
                 tmp.unlink(missing_ok=True)
             else:  # keep the zip
-                out_path = target_dir / f"chunk_{idx}.zip"
+                out_path = target_dir / f"{base}.zip"
                 if overwrite and out_path.exists():
                     out_path.unlink()
                 shutil.move(tmp, out_path)
         else:  # raw .h5ad etc.
             suffix = Path(link).suffix or ".bin"
-            out_path = target_dir / f"chunk_{idx}{suffix}"
+            out_path = target_dir / f"{base}{suffix}"
             if overwrite and out_path.exists():
                 out_path.unlink()
             shutil.move(tmp, out_path)
@@ -753,37 +826,51 @@ def build_embedding_df(
     -----
     *Works with AnnData ≥ 0.10 and Zarr ≥ 2.16.*
     """
+    from mmcontext.io._zarr_read import read_obs_and_obsm
+
     rows: list[dict] = []
 
     for p in link2path.values():
         p = Path(p)
         logger.info("Reading %s", p)
 
-        # -------------------------------------------------- open store ----
-        if p.suffix == ".h5ad":  # plain HDF5
-            adata = ad.read_h5ad(p, backed="r")
-        elif p.suffix == ".zip":  # zipped zarr
-            zroot = zarr.open(p, mode="r")
-            adata = ad.read_zarr(zroot)
-        elif p.suffix == ".zarr" or p.is_dir():  # dir zarr
-            # Fix corrupted null arrays before reading
-            remove_corrupted_null_arrays(p)
-            adata = ad.read_zarr(str(p))
+        # -------------------------------------------------- read embedding ----
+        # For the obs axis we only need obsm[layer_key] + the obs index, so use
+        # the lean reader: it pulls just that one layer (zarr: never touches X or
+        # other obsm layers; h5ad: backed read). The var axis still falls back to
+        # a full read (the lean reader is obs-only).
+        if axis == "obs":
+            if p.suffix not in (".h5ad", ".zip") and p.is_dir():
+                # Fix corrupted null arrays before reading (zarr directory only)
+                remove_corrupted_null_arrays(p)
+            try:
+                obs_df, emb_matrix = read_obs_and_obsm(p, obsm_key=layer_key)
+            except FileNotFoundError:
+                logger.warning("Skip unsupported file %s", p)
+                continue
+            if emb_matrix is None:
+                raise KeyError(f"Layer key '{layer_key}' not found in obsm for file {p}.")
+            tokens = obs_df.index.to_numpy()
         else:
-            logger.warning("Skip unsupported file %s", p)
-            continue
-
-        # Get embedding matrix with helpful error message if key is missing
-        try:
-            emb_matrix = adata.obsm[layer_key] if axis == "obs" else adata.varm[layer_key]
-        except KeyError as e:
-            available_keys = list(adata.obsm.keys()) if axis == "obs" else list(adata.varm.keys())
-            axis_name = "obsm" if axis == "obs" else "varm"
-            raise KeyError(
-                f"Layer key '{layer_key}' not found in adata.{axis_name} for file {p}.\n"
-                f"Available keys in adata.{axis_name}: {available_keys}"
-            ) from e
-        tokens = adata.obs.index.to_numpy() if axis == "obs" else adata.var.index.to_numpy()
+            if p.suffix == ".h5ad":  # plain HDF5
+                adata = ad.read_h5ad(p, backed="r")
+            elif p.suffix == ".zip":  # zipped zarr
+                zroot = zarr.open(p, mode="r")
+                adata = ad.read_zarr(zroot)
+            elif p.suffix == ".zarr" or p.is_dir():  # dir zarr
+                remove_corrupted_null_arrays(p)
+                adata = ad.read_zarr(str(p))
+            else:
+                logger.warning("Skip unsupported file %s", p)
+                continue
+            try:
+                emb_matrix = adata.varm[layer_key]
+            except KeyError as e:
+                raise KeyError(
+                    f"Layer key '{layer_key}' not found in adata.varm for file {p}.\n"
+                    f"Available keys in adata.varm: {list(adata.varm.keys())}"
+                ) from e
+            tokens = adata.var.index.to_numpy()
 
         # -------------------------------------------------- stream rows ----
         n_rows = emb_matrix.shape[0]
@@ -797,7 +884,8 @@ def build_embedding_df(
                 else:
                     vec = np.asarray(vec)
                 rows.append({"token": tokens[start + i], "embedding": vec})
-        adata.file.close() if hasattr(adata, "file") else None  # close backing
+        if axis != "obs" and hasattr(adata, "file"):
+            adata.file.close() if adata.file else None  # close backing
 
     df = pd.DataFrame(rows, columns=["token", "embedding"])
     logger.info("Built DataFrame with %d rows × %d-dim embeddings", len(df), len(df["embedding"].iloc[0]))
diff --git a/src/mmcontext/hub_utils.py b/src/mmcontext/hub_utils.py
index 513c1f6..0a73b6b 100644
--- a/src/mmcontext/hub_utils.py
+++ b/src/mmcontext/hub_utils.py
@@ -1,17 +1,187 @@
 """Utilities for uploading MMContext models to Hugging Face Hub with custom model cards."""
 
 import logging
+import os
 import shutil
 import tempfile
 from pathlib import Path
 from typing import Optional, Union
 
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, HfFolder
 from sentence_transformers import SentenceTransformer
 
 logger = logging.getLogger(__name__)
 
 
+# ---------------------------------------------------------------------------
+# Authentication / naming helpers (shared by training scripts)
+# ---------------------------------------------------------------------------
+def get_hf_token() -> str | None:
+    """Return a HuggingFace token from env vars or cached CLI credentials.
+
+    Checks ``HF_TOKEN_UPLOAD`` then ``HF_TOKEN``, then falls back to the token
+    cached by ``huggingface-cli login``.
+
+    Returns
+    -------
+    str or None
+        The token if found, otherwise ``None``.
+    """
+    token = os.getenv("HF_TOKEN_UPLOAD") or os.getenv("HF_TOKEN")
+    if token:
+        return token
+    try:
+        return HfFolder.get_token()
+    except Exception:
+        return None
+
+
+def get_hf_username() -> str:
+    """Auto-detect the HuggingFace username from a token or CLI login.
+
+    Returns
+    -------
+    str
+        The username of the authenticated account.
+
+    Raises
+    ------
+    ValueError
+        If no valid authentication is found or the token is invalid.
+    """
+    from huggingface_hub.errors import LocalTokenNotFoundError
+
+    api = HfApi()
+    token = get_hf_token()
+    try:
+        user_info = api.whoami(token=token)
+        if isinstance(user_info, dict):
+            username = user_info.get("name")
+            if username:
+                logger.info("Auto-detected HuggingFace username: %s", username)
+                return username
+            raise ValueError(f"Could not extract username from HuggingFace API response. Response: {user_info}")
+        raise ValueError(f"Unexpected response type from whoami(): {type(user_info)}")
+    except LocalTokenNotFoundError:
+        raise ValueError(
+            "No HuggingFace authentication found. Please either:\n"
+            "1. Set HF_TOKEN or HF_TOKEN_UPLOAD environment variable, or\n"
+            "2. Run 'huggingface-cli login' to authenticate"
+        ) from None
+    except Exception as e:
+        raise ValueError(
+            f"Failed to get HuggingFace username: {e}\nPlease check your HuggingFace token is valid."
+        ) from e
+
+
+def check_model_exists(model_name: str, username: str | None = None) -> bool:
+    """Check whether ``{username}/{model_name}`` exists on the Hub.
+
+    Parameters
+    ----------
+    model_name : str
+        Repository name (without the owner prefix).
+    username : str, optional
+        Owner; auto-detected via :func:`get_hf_username` when omitted.
+
+    Returns
+    -------
+    bool
+        ``True`` if the model exists, ``False`` otherwise (including on any error).
+    """
+    try:
+        if username is None:
+            username = get_hf_username()
+        HfApi().model_info(f"{username}/{model_name}")
+        return True
+    except Exception:
+        return False
+
+
+def generate_model_name(
+    text_encoder_name: str,
+    *,
+    obsm_key: str | None = None,
+    tag: str | None = None,
+) -> str:
+    """Build a model name like ``mmcontext-{encoder}-{obsm_key}[-{tag}]``.
+
+    Parameters
+    ----------
+    text_encoder_name : str
+        HuggingFace text-encoder id (e.g. ``"NeuML/pubmedbert-base-embeddings"``).
+    obsm_key : str, optional
+        Omics embedding key (e.g. ``"X_scvi_fm"``); the leading ``"X_"`` is
+        stripped. Omitted from the name when ``None`` (text-only run).
+    tag : str, optional
+        Optional suffix tag.
+
+    Returns
+    -------
+    str
+        The generated model name.
+    """
+    name = text_encoder_name.lower()
+    if "pubmedbert" in name:
+        encoder_str = "pubmedbert"
+    elif "biobert" in name:
+        encoder_str = "biobert"
+    elif "modernbert" in name:
+        encoder_str = "biomodern" if "bioclinical" in name else "modernbert"
+    elif "qwen" in name:
+        encoder_str = "qwen"
+    else:
+        encoder_str = text_encoder_name.split("/")[-1]
+
+    if obsm_key:
+        method = obsm_key[2:] if obsm_key.startswith("X_") else obsm_key
+        model_name = f"mmcontext-{encoder_str}-{method}"
+    else:
+        model_name = f"mmcontext-{encoder_str}"
+
+    if tag and tag.strip():
+        model_name = f"{model_name}-{tag.strip()}"
+    return model_name
+
+
+def generate_unique_model_name(
+    text_encoder_name: str,
+    *,
+    obsm_key: str | None = None,
+    tag: str | None = None,
+    username: str | None = None,
+) -> str:
+    """Like :func:`generate_model_name` but append ``-vN`` to avoid Hub clashes.
+
+    Parameters
+    ----------
+    text_encoder_name, obsm_key, tag
+        Forwarded to :func:`generate_model_name`.
+    username : str, optional
+        Owner to check against; auto-detected when omitted.
+
+    Returns
+    -------
+    str
+        A model name that does not currently exist under *username*.
+    """
+    base_name = generate_model_name(text_encoder_name, obsm_key=obsm_key, tag=tag)
+    if username is None:
+        username = get_hf_username()
+    if not check_model_exists(base_name, username):
+        return base_name
+    version = 2
+    while version <= 10:
+        versioned = f"{base_name}-v{version}"
+        if not check_model_exists(versioned, username):
+            logger.info("Model '%s' exists, using '%s' instead", base_name, versioned)
+            return versioned
+        version += 1
+    import time
+
+    return f"{base_name}-{int(time.time())}"
+
+
 # Embedding method descriptions
 EMBEDDING_DESCRIPTIONS = {
     "pca": "Principal Component Analysis (PCA) dimensionality reduction selection of genes (from gs10k) (50-dimensional)",
@@ -562,6 +732,42 @@ def insert_custom_content_after_yaml(existing_readme: str, custom_content: str)
     return "\n".join(combined_lines)
 
 
+def append_model_card_section(
+    model_dir: str | Path,
+    section_md: str,
+    *,
+    heading: str = "Training Details (MMContext)",
+) -> Path:
+    """Append a Markdown section to a saved model's ``README.md``.
+
+    ``SentenceTransformer.save`` writes a ``README.md`` (model card) into the
+    model directory. When that directory is later pushed via
+    ``model.push_to_hub(..., local_model_path=model_dir)``, the file is uploaded
+    verbatim — so appending here makes custom details (e.g. hard-negative mining
+    settings) show up on the Hub.
+
+    Parameters
+    ----------
+    model_dir : str or Path
+        Directory of a saved SentenceTransformer model.
+    section_md : str
+        Markdown body for the new section (without the heading).
+    heading : str, optional
+        Section heading (rendered as an ``##`` header).
+
+    Returns
+    -------
+    Path
+        Path to the README that was written.
+    """
+    readme = Path(model_dir) / "README.md"
+    existing = readme.read_text() if readme.exists() else ""
+    block = f"\n\n## {heading}\n\n{section_md.strip()}\n"
+    readme.write_text(existing + block)
+    logger.info("Appended '%s' section to %s", heading, readme)
+    return readme
+
+
 def get_model_info_from_config(config_path: str | Path) -> dict[str, str]:
     """
     Extract model information from training configuration.
diff --git a/src/mmcontext/io/__init__.py b/src/mmcontext/io/__init__.py
new file mode 100644
index 0000000..2ff385f
--- /dev/null
+++ b/src/mmcontext/io/__init__.py
@@ -0,0 +1,11 @@
+"""mmcontext.io — Data loading and vector storage utilities.
+
+This module provides disk-backed vector storage for omics embeddings
+and utilities for building stores from AnnData zarr archives.
+"""
+
+from ._adata_subset import collect_adata_subset
+from .prepare_store import build_namespaced_vector_store, prepare_vector_store
+from .vector_store import VectorStore
+
+__all__ = ["VectorStore", "build_namespaced_vector_store", "collect_adata_subset", "prepare_vector_store"]
diff --git a/src/mmcontext/io/_adata_subset.py b/src/mmcontext/io/_adata_subset.py
new file mode 100644
index 0000000..a327c71
--- /dev/null
+++ b/src/mmcontext/io/_adata_subset.py
@@ -0,0 +1,159 @@
+"""Collect a subset of AnnData observations from multiple Zarr/H5AD chunks.
+
+Moved from ``embed.dataset_utils`` during v2.0 cleanup — it belongs with the
+I/O utilities in ``mmcontext.io``.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from pathlib import Path
+
+import anndata as ad
+import numpy as np
+
+from mmcontext.file_utils import remove_corrupted_null_arrays
+from mmcontext.io._zarr_read import _open_zarr, _read_obs_names_zarr, read_obs_and_obsm
+
+logger = logging.getLogger(__name__)
+
+
+def collect_adata_subset(
+    download_dir: str | Path | None = None,
+    sample_ids: Sequence[str] | None = None,
+    *,
+    file_paths: list[Path] | None = None,
+    join_vars: str = "outer",
+    obsm_key: str | None = None,
+    obs_columns: list[str] | None = None,
+) -> ad.AnnData:
+    """Load a subset of observations from multiple Zarr/H5AD chunks.
+
+    Gather observations that may live in **multiple** Zarr/H5AD chunks into a
+    single in-memory :class:`~anndata.AnnData`.
+
+    Parameters
+    ----------
+    download_dir : str, Path, or None
+        Directory containing ``*.zarr/`` folders and/or ``*.h5ad`` files.
+        Ignored when *file_paths* is given.
+    sample_ids : Sequence[str] or None
+        Observation indices to collect. Order is preserved in the output.
+    file_paths : list[Path] or None
+        Explicit file paths to read from. Takes precedence over *download_dir*.
+    join_vars : {"outer", "inner"}, default "outer"
+        How to merge variables (genes) when concatenating chunks.
+    obsm_key : str or None
+        If given, use the memory-lean path: only the requested ``obsm`` layer
+        and the ``obs`` table are read per chunk. The returned AnnData has no
+        ``X`` (``n_vars == 0``). If ``None``, the full-read path is used.
+    obs_columns : list[str] or None
+        With *obsm_key*, restrict collected ``obs`` to these columns (plus
+        the index). Ignored when *obsm_key* is ``None``.
+
+    Returns
+    -------
+    AnnData
+        Combined object with ``obs_names`` matching the order of *sample_ids*.
+    """
+    if download_dir is None and file_paths is None:
+        raise ValueError("Either download_dir or file_paths must be provided")
+
+    if sample_ids is None:
+        raise ValueError("sample_ids must be provided")
+
+    if file_paths is not None:
+        paths_to_scan = file_paths
+    else:
+        download_dir = Path(download_dir).expanduser()
+        if not download_dir.exists():
+            raise FileNotFoundError(download_dir)
+        paths_to_scan = list(download_dir.glob("*.zarr")) + list(download_dir.glob("*.h5ad"))
+
+    remaining = set(sample_ids)
+    collected: list[ad.AnnData] = []
+
+    # Scan zarr stores first
+    for path in paths_to_scan:
+        if not remaining:
+            break
+        path = Path(path)
+        if path.suffix == ".zarr" or (path.is_dir() and path.name.endswith(".zarr")):
+            logger.debug("Scanning Zarr store %s", path)
+            remove_corrupted_null_arrays(path)
+
+            if obsm_key is not None:
+                obs_names = _read_obs_names_zarr(_open_zarr(path))
+                name_to_idx = {n: i for i, n in enumerate(obs_names)}
+                wanted = remaining.intersection(obs_names)
+                if wanted:
+                    logger.info("  ↳ %d obs found in %s", len(wanted), path.name)
+                    wanted_list = [n for n in obs_names if n in wanted]
+                    rows = np.array([name_to_idx[n] for n in wanted_list])
+                    obs_df, obsm = read_obs_and_obsm(path, obsm_key=obsm_key, rows=rows, obs_columns=obs_columns)
+                    sub = ad.AnnData(obs=obs_df)
+                    sub.obsm[obsm_key] = obsm
+                    collected.append(sub)
+                    remaining.difference_update(wanted)
+            else:
+                view = ad.read_zarr(path)
+                wanted = remaining.intersection(view.obs_names)
+                if wanted:
+                    logger.info("  ↳ %d obs found in %s", len(wanted), path.name)
+                    collected.append(view[list(wanted)].copy())
+                    remaining.difference_update(wanted)
+                if hasattr(view, "file") and view.file:
+                    view.file.close()
+
+    # Scan *.h5ad files
+    for path in paths_to_scan:
+        if not remaining:
+            break
+        path = Path(path)
+        if path.suffix == ".h5ad":
+            logger.debug("Scanning H5AD file %s", path.name)
+            view = ad.read_h5ad(path, backed="r")
+            obs_names = list(view.obs_names)
+            wanted = remaining.intersection(obs_names)
+            if hasattr(view, "file") and view.file:
+                view.file.close()
+
+            if wanted:
+                logger.info("  ↳ %d obs found in %s", len(wanted), path.name)
+                if obsm_key is not None:
+                    name_to_idx = {n: i for i, n in enumerate(obs_names)}
+                    wanted_list = [n for n in obs_names if n in wanted]
+                    rows = np.array([name_to_idx[n] for n in wanted_list])
+                    obs_df, obsm = read_obs_and_obsm(path, obsm_key=obsm_key, rows=rows, obs_columns=obs_columns)
+                    sub = ad.AnnData(obs=obs_df)
+                    sub.obsm[obsm_key] = obsm
+                    collected.append(sub)
+                else:
+                    full = ad.read_h5ad(path)
+                    collected.append(full[list(wanted)].copy())
+                remaining.difference_update(wanted)
+
+    if remaining:
+        logger.warning(
+            "collect_adata_subset | %d sample_ids were *not* found in available files: %s",
+            len(remaining),
+            sorted(remaining)[:10],
+        )
+
+    if not collected:
+        raise ValueError("None of the requested sample_ids were found in any chunk.")
+
+    logger.info("Concatenating %d partial AnnData objects …", len(collected))
+    adata = ad.concat(
+        collected,
+        axis=0,
+        join=join_vars,
+        merge="same",
+        label=None,
+        index_unique=None,
+    )
+
+    # Restore original order
+    adata = adata[sample_ids].copy()
+    return adata
diff --git a/src/mmcontext/io/_zarr_read.py b/src/mmcontext/io/_zarr_read.py
new file mode 100644
index 0000000..3be2ea7
--- /dev/null
+++ b/src/mmcontext/io/_zarr_read.py
@@ -0,0 +1,239 @@
+"""Low-level zarr/h5ad readers that load only what a caller needs.
+
+The functions here open an AnnData-formatted store *without* materialising the
+full :class:`~anndata.AnnData` object. For zarr stores this means reading only
+the ``obs`` group (cheap: ``N`` rows × a few columns), the requested ``obsm``
+layer, and — optionally — only a subset of rows of that layer. The large ``X``
+matrix and any unrequested ``obsm`` layers are never touched.
+
+This module deliberately keeps its top-level imports light (``numpy`` and
+``zarr`` only). ``anndata`` is imported lazily inside the functions that need it
+so that ``import mmcontext`` stays cheap.
+
+The zarr-handle helpers (:func:`_open_zarr`, :func:`_read_obs_names_zarr`,
+:func:`_get_obsm_zarr_array`) are the single source of truth reused by
+:mod:`mmcontext.io.prepare_store`.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import numpy as np
+import zarr
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Zarr handle helpers — open a store and read obs_names / obsm lazily
+# ---------------------------------------------------------------------------
+
+
+def _open_zarr(path: Path) -> zarr.Group:
+    """Open a zarr store from a directory or zip file (read-only)."""
+    path = Path(path)
+    if path.suffix == ".zip" and path.is_file():
+        return zarr.open(str(path), mode="r")
+    elif path.is_dir():
+        # Could be a zarr directory, or a directory containing a single zarr.
+        # v3 stores carry ``zarr.json``; v2 stores carry ``.zgroup``/``.zattrs``.
+        if (path / "zarr.json").exists() or (path / ".zgroup").exists() or (path / ".zattrs").exists():
+            return zarr.open(str(path), mode="r")
+        # Check for a single subdirectory that is the actual zarr
+        subdirs = [p for p in path.iterdir() if p.is_dir()]
+        if len(subdirs) == 1:
+            return zarr.open(str(subdirs[0]), mode="r")
+        raise ValueError(f"Cannot determine zarr root in {path}")
+    else:
+        raise FileNotFoundError(f"Path does not exist: {path}")
+
+
+def _read_obs_names_zarr(root: zarr.Group) -> list[str]:
+    """Read observation names from a zarr-backed AnnData store.
+
+    AnnData zarr layout stores the obs index under ``obs/_index`` (recent
+    anndata) or ``obs/index`` (older convention). In both cases the values
+    may be stored directly or via ``__categories``.
+    """
+    obs = root["obs"]
+
+    # Determine the index key name
+    if "_index" in obs.attrs:
+        idx_key = obs.attrs["_index"]
+    elif "__index_level_0__" in obs:
+        idx_key = "__index_level_0__"
+    else:
+        idx_key = "_index"
+
+    idx_elem = obs[idx_key]
+
+    # anndata ≥ 0.11 with zarr 3.x stores array elements as zarr Groups
+    # (encoded format with encoding-type metadata). Use read_elem which handles
+    # both old (plain zarr Array) and new (encoded zarr Group) layouts.
+    if isinstance(idx_elem, zarr.Group):
+        index_data = _read_elem(idx_elem)
+        return [str(v) for v in index_data]
+
+    # Legacy path: plain zarr Array
+    # Handle categorical encoding (older anndata versions)
+    if hasattr(idx_elem, "attrs") and "categories" in idx_elem.attrs:
+        cat_path = idx_elem.attrs["categories"]
+        codes = idx_elem[:]
+        categories = root[cat_path][:]
+        return [categories[c].decode() if isinstance(categories[c], bytes) else str(categories[c]) for c in codes]
+
+    values = idx_elem[:]
+    return [v.decode() if isinstance(v, bytes) else str(v) for v in values]
+
+
+def _get_obsm_zarr_array(root: zarr.Group, obsm_key: str) -> zarr.Array:
+    """Return a zarr Array handle for an obsm layer (no data read yet).
+
+    Callers can use ``.get_orthogonal_selection()`` or ``[rows]`` to read
+    only the rows they need, so the full matrix is never in memory.
+    """
+    if "obsm" not in root:
+        raise KeyError(f"No 'obsm' group in zarr store. Available: {list(root.keys())}")
+
+    obsm = root["obsm"]
+    if obsm_key not in obsm:
+        raise KeyError(f"Key '{obsm_key}' not in obsm. Available: {list(obsm.keys())}")
+
+    return obsm[obsm_key]
+
+
+# ---------------------------------------------------------------------------
+# Row-selection helper
+# ---------------------------------------------------------------------------
+
+
+def _select_rows_float32(arr: zarr.Array, rows: np.ndarray | None) -> np.ndarray:
+    """Read ``rows`` of a zarr array as ``float32``, or all rows if None.
+
+    Row indices are sorted before the read for sequential disk access and the
+    result is unsorted back into the caller's order, mirroring the strategy in
+    :func:`mmcontext.io.prepare_store.prepare_vector_store`.
+    """
+    if rows is None:
+        return arr[:].astype(np.float32, copy=False)
+
+    rows = np.asarray(rows)
+    if rows.size == 0:
+        return np.empty((0, arr.shape[1]), dtype=np.float32)
+
+    sort_order = np.argsort(rows)
+    sorted_rows = rows[sort_order]
+    # Use oindex for orthogonal row selection — compatible with zarr 2.x and 3.x
+    # (zarr 3.x deprecated get_orthogonal_selection in favour of arr.oindex).
+    selected = arr.oindex[sorted_rows, :].astype(np.float32, copy=False)
+    # Unsort back to the caller's original row order
+    unsort = np.argsort(sort_order)
+    return selected[unsort]
+
+
+# ---------------------------------------------------------------------------
+# High-level lazy reader
+# ---------------------------------------------------------------------------
+
+
+def read_obs_and_obsm(
+    path: str | Path,
+    *,
+    obsm_key: str | None = None,
+    rows: np.ndarray | list[int] | None = None,
+    obs_columns: list[str] | None = None,
+):
+    """Read only the obs table and (optionally) one obsm layer from a store.
+
+    Parameters
+    ----------
+    path
+        Path to a zarr store (directory or ``.zip``) or an ``.h5ad`` file.
+    obsm_key
+        Name of the single ``obsm`` layer to read. If ``None`` no obsm data
+        is read and the returned array is ``None``.
+    rows
+        Integer row indices to read (in the order they should be returned). If
+        ``None`` all rows are read. The returned obs table is sliced to the
+        same rows/order.
+    obs_columns
+        Subset of ``obs`` columns to keep. Columns not present in the store are
+        silently ignored. If ``None`` all obs columns are kept. The obs index
+        is always retained.
+
+    Returns
+    -------
+    obs : pandas.DataFrame
+        The (optionally row- and column-subset) obs table, index preserved.
+    obsm : numpy.ndarray | None
+        ``float32`` array of shape ``(len(rows), D)`` for the requested layer,
+        or ``None`` if *obsm_key* was ``None``.
+
+    Notes
+    -----
+    For zarr stores the full ``X`` matrix and any unrequested ``obsm`` layers
+    are never read. ``obs`` is read in full (it is small) and then subset. For
+    ``.h5ad`` files the store is opened with ``backed="r"`` so only the needed
+    elements are pulled into memory.
+    """
+    path = Path(path)
+    is_h5ad = path.suffix == ".h5ad"
+    if is_h5ad:
+        return _read_obs_and_obsm_h5ad(path, obsm_key=obsm_key, rows=rows, obs_columns=obs_columns)
+    return _read_obs_and_obsm_zarr(path, obsm_key=obsm_key, rows=rows, obs_columns=obs_columns)
+
+
+def _read_elem(elem):
+    """Import-light wrapper around anndata's ``read_elem``."""
+    try:
+        from anndata.io import read_elem
+    except ImportError:  # pragma: no cover - older anndata layout
+        from anndata.experimental import read_elem
+    return read_elem(elem)
+
+
+def _subset_obs(obs, rows, obs_columns):
+    """Apply column then row subsetting to an obs DataFrame."""
+    if obs_columns is not None:
+        keep = [c for c in obs_columns if c in obs.columns]
+        obs = obs[keep]
+    if rows is not None:
+        obs = obs.iloc[np.asarray(rows)]
+    return obs
+
+
+def _read_obs_and_obsm_zarr(path, *, obsm_key, rows, obs_columns):
+    root = _open_zarr(path)
+    obs = _read_elem(root["obs"])  # full obs table — cheap (N × few cols)
+    obs = _subset_obs(obs, rows, obs_columns)
+
+    obsm = None
+    if obsm_key is not None:
+        arr = _get_obsm_zarr_array(root, obsm_key)
+        obsm = _select_rows_float32(arr, None if rows is None else np.asarray(rows))
+    return obs, obsm
+
+
+def _read_obs_and_obsm_h5ad(path, *, obsm_key, rows, obs_columns):
+    import anndata as ad
+
+    # backed="r" keeps X on disk; obs/obsm are read into memory (small relative
+    # to X). We then slice to the requested rows/columns.
+    adata = ad.read_h5ad(path, backed="r")
+    try:
+        obs = adata.obs
+        obs = _subset_obs(obs, rows, obs_columns)
+        obsm = None
+        if obsm_key is not None:
+            if obsm_key not in adata.obsm:
+                raise KeyError(f"Key '{obsm_key}' not in obsm. Available: {list(adata.obsm.keys())}")
+            layer = np.asarray(adata.obsm[obsm_key])
+            if rows is not None:
+                layer = layer[np.asarray(rows)]
+            obsm = layer.astype(np.float32, copy=False)
+    finally:
+        if getattr(adata, "file", None) is not None:
+            adata.file.close()
+    return obs, obsm
diff --git a/src/mmcontext/io/prepare_store.py b/src/mmcontext/io/prepare_store.py
new file mode 100644
index 0000000..68a6146
--- /dev/null
+++ b/src/mmcontext/io/prepare_store.py
@@ -0,0 +1,334 @@
+"""Build a VectorStore from adata_link URLs in a HuggingFace dataset.
+
+The key design constraint is **memory efficiency**: zarr files are opened
+directly (without loading the full AnnData) so that only the requested
+``obsm`` layer and the obs index are read.  This makes it feasible to
+prepare stores from large datasets on a 16 GB laptop.
+
+Typical usage::
+
+    from datasets import load_dataset
+    from mmcontext.io import prepare_vector_store
+
+    ds = load_dataset("jo-mengr/cxg_schaefer_tiny", split="train")
+    store = prepare_vector_store(
+        ds,
+        obsm_key="X_scvi_fm",
+        output_path="data/store.mmap",
+        cache_dir="data/zarr_cache",
+    )
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import tempfile
+from collections.abc import Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse
+from zipfile import ZipFile, is_zipfile
+
+import numpy as np
+import requests
+from tqdm.auto import tqdm
+
+from ._zarr_read import _get_obsm_zarr_array, _open_zarr, _read_obs_names_zarr
+from .vector_store import VectorStore
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+# Zarr handle helpers (`_open_zarr`, `_read_obs_names_zarr`,
+# `_get_obsm_zarr_array`) live in `._zarr_read` and are imported above so there
+# is a single implementation shared with the loaders in `file_utils` and
+# `embed.dataset_utils`.
+
+
+# ---------------------------------------------------------------------------
+# Download helper
+# ---------------------------------------------------------------------------
+
+
+def _url_to_cache_name(url: str) -> str:
+    """Deterministic short name for a URL, used as cache directory name."""
+    return hashlib.sha256(url.encode()).hexdigest()[:16]
+
+
+def _download_zarr(url: str, cache_dir: Path) -> Path:
+    """Download a zarr archive (zip or directory) with caching.
+
+    Returns the local path to the zarr store (directory or zip).
+    """
+    cache_name = _url_to_cache_name(url)
+
+    # Check for already-extracted zarr directory
+    zarr_dir = cache_dir / f"{cache_name}.zarr"
+    if zarr_dir.is_dir():
+        logger.debug("Cache hit (zarr dir): %s", zarr_dir)
+        return zarr_dir
+
+    # Check for already-downloaded zip
+    zip_path = cache_dir / f"{cache_name}.zip"
+    if zip_path.is_file():
+        logger.debug("Cache hit (zip): %s", zip_path)
+        return _maybe_extract_zip(zip_path, zarr_dir)
+
+    # Download
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    # Zenodo draft→published URL fixup: published records use the non-draft
+    # API path, but datasets often store the draft URL from upload time.
+    download_url = url
+    parsed = urlparse(download_url)
+    host = (parsed.hostname or "").lower()
+    is_zenodo_host = host == "zenodo.org" or host.endswith(".zenodo.org")
+    if is_zenodo_host and "/draft/" in download_url:
+        download_url = download_url.replace("/draft/", "/")
+        logger.info("Converted Zenodo draft URL to published: %s", download_url)
+
+    logger.info("Downloading %s → %s", download_url, zip_path)
+
+    headers = {}
+    if is_zenodo_host:
+        headers["User-Agent"] = "Mozilla/5.0 (compatible; mmcontext/1.0)"
+
+    with requests.get(download_url, stream=True, timeout=(30, 600), headers=headers) as r:
+        r.raise_for_status()
+        total = int(r.headers.get("content-length", 0))
+        with (
+            open(zip_path, "wb") as f,
+            tqdm(total=total, unit="B", unit_scale=True, desc="Downloading", leave=False) as pbar,
+        ):
+            for chunk in r.iter_content(chunk_size=8 * 1024 * 1024):
+                f.write(chunk)
+                pbar.update(len(chunk))
+
+    return _maybe_extract_zip(zip_path, zarr_dir)
+
+
+def _maybe_extract_zip(zip_path: Path, zarr_dir: Path) -> Path:
+    """If *zip_path* is a zip, extract to *zarr_dir* and return it.
+
+    If it's already a plain zarr directory (mis-named), just rename.
+    """
+    if is_zipfile(zip_path):
+        logger.info("Extracting %s → %s", zip_path, zarr_dir)
+        with ZipFile(zip_path) as zf:
+            zf.extractall(zarr_dir)
+        # Optionally delete zip to save space
+        # zip_path.unlink()
+        return zarr_dir
+
+    # Not a zip — the "download" was already a zarr directory or similar
+    zip_path.rename(zarr_dir)
+    return zarr_dir
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def prepare_vector_store(
+    dataset: Dataset,
+    *,
+    obsm_key: str,
+    output_path: str | Path,
+    cache_dir: str | Path | None = None,
+    sample_id_column: str = "sample_idx",
+    adata_link_column: str = "adata_link",
+    overwrite: bool = False,
+) -> VectorStore:
+    """Build a VectorStore from adata_link URLs in a HuggingFace dataset.
+
+    For each unique ``adata_link`` in the dataset, the zarr file is downloaded
+    (with local caching), and only the ``obsm[obsm_key]`` layer plus the obs
+    index are read — the full AnnData is **never** loaded into memory.
+
+    The resulting VectorStore maps ``sample_id_column`` values to their
+    embedding vectors and is written as a memory-mapped file at *output_path*.
+
+    Parameters
+    ----------
+    dataset : datasets.Dataset
+        HuggingFace dataset with at least *sample_id_column* and
+        *adata_link_column* columns.
+    obsm_key : str
+        Key in ``adata.obsm`` to extract (e.g. ``"X_scvi_fm"``,
+        ``"X_pca"``).
+    output_path : str or Path
+        Where to write the ``.mmap`` file (a sidecar ``.index.json``
+        is created alongside it).
+    cache_dir : str, Path, or None
+        Directory for caching downloaded zarr files. Defaults to a
+        ``mmcontext_zarr_cache`` folder inside the system temp directory.
+    sample_id_column : str
+        Column in *dataset* containing sample IDs (must match
+        ``adata.obs_names``).
+    adata_link_column : str
+        Column in *dataset* containing URLs to zarr-backed AnnData files.
+    overwrite : bool
+        If False and *output_path* already exists, load and return the
+        existing store.
+
+    Returns
+    -------
+    VectorStore
+        Memory-mapped store ready to attach to an MMContextModule.
+    """
+    output_path = Path(output_path)
+    if output_path.exists() and not overwrite:
+        logger.info("VectorStore already exists at %s, loading.", output_path)
+        return VectorStore.load(output_path)
+
+    if cache_dir is None:
+        cache_dir = Path(tempfile.gettempdir()) / "mmcontext_zarr_cache"
+    cache_dir = Path(cache_dir)
+
+    # Collect unique adata links and which sample IDs come from each
+    link_to_sample_ids: dict[str, list[str]] = {}
+    for row in dataset:
+        link = row[adata_link_column]
+        sid = str(row[sample_id_column])
+        link_to_sample_ids.setdefault(link, []).append(sid)
+
+    logger.info(
+        "Preparing VectorStore: %d samples across %d adata files, obsm_key=%r",
+        sum(len(v) for v in link_to_sample_ids.values()),
+        len(link_to_sample_ids),
+        obsm_key,
+    )
+
+    # Process each adata file and collect (id, vector) pairs
+    all_ids: list[str] = []
+    all_vectors: list[np.ndarray] = []
+
+    for link, sample_ids in tqdm(link_to_sample_ids.items(), desc="Processing adata files"):
+        # Download / cache
+        if link.startswith(("http://", "https://")):
+            local_path = _download_zarr(link, cache_dir)
+        else:
+            local_path = Path(link)
+
+        # Open zarr and read only what we need
+        root = _open_zarr(local_path)
+        obs_names = _read_obs_names_zarr(root)
+        obs_name_to_idx = {name: i for i, name in enumerate(obs_names)}
+
+        # Resolve which rows we need
+        needed_rows = []
+        needed_ids = []
+        for sid in sample_ids:
+            if sid not in obs_name_to_idx:
+                raise KeyError(
+                    f"Sample ID '{sid}' not found in obs_names of {link}. First 5 obs_names: {obs_names[:5]}"
+                )
+            needed_rows.append(obs_name_to_idx[sid])
+            needed_ids.append(sid)
+
+        # Read only the needed rows from obsm — zarr chunks are paged on
+        # demand so we never materialise the full (N_total, D) matrix.
+        # Sort row indices for sequential disk access, then unsort.
+        sort_order = np.argsort(needed_rows)
+        sorted_rows = np.array(needed_rows)[sort_order]
+
+        obsm_array = _get_obsm_zarr_array(root, obsm_key)
+        selected = obsm_array.get_orthogonal_selection((sorted_rows, slice(None))).astype(
+            np.float32
+        )  # (len(needed_rows), D)
+
+        # Unsort back to original order
+        unsort = np.argsort(sort_order)
+        selected = selected[unsort]
+
+        all_ids.extend(needed_ids)
+        all_vectors.append(selected)
+
+    # Combine and build VectorStore
+    matrix = np.concatenate(all_vectors, axis=0)  # (N_total, D)
+    logger.info("Building VectorStore: %d vectors, dim=%d", matrix.shape[0], matrix.shape[1])
+
+    store = VectorStore.from_numpy(matrix, all_ids, path=output_path)
+    logger.info("VectorStore saved to %s", output_path)
+    return store
+
+
+def build_namespaced_vector_store(
+    per_dataset: dict[str, tuple[VectorStore, Sequence[str]]],
+    *,
+    output_path: str | Path,
+    namespace_sep: str = ":",
+) -> VectorStore:
+    """Merge several per-dataset stores into one with namespaced ids.
+
+    Multiple omics datasets typically number their samples independently
+    (``sample_idx`` "0", "1", … restarts in each dataset), so concatenating
+    their stores directly would collide. A single mmcontext pipeline holds
+    exactly one :class:`VectorStore`, so the per-dataset stores must be merged.
+    This helper re-keys every vector to ``{dataset_name}{sep}{sample_idx}`` and
+    writes one combined store.
+
+    The same namespacing must be applied to the dataset's ``sample_idx`` column
+    *before* building anchors, so the resulting ``omics:{name}{sep}{idx}``
+    anchors resolve against these keys. Build each per-dataset store on the
+    **original** ids first (so the lookup against ``adata.obs_names`` in
+    :func:`prepare_vector_store` succeeds), then pass it here.
+
+    Parameters
+    ----------
+    per_dataset : dict[str, tuple[VectorStore, Sequence[str]]]
+        Maps a dataset name to ``(store, original_sample_ids)``. The ids are
+        looked up in *store* via :meth:`VectorStore.batch_lookup`; duplicates
+        are de-duplicated while preserving order.
+    output_path : str or Path
+        Where to write the merged ``.mmap`` file (a sidecar ``.index.json`` is
+        written alongside it).
+    namespace_sep : str, default ":"
+        Separator between the dataset name and the original id.
+
+    Returns
+    -------
+    VectorStore
+        The merged store, keyed by ``{name}{sep}{orig_id}``.
+
+    Raises
+    ------
+    ValueError
+        If *per_dataset* is empty or the stores disagree on dimensionality.
+    """
+    if not per_dataset:
+        raise ValueError("per_dataset is empty: at least one store is required.")
+
+    all_ids: list[str] = []
+    all_vectors: list[np.ndarray] = []
+    dim: int | None = None
+
+    for name, (store, orig_ids) in per_dataset.items():
+        # Unique, order-preserving so rows line up with the namespaced ids.
+        unique_ids = list(dict.fromkeys(str(i) for i in orig_ids))
+        if not unique_ids:
+            continue
+        vectors = store.batch_lookup(unique_ids)  # (N, D)
+        if dim is None:
+            dim = vectors.shape[1]
+        elif vectors.shape[1] != dim:
+            raise ValueError(
+                f"Store '{name}' has dim {vectors.shape[1]} but a previous store has dim {dim}. "
+                "All omics datasets in one run must share the same obsm_key / dimensionality."
+            )
+        all_vectors.append(vectors)
+        all_ids.extend(f"{name}{namespace_sep}{sid}" for sid in unique_ids)
+
+    matrix = np.concatenate(all_vectors, axis=0)
+    logger.info(
+        "Merged %d namespaced VectorStores: %d vectors, dim=%d",
+        len(per_dataset),
+        matrix.shape[0],
+        matrix.shape[1],
+    )
+    return VectorStore.from_numpy(matrix, all_ids, path=output_path)
diff --git a/src/mmcontext/io/vector_store.py b/src/mmcontext/io/vector_store.py
new file mode 100644
index 0000000..50bec01
--- /dev/null
+++ b/src/mmcontext/io/vector_store.py
@@ -0,0 +1,403 @@
+"""VectorStore — disk-backed vector lookup using numpy memory-mapped files.
+
+Provides efficient storage and retrieval of high-dimensional embedding vectors
+(e.g., scVI latent spaces, Geneformer embeddings, gene count vectors) without
+loading the full matrix into RAM. Vectors are stored in a flat numpy memmap
+file on disk; the OS pages in only the rows that are actually accessed.
+
+Typical usage
+-------------
+**Create from AnnData:**
+
+>>> store = VectorStore.from_adata(adata, layer_key="X_scvi", axis="obs", path="cache/vectors.mmap")
+
+**Create from numpy:**
+
+>>> store = VectorStore.from_numpy(matrix, ids, path="cache/vectors.mmap")
+
+**Lookup:**
+
+>>> vec = store["cell_42"]  # single lookup → (D,)
+>>> batch = store.batch_lookup(ids)  # batch lookup  → (N, D)
+
+**Persistence:**
+
+>>> store2 = VectorStore.load("cache/vectors.mmap")  # reopen later
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+
+if TYPE_CHECKING:
+    import anndata as ad
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+class VectorStore:
+    """Disk-backed vector lookup using numpy memory-mapped files.
+
+    Stores embedding vectors in a flat ``(N, D)`` numpy memmap file on disk,
+    with a JSON sidecar file mapping string IDs to row indices. Only the rows
+    actually accessed are paged into RAM by the OS, making it feasible to work
+    with large embedding matrices (e.g., 300k × 10k) on machines with limited
+    memory.
+
+    Parameters
+    ----------
+    mmap : np.memmap
+        Memory-mapped array of shape ``(N, D)``.
+    index : dict[str, int]
+        Mapping from string IDs to row indices in the memmap.
+    path : Path
+        Path to the memmap file on disk.
+
+    Notes
+    -----
+    Do not instantiate directly — use the class methods
+    :meth:`from_numpy`, :meth:`from_dataframe`, :meth:`from_adata`,
+    :meth:`from_dict`, or :meth:`load`.
+    """
+
+    def __init__(
+        self,
+        mmap: np.memmap,
+        index: dict[str, int],
+        path: Path,
+    ) -> None:
+        self._mmap = mmap
+        self._index = index
+        self._path = path
+
+    # ------------------------------------------------------------------
+    # Construction class methods
+    # ------------------------------------------------------------------
+    @classmethod
+    def from_numpy(
+        cls,
+        matrix: np.ndarray,
+        ids: Sequence[str],
+        *,
+        path: str | Path,
+    ) -> VectorStore:
+        """Create a VectorStore from a numpy array and a list of IDs.
+
+        Parameters
+        ----------
+        matrix : np.ndarray, shape (N, D)
+            Embedding matrix. Row *i* is the vector for ``ids[i]``.
+        ids : sequence of str
+            Token/sample IDs. Must have the same length as ``matrix.shape[0]``.
+        path : str or Path
+            Where to write the memmap file. A sidecar ``<path>.index.json``
+            is written alongside it.
+
+        Returns
+        -------
+        VectorStore
+
+        Raises
+        ------
+        ValueError
+            If ``ids`` is empty, contains duplicates, or its length doesn't
+            match ``matrix.shape[0]``.
+        """
+        ids = list(ids)
+        cls._validate_ids_and_matrix(ids, matrix)
+
+        path = Path(path)
+        index = {sid: i for i, sid in enumerate(ids)}
+
+        # Write memmap
+        path.parent.mkdir(parents=True, exist_ok=True)
+        mmap = np.memmap(path, dtype=matrix.dtype, mode="w+", shape=matrix.shape)
+        mmap[:] = matrix
+        mmap.flush()
+
+        # Write sidecar index
+        cls._write_index(path, index, matrix.shape, matrix.dtype)
+
+        # Re-open as read-only
+        mmap = np.memmap(path, dtype=matrix.dtype, mode="r", shape=matrix.shape)
+
+        logger.info(
+            "Created VectorStore: %d vectors × %d dims (%s) at %s",
+            matrix.shape[0],
+            matrix.shape[1],
+            matrix.dtype,
+            path,
+        )
+        return cls(mmap=mmap, index=index, path=path)
+
+    @classmethod
+    def from_dataframe(
+        cls,
+        df: pd.DataFrame,
+        *,
+        path: str | Path,
+        id_col: str = "token",
+        embedding_col: str = "embedding",
+    ) -> VectorStore:
+        """Create a VectorStore from a pandas DataFrame.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            Must contain columns ``id_col`` (str IDs) and ``embedding_col``
+            (array-like vectors).
+        path : str or Path
+            Where to write the memmap file.
+        id_col : str
+            Column name for token/sample IDs. Default: ``"token"``.
+        embedding_col : str
+            Column name for embedding vectors. Default: ``"embedding"``.
+
+        Returns
+        -------
+        VectorStore
+        """
+        ids = df[id_col].tolist()
+        matrix = np.vstack(df[embedding_col].to_numpy())
+        return cls.from_numpy(matrix, ids, path=path)
+
+    @classmethod
+    def from_dict(
+        cls,
+        mapping: dict[str, np.ndarray],
+        *,
+        path: str | Path,
+    ) -> VectorStore:
+        """Create a VectorStore from a ``{id: vector}`` mapping.
+
+        Parameters
+        ----------
+        mapping : dict[str, np.ndarray]
+            Keys are string IDs, values are 1-D numpy arrays.
+        path : str or Path
+            Where to write the memmap file.
+
+        Returns
+        -------
+        VectorStore
+        """
+        ids = list(mapping.keys())
+        matrix = np.vstack(list(mapping.values()))
+        return cls.from_numpy(matrix, ids, path=path)
+
+    @classmethod
+    def from_adata(
+        cls,
+        adata: ad.AnnData,
+        *,
+        layer_key: str,
+        axis: Literal["obs", "var"] = "obs",
+        path: str | Path,
+    ) -> VectorStore:
+        """Create a VectorStore from an AnnData object.
+
+        Parameters
+        ----------
+        adata : anndata.AnnData
+            AnnData object containing embeddings.
+        layer_key : str
+            Key in ``.obsm`` (if *axis="obs"*) or ``.varm`` (if *axis="var"*).
+        axis : {"obs", "var"}
+            Which axis to extract embeddings from.
+        path : str or Path
+            Where to write the memmap file.
+
+        Returns
+        -------
+        VectorStore
+
+        Raises
+        ------
+        KeyError
+            If ``layer_key`` is not found in the specified axis.
+        """
+        if axis == "obs":
+            if layer_key not in adata.obsm:
+                raise KeyError(f"Key '{layer_key}' not found in adata.obsm. Available keys: {list(adata.obsm.keys())}")
+            matrix = np.asarray(adata.obsm[layer_key])
+            ids = adata.obs.index.tolist()
+        elif axis == "var":
+            if layer_key not in adata.varm:
+                raise KeyError(f"Key '{layer_key}' not found in adata.varm. Available keys: {list(adata.varm.keys())}")
+            matrix = np.asarray(adata.varm[layer_key])
+            ids = adata.var.index.tolist()
+        else:
+            raise ValueError(f"axis must be 'obs' or 'var', got '{axis}'")
+
+        # Ensure float32 if not already a float type
+        if not np.issubdtype(matrix.dtype, np.floating):
+            matrix = matrix.astype(np.float32)
+
+        return cls.from_numpy(matrix, ids, path=path)
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+    @classmethod
+    def load(cls, path: str | Path) -> VectorStore:
+        """Re-open a previously saved VectorStore from disk.
+
+        Parameters
+        ----------
+        path : str or Path
+            Path to the memmap file (the sidecar ``.index.json`` must exist
+            alongside it).
+
+        Returns
+        -------
+        VectorStore
+
+        Raises
+        ------
+        FileNotFoundError
+            If the memmap file or the index file does not exist.
+        """
+        path = Path(path)
+        index_path = Path(str(path) + ".index.json")
+
+        if not path.is_file():
+            raise FileNotFoundError(f"Memmap file not found: {path}")
+        if not index_path.is_file():
+            raise FileNotFoundError(f"Index file not found: {index_path}")
+
+        with open(index_path) as f:
+            meta = json.load(f)
+
+        index = meta["index"]
+        shape = tuple(meta["shape"])
+        dtype = np.dtype(meta["dtype"])
+
+        mmap = np.memmap(path, dtype=dtype, mode="r", shape=shape)
+
+        logger.info(
+            "Loaded VectorStore: %d vectors × %d dims (%s) from %s",
+            shape[0],
+            shape[1],
+            dtype,
+            path,
+        )
+        return cls(mmap=mmap, index=index, path=path)
+
+    # ------------------------------------------------------------------
+    # Lookup
+    # ------------------------------------------------------------------
+    def __getitem__(self, key: str) -> np.ndarray:
+        """Look up a single vector by ID.
+
+        Parameters
+        ----------
+        key : str
+            Token/sample ID.
+
+        Returns
+        -------
+        np.ndarray, shape (D,)
+            The embedding vector (read from memmap, not a copy).
+
+        Raises
+        ------
+        KeyError
+            If the ID is not in the store.
+        """
+        try:
+            idx = self._index[key]
+        except KeyError:
+            raise KeyError(f"ID '{key}' not found in VectorStore. Store contains {len(self._index)} entries.") from None
+        return np.array(self._mmap[idx])
+
+    def batch_lookup(self, ids: Sequence[str]) -> np.ndarray:
+        """Look up multiple vectors by ID.
+
+        Parameters
+        ----------
+        ids : sequence of str
+            Token/sample IDs.
+
+        Returns
+        -------
+        np.ndarray, shape (len(ids), D)
+            Stacked embedding vectors.
+
+        Raises
+        ------
+        KeyError
+            If any ID is not in the store.
+        """
+        indices = []
+        for sid in ids:
+            try:
+                indices.append(self._index[sid])
+            except KeyError:
+                raise KeyError(
+                    f"ID '{sid}' not found in VectorStore. Store contains {len(self._index)} entries."
+                ) from None
+        return np.array(self._mmap[indices])
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+    @property
+    def dim(self) -> int:
+        """Embedding dimensionality (D)."""
+        return self._mmap.shape[1]
+
+    @property
+    def dtype(self) -> np.dtype:
+        """Data type of stored vectors."""
+        return self._mmap.dtype
+
+    def __len__(self) -> int:
+        """Number of stored vectors."""
+        return len(self._index)
+
+    def __contains__(self, key: str) -> bool:
+        """Check if an ID is in the store."""
+        return key in self._index
+
+    def __repr__(self) -> str:
+        return f"VectorStore(n={len(self)}, dim={self.dim}, dtype={self.dtype}, path='{self._path}')"
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _validate_ids_and_matrix(ids: list[str], matrix: np.ndarray) -> None:
+        """Validate that ids and matrix are consistent."""
+        if len(ids) == 0:
+            raise ValueError("Empty ID list: at least one vector is required.")
+        if matrix.ndim != 2:
+            raise ValueError(f"Expected 2-D matrix, got {matrix.ndim}-D array with shape {matrix.shape}.")
+        if len(ids) != matrix.shape[0]:
+            raise ValueError(f"Length mismatch: {len(ids)} IDs but matrix has {matrix.shape[0]} rows.")
+        if len(set(ids)) != len(ids):
+            duplicates = [x for x in ids if ids.count(x) > 1]
+            raise ValueError(f"Duplicate IDs found: {sorted(set(duplicates))[:10]}. All IDs must be unique.")
+
+    @staticmethod
+    def _write_index(
+        path: Path,
+        index: dict[str, int],
+        shape: tuple[int, ...],
+        dtype: np.dtype,
+    ) -> None:
+        """Write the sidecar JSON index file."""
+        index_path = Path(str(path) + ".index.json")
+        meta = {
+            "shape": list(shape),
+            "dtype": str(dtype),
+            "index": index,
+        }
+        with open(index_path, "w") as f:
+            json.dump(meta, f)
diff --git a/src/mmcontext/models/__init__.py b/src/mmcontext/models/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/mmcontext/modules/__init__.py b/src/mmcontext/modules/__init__.py
new file mode 100644
index 0000000..eb7c0ef
--- /dev/null
+++ b/src/mmcontext/modules/__init__.py
@@ -0,0 +1,15 @@
+"""mmcontext.modules — Sentence-transformers pipeline modules.
+
+This package provides the modules that form the sentence-transformers pipeline
+for multimodal (text + omics) encoding:
+
+- :class:`MMContextModule` — InputModule: text encoder + omics pass-through
+- :class:`AdapterModule` — modality-aware projection to shared space
+- :class:`OmicsAttentionModule` — optional self-attention for var-based models
+"""
+
+from .adapter_module import AdapterModule
+from .mmcontext_module import MMContextModule
+from .omics_attention_module import OmicsAttentionModule
+
+__all__ = ["MMContextModule", "AdapterModule", "OmicsAttentionModule"]
diff --git a/src/mmcontext/modules/adapter_module.py b/src/mmcontext/modules/adapter_module.py
new file mode 100644
index 0000000..37d1cfe
--- /dev/null
+++ b/src/mmcontext/modules/adapter_module.py
@@ -0,0 +1,349 @@
+"""AdapterModule — modality-aware projection for the ST pipeline.
+
+This module sits after :class:`MMContextModule` in a sentence-transformers
+pipeline. It reads ``modality_ids`` from the features dict and applies
+separate learned projections for text and omics tokens, mapping them into
+a shared embedding space of dimension ``D_shared``.
+
+Architecture
+------------
+Two independent projection heads (``text_proj`` and ``omics_proj``) map
+embeddings from their respective input dimensions to a common output
+dimension. Three projection modes are supported:
+
+1. **Identity** — when ``force_identity=True`` and all dims match, both
+   projections act as ``nn.Identity``.
+2. **Linear → BatchNorm** — when ``hidden_dim`` is ``None`` or ``0``.
+3. **Linear → ReLU → Linear → BatchNorm** — the default MLP mode.
+
+Pad tokens (``modality_id=2``) pass through as zeros.
+
+Features dict contract::
+
+    # Input (from MMContextModule.forward):
+    {
+        "token_embeddings": Tensor[B, L, D_text or D_omics],
+        "attention_mask": Tensor[B, L],
+        "modality_ids": Tensor[B, L],  # 0=text, 1=omics, 2=pad
+    }
+
+    # Output (after AdapterModule.forward):
+    {
+        "token_embeddings": Tensor[B, L, D_shared],  # projected
+        "attention_mask": Tensor[B, L],  # unchanged
+        "modality_ids": Tensor[B, L],  # unchanged
+    }
+
+Example
+-------
+>>> adapter = AdapterModule(text_input_dim=768, omics_input_dim=512, shared_dim=256)
+>>> features = mmcontext_module.forward(mmcontext_module.preprocess(texts))
+>>> projected = adapter(features)
+>>> projected["token_embeddings"].shape
+torch.Size([B, L, 256])
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+# Try v5.4+ import paths first, fall back to v5.0 paths
+try:
+    from sentence_transformers.base.modules import Module
+except ImportError:
+    from sentence_transformers.models import Module
+
+logger = logging.getLogger(__name__)
+
+# Modality constants (must match mmcontext_module.py)
+MODALITY_TEXT = 0
+MODALITY_OMICS = 1
+MODALITY_PAD = 2
+
+
+def _build_projection(
+    input_dim: int,
+    output_dim: int,
+    hidden_dim: int | None,
+    force_identity: bool,
+) -> nn.Module:
+    """Build a single projection head.
+
+    Parameters
+    ----------
+    input_dim : int
+        Dimensionality of the incoming features.
+    output_dim : int
+        Dimensionality of the projected output.
+    hidden_dim : int or None
+        Hidden layer size. If None, skip the hidden layer.
+    force_identity : bool
+        If True and dims match and no hidden layer, use nn.Identity.
+
+    Returns
+    -------
+    nn.Module
+        The projection network.
+    """
+    if hidden_dim is None and input_dim == output_dim and force_identity:
+        return nn.Identity()
+    elif hidden_dim is None:
+        return nn.Sequential(
+            nn.Linear(input_dim, output_dim),
+            nn.LayerNorm(output_dim),
+        )
+    else:
+        return nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, output_dim),
+            nn.LayerNorm(output_dim),
+        )
+
+
+class AdapterModule(Module):
+    """Modality-aware projection module for sentence-transformers pipelines.
+
+    Maintains separate projection heads for text and omics tokens,
+    dispatching based on ``modality_ids`` in the features dict.
+
+    Parameters
+    ----------
+    text_input_dim : int
+        Dimensionality of text token embeddings (from the text encoder).
+    omics_input_dim : int
+        Dimensionality of omics token embeddings.
+    shared_dim : int
+        Output dimensionality for both projections (the shared space).
+    hidden_dim : int or None, optional
+        Hidden layer size in each MLP projection. If ``None``, uses a
+        single linear layer + BatchNorm instead. Default: 512.
+    force_identity : bool, optional
+        If True and all dims match and ``hidden_dim`` is None, both
+        projections act as identity. Default: False.
+    """
+
+    config_keys: list[str] = [
+        "text_input_dim",
+        "omics_input_dim",
+        "shared_dim",
+        "hidden_dim",
+        "force_identity",
+    ]
+    config_file_name: str = "adapter_module_config.json"
+
+    def __init__(
+        self,
+        text_input_dim: int,
+        omics_input_dim: int,
+        shared_dim: int,
+        hidden_dim: int | None = 512,
+        force_identity: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.text_input_dim = text_input_dim
+        self.omics_input_dim = omics_input_dim
+        self.shared_dim = shared_dim
+        self.hidden_dim = hidden_dim if hidden_dim else None
+        self.force_identity = force_identity
+
+        # Build independent projection heads
+        self.text_proj = _build_projection(text_input_dim, shared_dim, self.hidden_dim, force_identity)
+        self.omics_proj = _build_projection(omics_input_dim, shared_dim, self.hidden_dim, force_identity)
+
+    # ------------------------------------------------------------------
+    # Forward (Module abstract method)
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        features: dict[str, torch.Tensor | Any],
+        **kwargs,
+    ) -> dict[str, torch.Tensor | Any]:
+        """Project token embeddings using modality-specific heads.
+
+        Parameters
+        ----------
+        features : dict
+            Must contain ``token_embeddings`` (B, L, D), ``attention_mask``
+            (B, L), and ``modality_ids`` (B, L).
+
+        Returns
+        -------
+        dict
+            Updated features with ``token_embeddings`` projected to
+            ``(B, L, shared_dim)``.
+        """
+        token_embeddings = features["token_embeddings"]
+        modality_ids = features["modality_ids"]
+
+        B, L, D = token_embeddings.shape
+        device = token_embeddings.device
+
+        # Allocate output
+        output = torch.zeros(B, L, self.shared_dim, device=device, dtype=token_embeddings.dtype)
+
+        # Masks for each modality
+        text_mask = modality_ids == MODALITY_TEXT
+        omics_mask = modality_ids == MODALITY_OMICS
+        # pad (modality_id=2) stays zero — no action needed
+
+        # Project text tokens
+        if text_mask.any():
+            text_tokens = token_embeddings[text_mask]  # (N_text, D)
+            output[text_mask] = self._apply_projection(self.text_proj, text_tokens)
+
+        # Project omics tokens
+        if omics_mask.any():
+            omics_tokens = token_embeddings[omics_mask]  # (N_omics, D)
+            output[omics_mask] = self._apply_projection(self.omics_proj, omics_tokens)
+
+        features["token_embeddings"] = output
+        return features
+
+    def _apply_projection(self, proj: nn.Module, tokens: torch.Tensor) -> torch.Tensor:
+        """Apply a projection head, handling BatchNorm's 2D requirement.
+
+        BatchNorm1d expects (N, C) input. Since we gather tokens from
+        potentially scattered positions, they are already (N, D) — no
+        reshape needed.
+        """
+        return proj(tokens)
+
+    # ------------------------------------------------------------------
+    # Freezing
+    # ------------------------------------------------------------------
+    def freeze_text_proj(self) -> None:
+        """Freeze all parameters in the text projection head."""
+        for param in self.text_proj.parameters():
+            param.requires_grad = False
+        logger.info("Froze text projection head")
+
+    def unfreeze_text_proj(self) -> None:
+        """Unfreeze all parameters in the text projection head."""
+        for param in self.text_proj.parameters():
+            param.requires_grad = True
+        logger.info("Unfroze text projection head")
+
+    def freeze_omics_proj(self) -> None:
+        """Freeze all parameters in the omics projection head."""
+        for param in self.omics_proj.parameters():
+            param.requires_grad = False
+        logger.info("Froze omics projection head")
+
+    def unfreeze_omics_proj(self) -> None:
+        """Unfreeze all parameters in the omics projection head."""
+        for param in self.omics_proj.parameters():
+            param.requires_grad = True
+        logger.info("Unfroze omics projection head")
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+    def get_sentence_embedding_dimension(self) -> int:
+        """Return the shared output dimensionality."""
+        return self.shared_dim
+
+    # ------------------------------------------------------------------
+    # Save / Load
+    # ------------------------------------------------------------------
+    def save(
+        self,
+        output_path: str,
+        *args,
+        safe_serialization: bool = True,
+        **kwargs,
+    ) -> None:
+        """Save config and weights to disk.
+
+        Parameters
+        ----------
+        output_path : str
+            Directory where files will be written.
+        safe_serialization : bool
+            If True, use safetensors format for weights.
+        """
+        from safetensors.torch import save_model as save_safetensors_model
+
+        output_path = str(output_path)
+        os.makedirs(output_path, exist_ok=True)
+
+        self.save_config(output_path)
+
+        # Save weights
+        if safe_serialization:
+            save_safetensors_model(self, os.path.join(output_path, "model.safetensors"))
+        else:
+            torch.save(self.state_dict(), os.path.join(output_path, "pytorch_model.bin"))
+
+        logger.info("Saved AdapterModule to %s", output_path)
+
+    @classmethod
+    def load(
+        cls,
+        model_name_or_path: str,
+        subfolder: str = "",
+        token: bool | str | None = None,
+        cache_folder: str | None = None,
+        revision: str | None = None,
+        local_files_only: bool = False,
+        **kwargs,
+    ) -> AdapterModule:
+        """Load a saved AdapterModule from disk.
+
+        Parameters
+        ----------
+        model_name_or_path : str
+            Path to directory containing saved module files.
+
+        Returns
+        -------
+        AdapterModule
+        """
+        from safetensors.torch import load_model as load_safetensors_model
+
+        config = cls.load_config(
+            model_name_or_path,
+            subfolder=subfolder,
+            config_filename=cls.config_file_name,
+            token=token,
+            cache_folder=cache_folder,
+            revision=revision,
+            local_files_only=local_files_only,
+        )
+
+        module = cls(**config)
+
+        # Load weights
+        load_path = model_name_or_path
+        if subfolder:
+            load_path = os.path.join(model_name_or_path, subfolder)
+
+        safetensors_path = os.path.join(load_path, "model.safetensors")
+        bin_path = os.path.join(load_path, "pytorch_model.bin")
+
+        if os.path.isfile(safetensors_path):
+            load_safetensors_model(module, safetensors_path)
+        elif os.path.isfile(bin_path):
+            module.load_state_dict(torch.load(bin_path, map_location=torch.device("cpu")))
+        else:
+            logger.warning("No weight files found in %s — module uses random init.", load_path)
+
+        logger.info("Loaded AdapterModule from %s", model_name_or_path)
+        return module
+
+    # ------------------------------------------------------------------
+    # Repr
+    # ------------------------------------------------------------------
+    def __repr__(self) -> str:
+        return (
+            f"AdapterModule("
+            f"text={self.text_input_dim}→{self.shared_dim}, "
+            f"omics={self.omics_input_dim}→{self.shared_dim}, "
+            f"hidden={self.hidden_dim})"
+        )
diff --git a/src/mmcontext/modules/mmcontext_module.py b/src/mmcontext/modules/mmcontext_module.py
new file mode 100644
index 0000000..7a2f04c
--- /dev/null
+++ b/src/mmcontext/modules/mmcontext_module.py
@@ -0,0 +1,567 @@
+"""MMContextModule — InputModule for multimodal text + omics encoding.
+
+This module serves as the first module in a sentence-transformers (>=5.4)
+pipeline, handling both text preprocessing (via AutoTokenizer + AutoModel)
+and omics vector pass-through (via VectorStore or direct input). It produces
+a unified features dict that downstream modules (AdapterModule, Pooling)
+consume.
+
+Architecture
+------------
+The module operates in two modes depending on the input modality:
+
+**Text mode** (``modality="text"``):
+    Input strings are tokenized with AutoTokenizer, then forwarded through
+    AutoModel to produce contextual token embeddings.
+
+**Omics mode** (``modality="omics"``):
+    Omics vectors are either looked up from an attached :class:`VectorStore`
+    (using prefixed string IDs) or provided directly as numpy arrays (via dict
+    inputs). These vectors are stored as ``input_values`` and passed through to
+    ``token_embeddings`` in the forward output without any
+    learned transformation — the downstream AdapterModule handles projection.
+
+Features dict contract (after ``forward()``)::
+
+    {
+        "token_embeddings": Tensor[B, L, D],  # per-token representations
+        "attention_mask": Tensor[B, L],  # 1 = real, 0 = pad
+        "modality_ids": Tensor[B, L],  # 0 = text, 1 = omics
+    }
+
+Example
+-------
+>>> from mmcontext.modules import MMContextModule
+>>> module = MMContextModule("pubmedbert-base")
+>>> features = module.preprocess(["A cell with high EGFR expression."])
+>>> result = module.forward(features)
+>>> result["token_embeddings"].shape
+torch.Size([1, 8, 32])
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import numpy as np
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+# Try v5.4+ import paths first, fall back to v5.0 paths
+try:
+    from sentence_transformers.base.modules import InputModule
+except ImportError:
+    from sentence_transformers.models import InputModule
+
+from mmcontext.io import VectorStore
+
+logger = logging.getLogger(__name__)
+
+# Modality constants used in modality_ids tensor
+MODALITY_TEXT = 0
+MODALITY_OMICS = 1
+MODALITY_PAD = 2
+
+
+class MMContextModule(InputModule):
+    """Multimodal InputModule for text + omics encoding.
+
+    Extends :class:`sentence_transformers.InputModule` (v5.4+) to support
+    both text and continuous omics vectors in a single pipeline.
+
+    Parameters
+    ----------
+    model_name_or_path : str
+        Name or path for the text encoder (passed to ``AutoModel.from_pretrained``
+        and ``AutoTokenizer.from_pretrained``).
+    max_seq_length : int, optional
+        Maximum sequence length for text tokenization. Default: 512.
+    omics_prefix : str, optional
+        String prefix that marks an input as an omics sample ID to be resolved
+        via the attached :class:`VectorStore`. Default: ``"omics:"``.
+    tokenizer_args : dict, optional
+        Extra keyword arguments forwarded to ``AutoTokenizer.from_pretrained``.
+    model_args : dict, optional
+        Extra keyword arguments forwarded to ``AutoModel.from_pretrained``.
+    """
+
+    config_keys: list[str] = [
+        "model_name_or_path",
+        "max_seq_length",
+        "omics_prefix",
+    ]
+    config_file_name: str = "mmcontext_module_config.json"
+    save_in_root: bool = True
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        max_seq_length: int = 512,
+        omics_prefix: str = "omics:",
+        tokenizer_args: dict | None = None,
+        model_args: dict | None = None,
+    ) -> None:
+        super().__init__()
+
+        self.model_name_or_path = model_name_or_path
+        self._max_seq_length = max_seq_length
+        self.omics_prefix = omics_prefix
+
+        # Text encoder
+        model_args = model_args or {}
+        self.auto_model = AutoModel.from_pretrained(model_name_or_path, **model_args)
+
+        # Tokenizer (stored as self.tokenizer for InputModule compatibility)
+        tokenizer_args = tokenizer_args or {}
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, **tokenizer_args)
+
+        # VectorStore for omics ID resolution (not saved with model)
+        self._vector_store: VectorStore | None = None
+
+    # ------------------------------------------------------------------
+    # Modality metadata (v5.4+ API)
+    # ------------------------------------------------------------------
+    @property
+    def modalities(self) -> list[str]:
+        """Modalities this module can process.
+
+        Returns ``["text", "omics"]``. Note that ``"omics"`` is a custom
+        modality not in the standard ST set (text, image, audio, video).
+        """
+        return ["text", "omics"]
+
+    # ------------------------------------------------------------------
+    # Sequence length
+    # ------------------------------------------------------------------
+    @property
+    def max_seq_length(self) -> int:
+        """Maximum sequence length for text tokenization."""
+        return self._max_seq_length
+
+    @max_seq_length.setter
+    def max_seq_length(self, value: int) -> None:
+        self._max_seq_length = value
+
+    # ------------------------------------------------------------------
+    # VectorStore management
+    # ------------------------------------------------------------------
+    def set_vector_store(self, store: VectorStore) -> None:
+        """Attach a VectorStore for resolving omics sample IDs.
+
+        Parameters
+        ----------
+        store : VectorStore
+            The store mapping sample IDs to embedding vectors.
+        """
+        self._vector_store = store
+        logger.info("Attached VectorStore with %d entries, dim=%d", len(store), store.dim)
+
+    def remove_vector_store(self) -> None:
+        """Detach the current VectorStore."""
+        self._vector_store = None
+        logger.info("Detached VectorStore")
+
+    # ------------------------------------------------------------------
+    # Preprocess (v5.4+ API — replaces tokenize())
+    # ------------------------------------------------------------------
+    def preprocess(
+        self,
+        inputs: list,
+        prompt: str | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor | Any]:
+        """Preprocess inputs, routing by modality.
+
+        This is the v5.4+ entry point (replaces the deprecated ``tokenize()``).
+        The method detects the input modality and dispatches accordingly.
+
+        Parameters
+        ----------
+        inputs : list[str | dict]
+            Inputs to process. Three formats are supported:
+
+            1. **Plain strings** — preprocessed as text via AutoTokenizer.
+            2. **Prefixed strings** (e.g., ``"omics:cell_42"``) — the suffix
+               after the prefix is looked up in the attached VectorStore.
+            3. **Dicts** with ``"omics_values"`` key — the value is either a
+               single 1-D numpy array (obs case) or a list of 1-D arrays
+               (var case, variable-length gene sequences).
+
+        prompt : str, optional
+            Optional prompt to prepend to text inputs.
+        **kwargs
+            Additional keyword arguments (e.g. ``task``).
+
+        Returns
+        -------
+        dict[str, Tensor | Any]
+            Features dict. For text: ``{input_ids, attention_mask, modality}``.
+            For omics: ``{input_values, attention_mask, modality}``.
+
+        Raises
+        ------
+        ValueError
+            If the input list is empty, or if prefixed omics IDs are used
+            without an attached VectorStore.
+        KeyError
+            If an omics ID is not found in the VectorStore.
+        """
+        if len(inputs) == 0:
+            raise ValueError("Empty input list: at least one input is required.")
+
+        # Dispatch based on input type
+        if isinstance(inputs[0], dict) and "omics_values" in inputs[0]:
+            return self._preprocess_omics_direct(inputs)
+        elif isinstance(inputs[0], str) and inputs[0].startswith(self.omics_prefix):
+            return self._preprocess_omics_via_store(inputs)
+        else:
+            return self._preprocess_text(inputs, prompt=prompt)
+
+    def _preprocess_text(
+        self,
+        texts: list[str],
+        prompt: str | None = None,
+    ) -> dict[str, torch.Tensor | Any]:
+        """Preprocess plain text strings via AutoTokenizer."""
+        if prompt:
+            texts = [prompt + t for t in texts]
+
+        encoded = self.tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            max_length=self._max_seq_length,
+            return_tensors="pt",
+        )
+        encoded["modality"] = "text"
+        return encoded
+
+    def _preprocess_omics_via_store(self, texts: list[str]) -> dict[str, torch.Tensor | Any]:
+        """Resolve prefixed omics IDs through VectorStore."""
+        if self._vector_store is None:
+            raise ValueError(
+                "No VectorStore attached. Call set_vector_store() before "
+                "preprocessing omics IDs. Received input starting with "
+                f"'{self.omics_prefix}'."
+            )
+
+        # Strip prefix and look up vectors
+        prefix_len = len(self.omics_prefix)
+        sample_ids = [t[prefix_len:] for t in texts]
+
+        vectors = []
+        for sid in sample_ids:
+            vec = self._vector_store[sid]  # raises KeyError if missing
+            vectors.append(torch.from_numpy(vec).unsqueeze(0))  # (1, D)
+
+        # Stack into (B, 1, D) — each sample is a single obs-level vector
+        input_values = torch.stack(vectors, dim=0)  # (B, 1, D)
+        attention_mask = torch.ones(len(texts), 1, dtype=torch.long)
+
+        return {
+            "input_values": input_values,
+            "attention_mask": attention_mask,
+            "modality": "omics",
+        }
+
+    def _preprocess_omics_direct(self, inputs: list[dict[str, Any]]) -> dict[str, torch.Tensor | Any]:
+        """Package direct omics vectors into features dict.
+
+        Handles both obs (single vector per sample) and var (list of gene
+        vectors per sample, padded to max length in batch).
+        """
+        all_embeddings = []
+        lengths = []
+
+        for item in inputs:
+            values = item["omics_values"]
+
+            if isinstance(values, np.ndarray) and values.ndim == 1:
+                # Obs case: single vector → (1, D)
+                all_embeddings.append(torch.from_numpy(values).unsqueeze(0))
+                lengths.append(1)
+            elif isinstance(values, list | tuple):
+                # Var case: list of gene vectors → (N_genes, D)
+                gene_tensors = [torch.from_numpy(np.asarray(v)) for v in values]
+                all_embeddings.append(torch.stack(gene_tensors, dim=0))
+                lengths.append(len(values))
+            elif isinstance(values, np.ndarray) and values.ndim == 2:
+                # Var case: (N_genes, D) array
+                all_embeddings.append(torch.from_numpy(values))
+                lengths.append(values.shape[0])
+            else:
+                raise ValueError(
+                    f"Unsupported omics_values type: {type(values)}. "
+                    "Expected 1-D array, 2-D array, or list of 1-D arrays."
+                )
+
+        # Pad to max length in batch
+        max_len = max(lengths)
+        dim = all_embeddings[0].shape[-1]
+        batch_size = len(inputs)
+
+        input_values = torch.zeros(batch_size, max_len, dim)
+        attention_mask = torch.zeros(batch_size, max_len, dtype=torch.long)
+
+        for i, (emb, length) in enumerate(zip(all_embeddings, lengths, strict=False)):
+            input_values[i, :length] = emb
+            attention_mask[i, :length] = 1
+
+        return {
+            "input_values": input_values,
+            "attention_mask": attention_mask,
+            "modality": "omics",
+        }
+
+    # ------------------------------------------------------------------
+    # Forward (Module abstract method)
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        features: dict[str, torch.Tensor | Any],
+        **kwargs,
+    ) -> dict[str, torch.Tensor | Any]:
+        """Process features through the text encoder or pass omics through.
+
+        Parameters
+        ----------
+        features : dict
+            Features dict from :meth:`preprocess`. Must contain either
+            ``input_ids`` (text) or ``input_values`` (omics), plus
+            ``attention_mask`` and ``modality``.
+
+        Returns
+        -------
+        dict[str, Tensor | Any]
+            Updated features dict with ``token_embeddings``, ``attention_mask``,
+            and ``modality_ids``.
+        """
+        modality = features.get("modality", "text")
+
+        if modality == "text":
+            return self._forward_text(features)
+        else:
+            return self._forward_omics(features)
+
+    def _forward_text(self, features: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]:
+        """Run text through the transformer encoder."""
+        input_ids = features["input_ids"]
+        attention_mask = features.get("attention_mask")
+
+        # Forward through text encoder
+        model_output = self.auto_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        )
+
+        token_embeddings = model_output.last_hidden_state  # (B, L, D)
+
+        B, L = input_ids.shape
+        modality_ids = torch.full((B, L), MODALITY_TEXT, dtype=torch.long, device=input_ids.device)
+
+        features["token_embeddings"] = token_embeddings
+        features["modality_ids"] = modality_ids
+        return features
+
+    def _forward_omics(self, features: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]:
+        """Pass omics embeddings through unchanged.
+
+        Reads from ``input_values`` (set by preprocess) and writes to
+        ``token_embeddings`` (the standard key consumed by downstream modules).
+        Using ``input_values`` as the preprocess key ensures compatibility with
+        the ST training collator's ``collect_features`` suffix matching.
+        """
+        token_embeddings = features["input_values"]  # (B, L, D)
+
+        B, L = token_embeddings.shape[:2]
+        modality_ids = torch.full((B, L), MODALITY_OMICS, dtype=torch.long, device=token_embeddings.device)
+
+        features["token_embeddings"] = token_embeddings
+        features["modality_ids"] = modality_ids
+        return features
+
+    # ------------------------------------------------------------------
+    # Properties for downstream modules
+    # ------------------------------------------------------------------
+    def get_word_embedding_dimension(self) -> int:
+        """Return the hidden size of the text encoder.
+
+        This is used by downstream modules (Pooling, AdapterModule) to
+        determine the text embedding dimension.
+        """
+        return self.auto_model.config.hidden_size
+
+    # ------------------------------------------------------------------
+    # Freezing
+    # ------------------------------------------------------------------
+    def unfreeze_text_encoder(self) -> None:
+        """Unfreeze all text encoder parameters."""
+        for param in self.auto_model.parameters():
+            param.requires_grad = True
+        logger.info("Unfroze all text encoder parameters")
+
+    def _get_encoder_layers(self) -> torch.nn.ModuleList | None:
+        """Return the transformer encoder layer stack, if it can be located.
+
+        Handles both BERT-style (``encoder.layer``) and RoBERTa-style
+        (``roberta.encoder.layer``) architectures. Returns ``None`` when the
+        layer stack cannot be identified.
+        """
+        if hasattr(self.auto_model, "encoder") and hasattr(self.auto_model.encoder, "layer"):
+            return self.auto_model.encoder.layer
+        elif hasattr(self.auto_model, "roberta"):
+            return self.auto_model.roberta.encoder.layer
+        elif hasattr(self.auto_model, "bert"):
+            return self.auto_model.bert.encoder.layer
+        return None
+
+    def freeze_all_but_top_layers(self, num_trainable_layers: int) -> None:
+        """Freeze the text encoder, keeping only the top layers trainable.
+
+        This is the single entry point for text-encoder freezing. Passing
+        ``0`` freezes the entire encoder; passing ``N`` freezes everything
+        except the top ``N`` transformer layers. To make the whole encoder
+        trainable again, use :meth:`unfreeze_text_encoder`.
+
+        Freezes every parameter of the text encoder, then re-enables gradients
+        on the top ``num_trainable_layers`` transformer layers (counting from
+        the output end) plus the pooler, if present. This is the typical
+        "fine-tune only the last N layers" setup, which drastically reduces the
+        memory used by gradients and optimizer state.
+
+        Parameters
+        ----------
+        num_trainable_layers : int
+            Number of top transformer layers to keep trainable. ``0`` (or less)
+            freezes the entire encoder. Values larger than the available number
+            of layers keep all layers trainable.
+        """
+        # Freeze everything first.
+        for param in self.auto_model.parameters():
+            param.requires_grad = False
+
+        if num_trainable_layers <= 0:
+            logger.info("Froze all text encoder parameters (0 trainable top layers)")
+            return
+
+        layers = self._get_encoder_layers()
+        if layers is None:
+            logger.warning(
+                "Could not identify encoder layers; the text encoder remains fully frozen. "
+                "Use unfreeze_text_encoder() to train all parameters instead."
+            )
+            return
+
+        total = len(layers)
+        n = min(num_trainable_layers, total)
+        for layer in layers[total - n :]:
+            for param in layer.parameters():
+                param.requires_grad = True
+
+        # Keep the pooler trainable too, when one exists.
+        pooler = getattr(self.auto_model, "pooler", None)
+        if pooler is not None:
+            for param in pooler.parameters():
+                param.requires_grad = True
+
+        logger.info("Froze text encoder, kept top %d of %d layers trainable", n, total)
+
+    # ------------------------------------------------------------------
+    # Save / Load (Module abstract methods)
+    # ------------------------------------------------------------------
+    def save(
+        self,
+        output_path: str,
+        *args,
+        safe_serialization: bool = True,
+        **kwargs,
+    ) -> None:
+        """Save module config, text encoder weights, and tokenizer.
+
+        The VectorStore is intentionally NOT saved — it must be reattached
+        after loading via :meth:`set_vector_store`.
+
+        Parameters
+        ----------
+        output_path : str
+            Directory where files will be written.
+        safe_serialization : bool
+            If True, use safetensors format for weights.
+        """
+        output_path = str(output_path)
+        os.makedirs(output_path, exist_ok=True)
+
+        # Save config
+        self.save_config(output_path)
+
+        # Save text encoder
+        self.auto_model.save_pretrained(output_path, safe_serialization=safe_serialization)
+
+        # Save tokenizer
+        self.save_tokenizer(output_path)
+
+        logger.info("Saved MMContextModule to %s", output_path)
+
+    @classmethod
+    def load(
+        cls,
+        model_name_or_path: str,
+        subfolder: str = "",
+        token: bool | str | None = None,
+        cache_folder: str | None = None,
+        revision: str | None = None,
+        local_files_only: bool = False,
+        **kwargs,
+    ) -> MMContextModule:
+        """Load a saved MMContextModule from disk.
+
+        Parameters
+        ----------
+        model_name_or_path : str
+            Path to directory containing saved module files.
+        subfolder : str
+            Optional subdirectory.
+        **kwargs
+            Additional arguments (ignored, for API compatibility).
+
+        Returns
+        -------
+        MMContextModule
+        """
+        config = cls.load_config(
+            model_name_or_path,
+            subfolder=subfolder,
+            config_filename=cls.config_file_name,
+            token=token,
+            cache_folder=cache_folder,
+            revision=revision,
+            local_files_only=local_files_only,
+        )
+
+        load_path = model_name_or_path
+        if subfolder:
+            load_path = os.path.join(model_name_or_path, subfolder)
+
+        module = cls(
+            model_name_or_path=load_path,
+            max_seq_length=config.get("max_seq_length", 512),
+            omics_prefix=config.get("omics_prefix", "omics:"),
+        )
+
+        logger.info("Loaded MMContextModule from %s", load_path)
+        return module
+
+    # ------------------------------------------------------------------
+    # Repr
+    # ------------------------------------------------------------------
+    def __repr__(self) -> str:
+        return (
+            f"MMContextModule("
+            f"model={self.model_name_or_path}, "
+            f"max_seq_length={self.max_seq_length}, "
+            f"omics_prefix='{self.omics_prefix}', "
+            f"store={'attached' if self._vector_store else 'none'}"
+            f")"
+        )
diff --git a/src/mmcontext/modules/omics_attention_module.py b/src/mmcontext/modules/omics_attention_module.py
new file mode 100644
index 0000000..8120cbe
--- /dev/null
+++ b/src/mmcontext/modules/omics_attention_module.py
@@ -0,0 +1,311 @@
+"""OmicsAttentionModule — optional self-attention for omics tokens.
+
+This module sits between :class:`MMContextModule` and :class:`AdapterModule`
+in a sentence-transformers pipeline. It applies multi-head self-attention
+ONLY to omics tokens (modality_id=1), leaving text tokens unchanged.
+
+This is useful for var-level models where each sample has multiple gene
+embeddings that benefit from contextual mixing before projection into the
+shared space.
+
+Architecture
+------------
+A stack of standard ``nn.TransformerEncoderLayer`` modules with pre-LayerNorm.
+Only omics tokens are gathered, attended, and scattered back.  Text tokens
+(modality_id=0) and pad tokens (modality_id=2) pass through untouched.
+
+For obs-level inputs (L=1), self-attention is trivial — the single token
+attends only to itself, which is essentially a feedforward pass. This is
+fine: the module adds negligible overhead for obs-level data and can be
+left in the pipeline without harm.
+
+Features dict contract::
+
+    # Input (from MMContextModule.forward):
+    {
+        "token_embeddings": Tensor[B, L, D],
+        "attention_mask": Tensor[B, L],
+        "modality_ids": Tensor[B, L],  # 0=text, 1=omics, 2=pad
+    }
+
+    # Output (after OmicsAttentionModule.forward):
+    {
+        "token_embeddings": Tensor[B, L, D],  # omics tokens attended
+        "attention_mask": Tensor[B, L],  # unchanged
+        "modality_ids": Tensor[B, L],  # unchanged
+    }
+
+Example
+-------
+>>> attn = OmicsAttentionModule(input_dim=768, num_heads=8, num_layers=2)
+>>> features = mmcontext_module.forward(...)
+>>> features = attn(features)
+>>> features = adapter_module(features)
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+# Try v5.4+ import paths first, fall back to v5.0 paths
+try:
+    from sentence_transformers.base.modules import Module
+except ImportError:
+    from sentence_transformers.models import Module
+
+logger = logging.getLogger(__name__)
+
+# Modality constants (must match mmcontext_module.py / adapter_module.py)
+MODALITY_TEXT = 0
+MODALITY_OMICS = 1
+MODALITY_PAD = 2
+
+
+class OmicsAttentionModule(Module):
+    """Self-attention module applied only to omics tokens.
+
+    Maintains a stack of ``TransformerEncoderLayer`` blocks. Only omics
+    tokens (modality_id=1) are gathered, fed through the attention layers,
+    and scattered back into the output tensor. Text and pad tokens pass
+    through unchanged.
+
+    Parameters
+    ----------
+    input_dim : int
+        Dimensionality of the token embeddings (must match the upstream
+        encoder's hidden size).
+    num_heads : int, optional
+        Number of attention heads. Must evenly divide ``input_dim``.
+        Default: 4.
+    num_layers : int, optional
+        Number of stacked TransformerEncoderLayer blocks. Default: 1.
+    feedforward_dim : int or None, optional
+        Hidden size of the feedforward network inside each layer.
+        If None, defaults to ``4 * input_dim``. Default: None.
+    dropout : float, optional
+        Dropout probability for attention weights and feedforward layers.
+        Default: 0.1.
+    """
+
+    config_keys: list[str] = [
+        "input_dim",
+        "num_heads",
+        "num_layers",
+        "feedforward_dim",
+        "dropout",
+    ]
+    config_file_name: str = "omics_attention_module_config.json"
+
+    def __init__(
+        self,
+        input_dim: int,
+        num_heads: int = 4,
+        num_layers: int = 1,
+        feedforward_dim: int | None = None,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.feedforward_dim = feedforward_dim if feedforward_dim else 4 * input_dim
+        self.dropout = dropout
+
+        # Build attention stack
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=input_dim,
+            nhead=num_heads,
+            dim_feedforward=self.feedforward_dim,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,  # pre-LayerNorm for training stability
+        )
+        self.attention = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=num_layers,
+        )
+
+    # ------------------------------------------------------------------
+    # Forward (Module abstract method)
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        features: dict[str, torch.Tensor | Any],
+        **kwargs,
+    ) -> dict[str, torch.Tensor | Any]:
+        """Apply self-attention to omics tokens only.
+
+        Text tokens and pad tokens pass through unchanged. Omics tokens
+        are gathered per sample, attended, and scattered back.
+
+        Parameters
+        ----------
+        features : dict
+            Must contain ``token_embeddings`` (B, L, D), ``attention_mask``
+            (B, L), and ``modality_ids`` (B, L).
+
+        Returns
+        -------
+        dict
+            Updated features with attended omics token embeddings.
+        """
+        token_embeddings = features["token_embeddings"]
+        modality_ids = features["modality_ids"]
+        attention_mask = features["attention_mask"]
+
+        B, L, D = token_embeddings.shape
+
+        # Start with a copy of the input
+        output = token_embeddings.clone()
+
+        # Check if there are any omics tokens at all
+        omics_mask = modality_ids == MODALITY_OMICS  # (B, L)
+        if not omics_mask.any():
+            return features
+
+        # Process each sample independently — omics sequences can have
+        # different lengths across samples in the batch
+        for b in range(B):
+            omics_positions = omics_mask[b].nonzero(as_tuple=True)[0]  # positions of omics tokens
+            if len(omics_positions) == 0:
+                continue
+
+            # Gather omics tokens for this sample: (N_omics, D)
+            omics_tokens = token_embeddings[b, omics_positions, :].unsqueeze(0)  # (1, N_omics, D)
+
+            # Build attention mask for these tokens
+            # The overall attention_mask tells us which positions are real
+            omics_attn_mask = attention_mask[b, omics_positions].unsqueeze(0)  # (1, N_omics)
+
+            # Convert to the format nn.TransformerEncoder expects:
+            # src_key_padding_mask: True = ignore, False = attend
+            padding_mask = omics_attn_mask == 0  # (1, N_omics)
+
+            # Apply attention
+            attended = self.attention(
+                omics_tokens,
+                src_key_padding_mask=padding_mask,
+            )  # (1, N_omics, D)
+
+            # Scatter back into output
+            output[b, omics_positions, :] = attended.squeeze(0)
+
+        # Zero out pad positions to be safe
+        pad_mask = modality_ids == MODALITY_PAD
+        if pad_mask.any():
+            output[pad_mask] = 0.0
+
+        features = dict(features)  # shallow copy to avoid mutating input
+        features["token_embeddings"] = output
+        return features
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+    def get_sentence_embedding_dimension(self) -> int:
+        """Return the output dimensionality (same as input — attention preserves dim)."""
+        return self.input_dim
+
+    # ------------------------------------------------------------------
+    # Save / Load
+    # ------------------------------------------------------------------
+    def save(
+        self,
+        output_path: str,
+        *args,
+        safe_serialization: bool = True,
+        **kwargs,
+    ) -> None:
+        """Save config and weights to disk.
+
+        Parameters
+        ----------
+        output_path : str
+            Directory where files will be written.
+        safe_serialization : bool
+            If True, use safetensors format for weights.
+        """
+        from safetensors.torch import save_model as save_safetensors_model
+
+        output_path = str(output_path)
+        os.makedirs(output_path, exist_ok=True)
+
+        self.save_config(output_path)
+
+        if safe_serialization:
+            save_safetensors_model(self, os.path.join(output_path, "model.safetensors"))
+        else:
+            torch.save(self.state_dict(), os.path.join(output_path, "pytorch_model.bin"))
+
+        logger.info("Saved OmicsAttentionModule to %s", output_path)
+
+    @classmethod
+    def load(
+        cls,
+        model_name_or_path: str,
+        subfolder: str = "",
+        token: bool | str | None = None,
+        cache_folder: str | None = None,
+        revision: str | None = None,
+        local_files_only: bool = False,
+        **kwargs,
+    ) -> OmicsAttentionModule:
+        """Load a saved OmicsAttentionModule from disk.
+
+        Parameters
+        ----------
+        model_name_or_path : str
+            Path to directory containing saved module files.
+
+        Returns
+        -------
+        OmicsAttentionModule
+        """
+        from safetensors.torch import load_model as load_safetensors_model
+
+        config = cls.load_config(
+            model_name_or_path,
+            subfolder=subfolder,
+            config_filename=cls.config_file_name,
+            token=token,
+            cache_folder=cache_folder,
+            revision=revision,
+            local_files_only=local_files_only,
+        )
+
+        module = cls(**config)
+
+        load_path = model_name_or_path
+        if subfolder:
+            load_path = os.path.join(model_name_or_path, subfolder)
+
+        safetensors_path = os.path.join(load_path, "model.safetensors")
+        bin_path = os.path.join(load_path, "pytorch_model.bin")
+
+        if os.path.isfile(safetensors_path):
+            load_safetensors_model(module, safetensors_path)
+        elif os.path.isfile(bin_path):
+            module.load_state_dict(torch.load(bin_path, map_location=torch.device("cpu")))
+        else:
+            logger.warning("No weight files found in %s — module uses random init.", load_path)
+
+        logger.info("Loaded OmicsAttentionModule from %s", model_name_or_path)
+        return module
+
+    # ------------------------------------------------------------------
+    # Repr
+    # ------------------------------------------------------------------
+    def __repr__(self) -> str:
+        return (
+            f"OmicsAttentionModule("
+            f"dim={self.input_dim}, "
+            f"heads={self.num_heads}, "
+            f"layers={self.num_layers}, "
+            f"ff={self.feedforward_dim})"
+        )
diff --git a/src/mmcontext/training.py b/src/mmcontext/training.py
new file mode 100644
index 0000000..f60266e
--- /dev/null
+++ b/src/mmcontext/training.py
@@ -0,0 +1,680 @@
+"""Config-driven multi-dataset training helpers for the ST pipeline.
+
+This module backs ``scripts/train_config.py``. It keeps the script thin by
+encapsulating the reusable, testable parts of multi-dataset training on the
+sentence-transformers v5.4+ pipeline (``MMContextModule`` + ``AdapterModule`` +
+``Pooling`` + ``Normalize``):
+
+* :class:`TrainConfig` and :func:`load_config` — a plain-YAML config (no Hydra).
+* :func:`build_merged_vector_store` — build one namespaced
+  :class:`~mmcontext.io.VectorStore` covering every bimodal omics dataset.
+* :func:`assemble_training_data` — turn raw datasets into the
+  ``train_dataset`` / ``eval_dataset`` / ``loss`` dicts (+ evaluators) that
+  :class:`~sentence_transformers.SentenceTransformerTrainer` expects.
+
+Omics datasets are processed via :func:`~mmcontext.embed.prepare_dataset`; bio
+datasets are assumed to already be in ``anchor`` / ``positive`` / ``negative_*``
+form and are passed through unchanged.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from datasets import Dataset, DatasetDict
+    from sentence_transformers import SentenceTransformer
+
+    from mmcontext.io import VectorStore
+
+logger = logging.getLogger(__name__)
+
+NAMESPACE_SEP = ":"
+
+
+# ---------------------------------------------------------------------------
+# Config dataclasses
+# ---------------------------------------------------------------------------
+@dataclass
+class DatasetConfig:
+    """One training dataset entry (omics or bio)."""
+
+    id: str
+    name: str
+    type: str = "multiplets"  # -> get_loss / get_evaluator
+    # Omics-only fields
+    modality: str = "bimodal"  # "bimodal" | "text"
+    primary_cell_sentence: str = "cell_sentence_1"
+    positive_col: str = "positive"
+    use_hard_negatives: bool = True
+    truncate: bool = False
+    truncate_kwargs: dict | None = None
+    # Bio-only field
+    revision: str | None = None
+
+
+@dataclass
+class TextEncoderConfig:
+    """Text encoder settings."""
+
+    name: str = "NeuML/pubmedbert-base-embeddings"
+    max_seq_length: int = 512
+
+
+@dataclass
+class AdapterFreezingConfig:
+    """Per-modality adapter freeze/unfreeze schedule (maps to UnfreezeAdapterCallback)."""
+
+    freeze_text_adapter: bool = False
+    freeze_omics_adapter: bool = False
+    unfreeze_text_adapter_epoch: float | None = None
+    unfreeze_omics_adapter_epoch: float | None = None
+
+
+@dataclass
+class MiningConfig:
+    """Hard-negative mining settings (``sentence_transformers.util.mine_hard_negatives``).
+
+    When ``enabled``, each dataset's own hard negatives are dropped and
+    model-specific negatives are mined instead: the ``anchor`` column (omics ids
+    resolved through the attached VectorStore for bimodal datasets, or text for
+    the text modality) is encoded as the query and the ``positive`` texts form
+    the corpus; the closest-but-not-true positives become the mined negatives.
+
+    This only makes sense when continuing from an already-trained model (set
+    ``model:`` in the config) — mining against a freshly initialised bimodal
+    model produces near-random negatives. Mirrors ``scripts/finetune_tiny.py``.
+    """
+
+    enabled: bool = False
+    num_negatives: int = 3
+    range_min: int = 1
+    range_max: int | None = None
+    max_score: float | None = 0.95
+    relative_margin: float | None = None
+    sampling_strategy: str = "top"
+    use_faiss: bool = False
+
+
+@dataclass
+class TrainerConfig:
+    """SentenceTransformerTrainingArguments-facing settings."""
+
+    output_dir: str = "outputs/multi"
+    num_train_epochs: float = 2
+    per_device_train_batch_size: int = 64
+    per_device_eval_batch_size: int = 64
+    learning_rate: float = 2e-5
+    warmup_ratio: float = 0.1
+    bf16: bool = False
+    fp16: bool = False
+    eval_strategy: str = "steps"
+    eval_steps: int | None = 200
+    save_strategy: str = "steps"
+    save_steps: int = 2000
+    save_total_limit: int = 2
+    logging_steps: int = 50
+    max_grad_norm: float = 1.0
+    gradient_checkpointing: bool = False
+    dataloader_num_workers: int = 0
+    use_mps_device: bool = False
+    use_cosine_scheduler: bool = False
+    wandb_project: str | None = None
+    wandb_run_name: str | None = None
+
+
+@dataclass
+class TrainConfig:
+    """Top-level training config (one YAML file)."""
+
+    text_encoder: TextEncoderConfig = field(default_factory=TextEncoderConfig)
+    trainer: TrainerConfig = field(default_factory=TrainerConfig)
+
+    model: str | None = None  # HF repo / local path to continue from, or None
+    tag: str | None = None
+    push_to_hub: bool = False
+
+    obsm_key: str = "X_scvi_fm"
+    shared_dim: int = 512
+    adapter_hidden_dim: int | None = None
+
+    freeze_text_encoder: bool = False
+    unfreeze_last_n_layers: int = 0
+    unfreeze_epoch: float = 1
+    adapter_freezing: AdapterFreezingConfig | None = None
+    mining: MiningConfig | None = None
+
+    gene_special_token: str | None = None
+    delimiter: str = " "
+
+    cache_dir: str | None = None
+    adata_cache_dir: str = "data/from_nxtcloud"
+    vector_store_dir: str = "data/vector_stores"
+    force_refresh_cache: bool = False  # re-download HF datasets
+    overwrite_vector_store: bool = False  # rebuild the VectorStore (e.g. after switching obsm_key)
+
+    omics_datasets: list[DatasetConfig] = field(default_factory=list)
+    bio_datasets: list[DatasetConfig] = field(default_factory=list)
+
+
+def _coerce_scalar(value: Any, type_str: str) -> Any:
+    """Coerce a scalar to the field's primitive type from its annotation string.
+
+    Handles the common YAML gotcha where ``2e-5`` (no decimal point) is parsed
+    as a string rather than a float. Annotations are strings here because this
+    module uses ``from __future__ import annotations``.
+    """
+    if value is None:
+        return None
+    if type_str.startswith("bool"):
+        return value
+    try:
+        if "float" in type_str:
+            return float(value)
+        if "int" in type_str:
+            return int(value)
+    except (TypeError, ValueError):
+        return value
+    return value
+
+
+def _filter_kwargs(cls: type, data: dict[str, Any]) -> dict[str, Any]:
+    """Keep only keys that are fields of *cls* (warn on the rest), coercing scalars."""
+    fields = cls.__dataclass_fields__  # type: ignore[attr-defined]
+    unknown = set(data) - set(fields)
+    if unknown:
+        logger.warning("Ignoring unknown config keys for %s: %s", cls.__name__, sorted(unknown))
+    return {k: _coerce_scalar(v, str(fields[k].type)) for k, v in data.items() if k in fields}
+
+
+def config_from_dict(data: dict[str, Any]) -> TrainConfig:
+    """Build a :class:`TrainConfig` from a plain (YAML-loaded) dict."""
+    data = dict(data or {})
+
+    text_encoder = TextEncoderConfig(**_filter_kwargs(TextEncoderConfig, data.pop("text_encoder", {}) or {}))
+    trainer = TrainerConfig(**_filter_kwargs(TrainerConfig, data.pop("trainer", {}) or {}))
+
+    adapter_freezing_raw = data.pop("adapter_freezing", None)
+    adapter_freezing = (
+        AdapterFreezingConfig(**_filter_kwargs(AdapterFreezingConfig, adapter_freezing_raw))
+        if adapter_freezing_raw
+        else None
+    )
+
+    mining_raw = data.pop("mining", None)
+    mining = MiningConfig(**_filter_kwargs(MiningConfig, mining_raw)) if mining_raw else None
+
+    omics = [DatasetConfig(**_filter_kwargs(DatasetConfig, d)) for d in (data.pop("omics_datasets", None) or [])]
+    bio = [DatasetConfig(**_filter_kwargs(DatasetConfig, d)) for d in (data.pop("bio_datasets", None) or [])]
+
+    return TrainConfig(
+        text_encoder=text_encoder,
+        trainer=trainer,
+        adapter_freezing=adapter_freezing,
+        mining=mining,
+        omics_datasets=omics,
+        bio_datasets=bio,
+        **_filter_kwargs(TrainConfig, data),
+    )
+
+
+def load_config(path: str | Path) -> TrainConfig:
+    """Load and validate a YAML training config.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path to the YAML config file.
+
+    Returns
+    -------
+    TrainConfig
+    """
+    import yaml
+
+    with open(path) as f:
+        raw = yaml.safe_load(f)
+    return config_from_dict(raw)
+
+
+# ---------------------------------------------------------------------------
+# Raw dataset loading
+# ---------------------------------------------------------------------------
+def load_raw_datasets(cfg: TrainConfig) -> tuple[dict[str, DatasetDict], dict[str, DatasetDict]]:
+    """Load every omics and bio dataset referenced by *cfg* from the Hub.
+
+    Parameters
+    ----------
+    cfg : TrainConfig
+        Run config.
+
+    Returns
+    -------
+    tuple[dict[str, DatasetDict], dict[str, DatasetDict]]
+        ``(omics_raw, bio_raw)`` keyed by ``DatasetConfig.name``.
+    """
+    from datasets import load_dataset
+    from datasets.exceptions import DatasetNotFoundError
+
+    download_mode = "force_redownload" if cfg.force_refresh_cache else "reuse_dataset_if_exists"
+
+    omics_raw: dict[str, DatasetDict] = {}
+    for dcfg in cfg.omics_datasets:
+        logger.info("Loading omics dataset '%s' from %s", dcfg.name, dcfg.id)
+        omics_raw[dcfg.name] = load_dataset(dcfg.id, cache_dir=cfg.cache_dir, download_mode=download_mode)
+
+    bio_raw: dict[str, DatasetDict] = {}
+    for dcfg in cfg.bio_datasets:
+        logger.info("Loading bio dataset '%s' from %s (revision=%s)", dcfg.name, dcfg.id, dcfg.revision)
+        try:
+            bio_raw[dcfg.name] = load_dataset(
+                dcfg.id, revision=dcfg.revision, cache_dir=cfg.cache_dir, download_mode=download_mode
+            )
+        except DatasetNotFoundError as e:
+            if dcfg.revision is None:
+                raise
+            # The pinned revision (e.g. "hard_negatives") may not exist for every
+            # dataset; the main branch is expected to carry a negatives column too.
+            logger.warning(
+                "Revision %r not found for bio dataset '%s' (%s); falling back to the default branch.",
+                dcfg.revision,
+                dcfg.name,
+                e,
+            )
+            bio_raw[dcfg.name] = load_dataset(dcfg.id, cache_dir=cfg.cache_dir, download_mode=download_mode)
+
+    return omics_raw, bio_raw
+
+
+# ---------------------------------------------------------------------------
+# Vector store
+# ---------------------------------------------------------------------------
+def build_merged_vector_store(
+    cfg: TrainConfig,
+    omics_raw: dict[str, DatasetDict],
+    *,
+    overwrite: bool = False,
+) -> tuple[VectorStore | None, int | None]:
+    """Build one namespaced VectorStore covering all bimodal omics datasets.
+
+    Each per-dataset store is built on the **original** ``sample_idx`` values
+    (so the lookup against ``adata.obs_names`` succeeds), then re-keyed to
+    ``{name}:{sample_idx}`` and merged. Returns ``(None, None)`` when no omics
+    dataset uses the bimodal modality.
+
+    Parameters
+    ----------
+    cfg : TrainConfig
+        Run config (uses ``obsm_key``, ``adata_cache_dir``, ``vector_store_dir``).
+    omics_raw : dict[str, DatasetDict]
+        Raw omics datasets keyed by ``DatasetConfig.name``.
+    overwrite : bool, default False
+        Rebuild per-dataset stores even if a cached ``.mmap`` exists.
+
+    Returns
+    -------
+    tuple[VectorStore or None, int or None]
+        The merged store and its dimensionality, or ``(None, None)``.
+    """
+    from datasets import concatenate_datasets
+
+    from mmcontext.io import build_namespaced_vector_store, prepare_vector_store
+
+    bimodal = [d for d in cfg.omics_datasets if d.modality == "bimodal"]
+    if not bimodal:
+        return None, None
+
+    store_dir = Path(cfg.vector_store_dir)
+    store_dir.mkdir(parents=True, exist_ok=True)
+
+    # Encode obsm_key in the filename so switching keys on the same data uses a
+    # distinct cache file instead of silently reusing the previous key's vectors.
+    key_slug = cfg.obsm_key.replace("/", "_")
+    per_dataset: dict[str, tuple[VectorStore, list[str]]] = {}
+    for dcfg in bimodal:
+        dsdict = omics_raw[dcfg.name]
+        splits = list(dsdict.values())
+        combined = concatenate_datasets(splits) if len(splits) > 1 else splits[0]
+        store = prepare_vector_store(
+            combined,
+            obsm_key=cfg.obsm_key,
+            output_path=store_dir / f"{dcfg.name}__{key_slug}.mmap",
+            cache_dir=cfg.adata_cache_dir,
+            overwrite=overwrite,
+        )
+        orig_ids = [str(x) for x in combined["sample_idx"]]
+        per_dataset[dcfg.name] = (store, orig_ids)
+
+    merged = build_namespaced_vector_store(
+        per_dataset,
+        output_path=store_dir / "merged.mmap",
+        namespace_sep=NAMESPACE_SEP,
+    )
+    return merged, merged.dim
+
+
+# ---------------------------------------------------------------------------
+# Dataset assembly
+# ---------------------------------------------------------------------------
+@dataclass
+class AssembledData:
+    """Container for the dicts the trainer consumes."""
+
+    train_datasets: dict[str, Dataset]
+    eval_datasets: dict[str, Dataset]
+    losses: dict[str, Any]
+    evaluators: list[Any]
+
+
+def _pick_val_split(dsdict: DatasetDict) -> Dataset | None:
+    """Return the validation split of a DatasetDict, if present."""
+    for key in ("val", "validation", "dev", "test"):
+        if key in dsdict:
+            return dsdict[key]
+    return None
+
+
+def _namespace_sample_ids(ds: Dataset, name: str, sample_id_col: str = "sample_idx") -> Dataset:
+    """Prefix each ``sample_idx`` with ``{name}:`` so anchors match store keys."""
+    return ds.map(
+        lambda row: {sample_id_col: f"{name}{NAMESPACE_SEP}{row[sample_id_col]}"},
+        desc=f"Namespacing sample_idx for {name}",
+    )
+
+
+def _select_training_columns(ds: Dataset, name: str) -> Dataset:
+    """Keep only ``anchor``, ``positive`` and ``negative*`` columns.
+
+    The ST trainer feeds every dataset column to the loss as text, so stray
+    metadata columns would be misread as extra negatives. Bio datasets are
+    passed through raw, so this drops anything that is not a training input.
+    A single ``negative`` column is supported by both the loss and the
+    TripletEvaluator (which prefers ``negative`` over ``negative_1``).
+    """
+    keep = [c for c in ds.column_names if c in ("anchor", "positive") or c.startswith("negative")]
+    missing = {"anchor", "positive"} - set(keep)
+    if missing:
+        raise KeyError(
+            f"Bio dataset '{name}' is missing required column(s) {sorted(missing)}. "
+            f"Available columns: {ds.column_names}"
+        )
+    dropped = [c for c in ds.column_names if c not in keep]
+    if dropped:
+        logger.info("Bio dataset '%s': keeping %s, dropping %s", name, keep, dropped)
+    return ds.select_columns(keep)
+
+
+def _mine_negatives(
+    pairs: Dataset,
+    model: SentenceTransformer,
+    mining: MiningConfig,
+    *,
+    name: str,
+    batch_size: int,
+) -> Dataset:
+    """Mine model-specific hard negatives for an ``anchor``/``positive`` dataset.
+
+    Wraps :func:`sentence_transformers.util.mine_hard_negatives` with the
+    ``n-tuple`` output format (``anchor`` / ``positive`` / ``negative_1`` …),
+    which both :class:`MultipleNegativesRankingLoss` and the ``TripletEvaluator``
+    consume. Only the ``anchor``/``positive`` columns are used; the model is put
+    in eval mode for the encode pass and its previous mode is restored.
+
+    Parameters
+    ----------
+    pairs : datasets.Dataset
+        Dataset with at least ``anchor`` and ``positive`` columns.
+    model : SentenceTransformer
+        Pretrained pipeline (VectorStore already attached for bimodal anchors).
+    mining : MiningConfig
+        Mining parameters.
+    name : str
+        Dataset label, used only for logging.
+    batch_size : int
+        Encode batch size for the similarity search.
+
+    Returns
+    -------
+    datasets.Dataset
+        The mined ``n-tuple`` dataset.
+    """
+    from sentence_transformers.util import mine_hard_negatives
+
+    pairs = pairs.select_columns(["anchor", "positive"])
+    logger.info(
+        "Mining hard negatives for '%s': num_negatives=%d, range=[%d, %s], max_score=%s, "
+        "relative_margin=%s, sampling=%s, faiss=%s",
+        name,
+        mining.num_negatives,
+        mining.range_min,
+        mining.range_max,
+        mining.max_score,
+        mining.relative_margin,
+        mining.sampling_strategy,
+        mining.use_faiss,
+    )
+    was_training = model.training
+    model.eval()
+    try:
+        mined = mine_hard_negatives(
+            pairs,
+            model,
+            anchor_column_name="anchor",
+            positive_column_name="positive",
+            num_negatives=mining.num_negatives,
+            range_min=mining.range_min,
+            range_max=mining.range_max,
+            max_score=mining.max_score,
+            relative_margin=mining.relative_margin,
+            sampling_strategy=mining.sampling_strategy,
+            output_format="n-tuple",
+            batch_size=batch_size,
+            use_faiss=mining.use_faiss,
+            verbose=True,
+        )
+    finally:
+        if was_training:
+            model.train()
+    logger.info("Mined '%s': %d rows (from %d pairs), columns=%s", name, len(mined), len(pairs), mined.column_names)
+    return mined
+
+
+def assemble_training_data(
+    cfg: TrainConfig,
+    model: SentenceTransformer,
+    omics_raw: dict[str, DatasetDict],
+    bio_raw: dict[str, DatasetDict] | None = None,
+    *,
+    log_backend: str = "wandb",
+) -> AssembledData:
+    """Build the train/eval/loss dicts (+ evaluators) for multi-dataset training.
+
+    Omics datasets are namespaced (bimodal) and passed through
+    :func:`~mmcontext.embed.prepare_dataset`. Bio datasets are assumed to be in
+    final ``anchor`` / ``positive`` / ``negative_*`` form and passed through;
+    a small slice of their train split is reused for evaluation.
+
+    When ``cfg.mining`` is enabled, each dataset's own hard negatives are dropped
+    and model-specific negatives are mined from the anchor/positive pairs with
+    :func:`sentence_transformers.util.mine_hard_negatives` (see
+    :func:`_mine_negatives`). Both the train datasets and the evaluator's
+    validation sets are mined so they reflect the same negative distribution.
+
+    Parameters
+    ----------
+    cfg : TrainConfig
+        Run config.
+    model : SentenceTransformer
+        Built/loaded pipeline (its first module must already have the merged
+        VectorStore attached for the bimodal path); passed to ``get_loss``.
+    omics_raw : dict[str, DatasetDict]
+        Raw omics datasets keyed by ``DatasetConfig.name``.
+    bio_raw : dict[str, DatasetDict], optional
+        Raw bio datasets keyed by ``DatasetConfig.name``.
+    log_backend : str, default "wandb"
+        Forwarded to :func:`~mmcontext.utils.get_loss`.
+
+    Returns
+    -------
+    AssembledData
+        ``train_datasets``, ``eval_datasets``, ``losses``, ``evaluators``.
+    """
+    from mmcontext.embed import prepare_dataset
+    from mmcontext.utils import get_evaluator, get_loss
+
+    bio_raw = bio_raw or {}
+    train_datasets: dict[str, Dataset] = {}
+    eval_datasets: dict[str, Dataset] = {}
+    losses: dict[str, Any] = {}
+    evaluators: list[Any] = []
+    logging_steps = cfg.trainer.logging_steps
+    train_bs = cfg.trainer.per_device_train_batch_size
+    eval_bs = cfg.trainer.per_device_eval_batch_size
+    mining = cfg.mining if (cfg.mining is not None and cfg.mining.enabled) else None
+
+    # -- Omics datasets -------------------------------------------------------
+    for dcfg in cfg.omics_datasets:
+        dsdict = omics_raw[dcfg.name]
+        train_split = dsdict["train"]
+        val_split = _pick_val_split(dsdict)
+
+        if dcfg.modality == "bimodal":
+            train_split = _namespace_sample_ids(train_split, dcfg.name)
+            if val_split is not None:
+                val_split = _namespace_sample_ids(val_split, dcfg.name)
+
+        # With mining on, drop the dataset's own hard negatives and mine
+        # model-specific ones from the anchor/positive pairs instead.
+        prep_kwargs = {
+            "purpose": "train",
+            "modality": dcfg.modality,
+            "primary_cell_sentence": dcfg.primary_cell_sentence,
+            "positive_col": dcfg.positive_col,
+            "use_hard_negatives": False if mining else dcfg.use_hard_negatives,
+            "truncate": dcfg.truncate,
+            "truncate_kwargs": dcfg.truncate_kwargs,
+        }
+        train_ready = prepare_dataset(train_split, **prep_kwargs)
+        if mining:
+            train_ready = _mine_negatives(train_ready, model, mining, name=dcfg.name, batch_size=train_bs)
+        train_datasets[dcfg.name] = train_ready
+
+        losses[dcfg.name] = get_loss(
+            model=model,
+            dataset_type=dcfg.type,
+            dataset_name=dcfg.name,
+            log_backend=log_backend,
+            logging_steps=logging_steps,
+        )
+
+        if val_split is not None:
+            val_ready = prepare_dataset(val_split, **prep_kwargs)
+            if mining:
+                val_ready = _mine_negatives(val_ready, model, mining, name=f"{dcfg.name} (eval)", batch_size=eval_bs)
+            eval_datasets[dcfg.name] = val_ready
+            evaluators.append(
+                get_evaluator(
+                    dataset_type=dcfg.type,
+                    dataset=val_ready,
+                    batch_size=eval_bs,
+                    current_eval_name=dcfg.name,
+                )
+            )
+
+    # -- Bio datasets (already anchor/positive/negative) ----------------------
+    for dcfg in cfg.bio_datasets:
+        dsdict = bio_raw[dcfg.name]
+        if "train" not in dsdict:
+            logger.warning("Bio dataset '%s' has no 'train' split, skipping", dcfg.name)
+            continue
+        train_split = _select_training_columns(dsdict["train"], dcfg.name)
+        if mining:
+            # Drop the dataset's own negatives and mine model-specific ones.
+            train_split = _mine_negatives(train_split, model, mining, name=dcfg.name, batch_size=train_bs)
+        train_datasets[dcfg.name] = train_split
+
+        losses[dcfg.name] = get_loss(
+            model=model,
+            dataset_type=dcfg.type,
+            dataset_name=dcfg.name,
+            log_backend=log_backend,
+            logging_steps=logging_steps,
+        )
+
+        # Reuse a slice of train as a sanity-check eval set (mirrors old train.py).
+        # When mining, this slice already carries the mined negatives.
+        val_size = min(1000, len(train_split))
+        val_split = train_split.select(range(val_size))
+        eval_datasets[dcfg.name] = val_split
+        evaluators.append(
+            get_evaluator(
+                dataset_type=dcfg.type,
+                dataset=val_split,
+                batch_size=eval_bs,
+                current_eval_name=dcfg.name,
+            )
+        )
+
+    return AssembledData(
+        train_datasets=train_datasets,
+        eval_datasets=eval_datasets,
+        losses=losses,
+        evaluators=evaluators,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Model-card details
+# ---------------------------------------------------------------------------
+def format_training_details(cfg: TrainConfig) -> str:
+    """Render a Markdown summary of how the model was (fine)tuned.
+
+    Intended for the Hugging Face model card so it is clear from the repo whether
+    the model was finetuned with hard-negative mining and with which settings.
+
+    Parameters
+    ----------
+    cfg : TrainConfig
+        Run config.
+
+    Returns
+    -------
+    str
+        A Markdown bullet list.
+    """
+    lines: list[str] = []
+    if cfg.model:
+        lines.append(f"- **Continued from**: `{cfg.model}`")
+    else:
+        lines.append(f"- **Base text encoder**: `{cfg.text_encoder.name}`")
+    if cfg.omics_datasets:
+        lines.append(f"- **Omics embedding key (obsm)**: `{cfg.obsm_key}`")
+    lines.append(f"- **Shared embedding dimension**: {cfg.shared_dim}")
+
+    mining = cfg.mining
+    if mining is not None and mining.enabled:
+        lines.append(
+            "- **Hard-negative mining**: ✅ enabled "
+            "(`sentence_transformers.util.mine_hard_negatives`, `n-tuple` output)"
+        )
+        lines.append(f"    - `num_negatives`: {mining.num_negatives}")
+        lines.append(f"    - `range_min` / `range_max`: {mining.range_min} / {mining.range_max}")
+        lines.append(f"    - `max_score`: {mining.max_score}")
+        lines.append(f"    - `relative_margin`: {mining.relative_margin}")
+        lines.append(f"    - `sampling_strategy`: {mining.sampling_strategy}")
+        lines.append(f"    - `use_faiss`: {mining.use_faiss}")
+        lines.append(
+            "    - Negatives are mined per-anchor in the model's current shared embedding space "
+            "(omics anchors resolved via the VectorStore for bimodal datasets), replacing any "
+            "dataset-provided hard negatives."
+        )
+    else:
+        lines.append("- **Hard-negative mining**: ❌ not used (dataset-provided hard negatives, if any)")
+
+    train_names = [d.name for d in cfg.omics_datasets] + [d.name for d in cfg.bio_datasets]
+    if train_names:
+        lines.append(f"- **Training datasets**: {', '.join(train_names)}")
+    return "\n".join(lines)
diff --git a/src/mmcontext/utils.py b/src/mmcontext/utils.py
index b6dfdb0..695d66d 100644
--- a/src/mmcontext/utils.py
+++ b/src/mmcontext/utils.py
@@ -1,26 +1,21 @@
 from __future__ import annotations
 
 import inspect
-import json
 import logging
 import os
 import re
-import shutil
 from collections.abc import Iterable
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 import anndata
 import numpy as np
 import pandas as pd
-import scanpy as sc
 import scipy.sparse as sp
 import torch
-from omegaconf import DictConfig
 from sentence_transformers import evaluation, losses
 from torch import nn
-from tqdm import tqdm
 
 try:
     import wandb  # optional; wrapper no-ops if not available or not initialized
@@ -90,7 +85,7 @@ def _truncate_batch(batch):
             if filter_strings:
                 original_count = len(tokens)
                 filtered_tokens = []
-                removed_counts = {filter_str: 0 for filter_str in filter_strings}
+                removed_counts = dict.fromkeys(filter_strings, 0)
 
                 for token in tokens:
                     should_keep = True
@@ -1332,77 +1327,3 @@ def _resolve_negatives(batch):
         raise TypeError("resolve_negative_indices expects a Dataset or DatasetDict")
 
     return processed_dataset
-
-
-'''
-def prepare_omics_resources(hf_ds, *, prefix: str = "sample_idx:"):
-    """
-    Parameters
-    ----------
-    hf_ds : datasets.Dataset
-        Must expose ``sample_idx`` and ``data_representation``.
-    prefix : str, optional
-        Tag added in front of each key in the lookup (default ``"sample_idx:"``).
-
-    Returns
-    -------
-    embedding_matrix : np.ndarray  (num_ids, dim)
-    lookup           : dict[str, int]
-    """
-    lookup = {}
-    vectors = []
-    # loop over row with progress bar
-    for row in tqdm(hf_ds, desc="Preparing omics resources", unit="row"):
-        sid = row["sample_idx"]
-        vec = row["data_representation"]
-        row_id = len(lookup)  # 0‥N–1 in encounter order
-        lookup[f"{prefix}{sid}"] = row_id
-        vectors.append(vec)
-
-    embedding_matrix = np.stack(vectors).astype("float32")
-    return embedding_matrix, lookup
-'''
-
-
-def copy_resolved_config(cfg, hydra_output_dir: Path, named_output_dir: Path) -> None:
-    """
-    Copy the resolved config to the output directory.
-
-    Serialize the *instantiated* Hydra config twice: once in the Hydra run
-    directory and once in the named output directory so downstream users can
-    locate it without knowing the run-dir.
-
-    The config is enriched with a few runtime fields first (git SHA, SLURM
-    job-ID, command line, etc.).
-    """
-    import subprocess
-
-    from hydra.utils import to_absolute_path
-    from omegaconf import OmegaConf
-
-    cfg_dict = OmegaConf.to_container(cfg, resolve=True)
-
-    # --- enrich with runtime metadata -------------------------------------
-    git_sha = subprocess.run(
-        ["git", "rev-parse", "HEAD"],
-        capture_output=True,
-        text=True,
-        check=False,
-    ).stdout.strip()
-    cfg_dict["_meta"] = {
-        "git_sha": git_sha,
-        "cmd": " ".join([to_absolute_path("main.py"), *os.sys.argv[1:]]),
-        "slurm_job_id": os.getenv("SLURM_JOB_ID") if cfg.slurm.store_id else None,
-    }
-
-    # -----------------------------------------------------------------------
-    for target_dir in (hydra_output_dir, named_output_dir):
-        target_dir.mkdir(parents=True, exist_ok=True)
-        out = target_dir / "resolved_config.json"
-        with out.open("w") as fp:
-            json.dump(cfg_dict, fp, indent=2)
-        logger.info("Wrote resolved Hydra config → %s", out)
-
-    # Also copy *raw* yaml files for debugging
-    orig_conf_dir = Path("conf").absolute()
-    shutil.copytree(orig_conf_dir, named_output_dir / "conf_raw", dirs_exist_ok=True)
diff --git a/tests/conftest.py b/tests/conftest.py
index f21c0ac..61eb01c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -4,6 +4,7 @@
 from typing import Optional, Union
 from unittest.mock import patch
 
+import anndata as ad
 import datasets
 import numpy as np
 import pandas as pd
@@ -13,7 +14,16 @@
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.models.auto import configuration_auto, modeling_auto
 
-from mmcontext.mmcontextencoder import MMContextEncoder, MMContextProcessor
+from mmcontext._legacy.mmcontextencoder import MMContextEncoder
+
+# -----------------------------------------------------------------------------
+# Zarr v3 format
+# -----------------------------------------------------------------------------
+# anndata still defaults to writing the zarr v2 on-disk format. Force v3 so test
+# fixtures that call ``adata.write_zarr(...)`` produce v3 stores, matching the
+# format produced in production and exercising the v3 read path. Loaded for the
+# whole test session because every test dir inherits this root conftest.
+ad.settings.zarr_write_format = 3
 
 # -----------------------------------------------------------------------------
 # Logging
diff --git a/tests/test_adapter_callback.py b/tests/test_adapter_callback.py
index e69de29..2177996 100644
--- a/tests/test_adapter_callback.py
+++ b/tests/test_adapter_callback.py
@@ -0,0 +1,363 @@
+"""Tests for UnfreezeTextEncoderCallback and UnfreezeAdapterCallback."""
+
+import logging
+from types import SimpleNamespace
+
+import pytest
+import torch.nn as nn
+
+from mmcontext.callback import UnfreezeAdapterCallback, UnfreezeTextEncoderCallback
+from mmcontext.modules.adapter_module import AdapterModule
+from mmcontext.modules.mmcontext_module import MMContextModule
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_state(epoch: float) -> SimpleNamespace:
+    return SimpleNamespace(epoch=epoch)
+
+
+def _all_frozen(module: nn.Module) -> bool:
+    return all(not p.requires_grad for p in module.parameters())
+
+
+def _all_trainable(module: nn.Module) -> bool:
+    return all(p.requires_grad for p in module.parameters())
+
+
+class _FakePipeline(nn.Module):
+    """Lightweight SentenceTransformer-like container that supports model[i] indexing."""
+
+    def __init__(self, *modules: nn.Module):
+        super().__init__()
+        self._modules_list = nn.ModuleList(modules)
+
+    def __getitem__(self, idx: int) -> nn.Module:
+        return self._modules_list[idx]
+
+    def children(self):
+        return iter(self._modules_list)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture()
+def mmcontext_module() -> MMContextModule:
+    """MMContextModule with patched auto_model (already patched session-wide in conftest)."""
+    return MMContextModule("bert-base-uncased")
+
+
+@pytest.fixture()
+def adapter_module() -> AdapterModule:
+    return AdapterModule(text_input_dim=32, omics_input_dim=32, shared_dim=16, hidden_dim=None)
+
+
+@pytest.fixture()
+def pipeline_without_attention(mmcontext_module, adapter_module) -> _FakePipeline:
+    """MMContextModule → AdapterModule (adapter at index 1)."""
+    return _FakePipeline(mmcontext_module, adapter_module)
+
+
+@pytest.fixture()
+def pipeline_with_attention(mmcontext_module, adapter_module) -> _FakePipeline:
+    """MMContextModule → dummy attention module → AdapterModule (adapter at index 2)."""
+    dummy_attention = nn.Linear(16, 16)  # not an AdapterModule, stands in for OmicsAttentionModule
+    return _FakePipeline(mmcontext_module, dummy_attention, adapter_module)
+
+
+@pytest.fixture()
+def pipeline_no_adapter(mmcontext_module) -> _FakePipeline:
+    """Pipeline without any AdapterModule."""
+    return _FakePipeline(mmcontext_module)
+
+
+# ---------------------------------------------------------------------------
+# UnfreezeTextEncoderCallback
+# ---------------------------------------------------------------------------
+
+
+class TestUnfreezeTextEncoderCallback:
+    def _call_epoch(self, callback, model, epoch: float):
+        """Simulate on_epoch_begin for a given epoch."""
+        state = _make_state(epoch)
+        callback.on_epoch_begin(args=None, state=state, control=None, model=model)
+
+    def test_text_encoder_frozen_before_unfreeze_epoch(self, pipeline_without_attention):
+        """Text encoder parameters must stay frozen before the unfreeze epoch."""
+        model = pipeline_without_attention
+        model[0].freeze_all_but_top_layers(0)
+
+        cb = UnfreezeTextEncoderCallback(unfreeze_epoch=2.0)
+        self._call_epoch(cb, model, epoch=0.0)
+        self._call_epoch(cb, model, epoch=1.0)
+
+        assert _all_frozen(model[0].auto_model), "Text encoder should still be frozen before epoch 2"
+        assert not cb.unfrozen
+
+    def test_text_encoder_unfreezes_at_target_epoch(self, pipeline_without_attention):
+        """Text encoder should be unfrozen exactly at the configured epoch."""
+        model = pipeline_without_attention
+        model[0].freeze_all_but_top_layers(0)
+
+        cb = UnfreezeTextEncoderCallback(unfreeze_epoch=2.0)
+        self._call_epoch(cb, model, epoch=2.0)
+
+        assert _all_trainable(model[0].auto_model), "Text encoder should be unfrozen at epoch 2"
+        assert cb.unfrozen
+
+    def test_text_encoder_unfreezes_past_target_epoch(self, pipeline_without_attention):
+        """Unfreezing also triggers when epoch exceeds the target (e.g., epoch 3 > target 2)."""
+        model = pipeline_without_attention
+        model[0].freeze_all_but_top_layers(0)
+
+        cb = UnfreezeTextEncoderCallback(unfreeze_epoch=2.0)
+        self._call_epoch(cb, model, epoch=3.0)
+
+        assert _all_trainable(model[0].auto_model)
+        assert cb.unfrozen
+
+    def test_text_encoder_unfreezes_only_once(self, pipeline_without_attention):
+        """The unfreeze action should happen exactly once regardless of subsequent epochs."""
+        model = pipeline_without_attention
+        model[0].freeze_all_but_top_layers(0)
+
+        cb = UnfreezeTextEncoderCallback(unfreeze_epoch=1.0)
+        self._call_epoch(cb, model, epoch=1.0)
+
+        # Manually re-freeze to check the callback doesn't unfreeze again
+        model[0].freeze_all_but_top_layers(0)
+        self._call_epoch(cb, model, epoch=2.0)
+
+        assert _all_frozen(model[0].auto_model), "Should not re-unfreeze after unfrozen=True is set"
+
+    def test_warns_when_no_compatible_module(self, caplog):
+        """A warning is logged when model[0] has no auto_model attribute."""
+        plain_nn = nn.Linear(4, 4)
+        model = _FakePipeline(plain_nn)
+
+        cb = UnfreezeTextEncoderCallback(unfreeze_epoch=0.0)
+        with caplog.at_level(logging.WARNING):
+            self._call_epoch(cb, model, epoch=0.0)
+
+        assert any("compatible text encoder" in r.message for r in caplog.records)
+        assert not cb.unfrozen
+
+    def test_different_unfreeze_epochs(self, pipeline_without_attention):
+        """Two callbacks with different epochs unfreeze independently."""
+        model = pipeline_without_attention
+        model[0].freeze_all_but_top_layers(0)
+
+        cb_epoch1 = UnfreezeTextEncoderCallback(unfreeze_epoch=1.0)
+        cb_epoch3 = UnfreezeTextEncoderCallback(unfreeze_epoch=3.0)
+
+        # At epoch 1 — only the epoch-1 callback fires
+        self._call_epoch(cb_epoch1, model, epoch=1.0)
+        assert cb_epoch1.unfrozen
+        assert not cb_epoch3.unfrozen
+
+        model[0].freeze_all_but_top_layers(0)  # re-freeze for the next check
+        # At epoch 3 — the epoch-3 callback fires
+        self._call_epoch(cb_epoch3, model, epoch=3.0)
+        assert cb_epoch3.unfrozen
+
+
+# ---------------------------------------------------------------------------
+# UnfreezeAdapterCallback
+# ---------------------------------------------------------------------------
+
+
+class TestUnfreezeAdapterCallback:
+    def _train_begin(self, callback, model):
+        callback.on_train_begin(args=None, state=None, control=None, model=model)
+
+    def _call_epoch(self, callback, model, epoch: float):
+        state = _make_state(epoch)
+        callback.on_epoch_begin(args=None, state=state, control=None, model=model)
+
+    # --- initialization ---
+
+    def test_freezes_both_on_train_begin(self, pipeline_without_attention, adapter_module):
+        cb = UnfreezeAdapterCallback(freeze_text_adapter=True, freeze_omics_adapter=True)
+        self._train_begin(cb, pipeline_without_attention)
+
+        assert _all_frozen(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+        assert cb.adapters_initialized
+
+    def test_freezes_only_text_on_train_begin(self, pipeline_without_attention, adapter_module):
+        cb = UnfreezeAdapterCallback(freeze_text_adapter=True, freeze_omics_adapter=False)
+        self._train_begin(cb, pipeline_without_attention)
+
+        assert _all_frozen(adapter_module.text_proj)
+        assert _all_trainable(adapter_module.omics_proj)
+
+    def test_freezes_only_omics_on_train_begin(self, pipeline_without_attention, adapter_module):
+        cb = UnfreezeAdapterCallback(freeze_text_adapter=False, freeze_omics_adapter=True)
+        self._train_begin(cb, pipeline_without_attention)
+
+        assert _all_trainable(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+
+    # --- epoch-based unfreezing ---
+
+    def test_text_adapter_unfreezes_at_epoch_1(self, pipeline_without_attention, adapter_module):
+        cb = UnfreezeAdapterCallback(
+            freeze_text_adapter=True,
+            freeze_omics_adapter=True,
+            unfreeze_text_adapter_epoch=1.0,
+        )
+        self._train_begin(cb, pipeline_without_attention)
+
+        # Before epoch 1 — still frozen
+        self._call_epoch(cb, pipeline_without_attention, epoch=0.0)
+        assert _all_frozen(adapter_module.text_proj)
+        assert not cb.text_adapter_unfrozen
+
+        # At epoch 1 — text unfrozen, omics still frozen
+        self._call_epoch(cb, pipeline_without_attention, epoch=1.0)
+        assert _all_trainable(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+        assert cb.text_adapter_unfrozen
+        assert not cb.omics_adapter_unfrozen
+
+    def test_omics_adapter_unfreezes_at_epoch_2(self, pipeline_without_attention, adapter_module):
+        cb = UnfreezeAdapterCallback(
+            freeze_text_adapter=True,
+            freeze_omics_adapter=True,
+            unfreeze_text_adapter_epoch=1.0,
+            unfreeze_omics_adapter_epoch=2.0,
+        )
+        self._train_begin(cb, pipeline_without_attention)
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=1.0)  # unfreeze text
+        assert cb.text_adapter_unfrozen
+        assert not cb.omics_adapter_unfrozen
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=2.0)  # unfreeze omics
+        assert _all_trainable(adapter_module.omics_proj)
+        assert cb.omics_adapter_unfrozen
+
+    def test_adapters_stay_frozen_without_unfreeze_epoch(self, pipeline_without_attention, adapter_module):
+        """When unfreeze_*_epoch is None, frozen adapters remain frozen forever."""
+        cb = UnfreezeAdapterCallback(freeze_text_adapter=True, freeze_omics_adapter=True)
+        self._train_begin(cb, pipeline_without_attention)
+
+        for epoch in [0.0, 1.0, 5.0, 100.0]:
+            self._call_epoch(cb, pipeline_without_attention, epoch=epoch)
+
+        assert _all_frozen(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+
+    def test_unfreezes_only_once(self, pipeline_without_attention, adapter_module):
+        """Each adapter is unfrozen at most once."""
+        cb = UnfreezeAdapterCallback(
+            freeze_text_adapter=True,
+            freeze_omics_adapter=True,
+            unfreeze_text_adapter_epoch=1.0,
+            unfreeze_omics_adapter_epoch=1.0,
+        )
+        self._train_begin(cb, pipeline_without_attention)
+        self._call_epoch(cb, pipeline_without_attention, epoch=1.0)
+
+        assert cb.text_adapter_unfrozen
+        assert cb.omics_adapter_unfrozen
+
+        # Re-freeze manually and simulate another epoch — callback should not unfreeze again
+        adapter_module.freeze_text_proj()
+        adapter_module.freeze_omics_proj()
+        self._call_epoch(cb, pipeline_without_attention, epoch=2.0)
+
+        assert _all_frozen(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+
+    # --- adapter at position 2 (with OmicsAttentionModule) ---
+
+    def test_finds_adapter_at_position_2(self, pipeline_with_attention, adapter_module):
+        """Adapter is found even when it's at index 2 (behind a dummy attention module)."""
+        cb = UnfreezeAdapterCallback(
+            freeze_text_adapter=True,
+            freeze_omics_adapter=True,
+            unfreeze_text_adapter_epoch=1.0,
+            unfreeze_omics_adapter_epoch=2.0,
+        )
+        self._train_begin(cb, pipeline_with_attention)
+
+        assert cb.adapters_initialized
+        assert cb._adapter is adapter_module
+
+        self._call_epoch(cb, pipeline_with_attention, epoch=1.0)
+        assert _all_trainable(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+
+        self._call_epoch(cb, pipeline_with_attention, epoch=2.0)
+        assert _all_trainable(adapter_module.omics_proj)
+
+    # --- no adapter present ---
+
+    def test_warns_when_no_adapter(self, pipeline_no_adapter, caplog):
+        """A warning is emitted and adapters_initialized stays False when no AdapterModule found."""
+        cb = UnfreezeAdapterCallback(freeze_text_adapter=True, freeze_omics_adapter=True)
+
+        with caplog.at_level(logging.WARNING):
+            self._train_begin(cb, pipeline_no_adapter)
+
+        assert any("No AdapterModule found" in r.message for r in caplog.records)
+        assert not cb.adapters_initialized
+
+    def test_epoch_begin_noop_without_initialization(self, pipeline_no_adapter):
+        """on_epoch_begin is a no-op when adapters_initialized is False."""
+        cb = UnfreezeAdapterCallback(
+            freeze_text_adapter=True,
+            unfreeze_text_adapter_epoch=0.0,
+        )
+        # Deliberately skip on_train_begin so adapters_initialized stays False
+        state = _make_state(5.0)
+        result = cb.on_epoch_begin(args=None, state=state, control=None, model=pipeline_no_adapter)
+        assert result is None  # control passed through unchanged
+
+    # --- different epoch scenarios ---
+
+    def test_text_and_omics_unfreeze_at_different_epochs(self, pipeline_without_attention, adapter_module):
+        """Text unfreezes at epoch 1, omics at epoch 3 — verify each fires independently."""
+        cb = UnfreezeAdapterCallback(
+            freeze_text_adapter=True,
+            freeze_omics_adapter=True,
+            unfreeze_text_adapter_epoch=1.0,
+            unfreeze_omics_adapter_epoch=3.0,
+        )
+        self._train_begin(cb, pipeline_without_attention)
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=0.0)
+        assert _all_frozen(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=1.0)
+        assert _all_trainable(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=2.0)
+        assert _all_trainable(adapter_module.text_proj)
+        assert _all_frozen(adapter_module.omics_proj)
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=3.0)
+        assert _all_trainable(adapter_module.text_proj)
+        assert _all_trainable(adapter_module.omics_proj)
+
+    def test_fractional_unfreeze_epoch(self, pipeline_without_attention, adapter_module):
+        """Fractional unfreeze epochs work correctly (e.g., epoch 0.5)."""
+        cb = UnfreezeAdapterCallback(
+            freeze_omics_adapter=True,
+            unfreeze_omics_adapter_epoch=0.5,
+        )
+        self._train_begin(cb, pipeline_without_attention)
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=0.0)
+        assert _all_frozen(adapter_module.omics_proj)
+
+        self._call_epoch(cb, pipeline_without_attention, epoch=0.5)
+        assert _all_trainable(adapter_module.omics_proj)
diff --git a/tests/test_adapter_module.py b/tests/test_adapter_module.py
new file mode 100644
index 0000000..80692f4
--- /dev/null
+++ b/tests/test_adapter_module.py
@@ -0,0 +1,388 @@
+"""Tests for AdapterModule — modality-aware projection for the ST pipeline.
+
+The AdapterModule sits after MMContextModule in the sentence-transformers
+pipeline. It reads ``modality_ids`` from the features dict and applies
+separate learned projections for text (modality_id=0) and omics
+(modality_id=1) tokens, mapping them into a shared embedding space.
+
+These tests define the contract and are written FIRST (TDD).
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import pytest
+import torch
+
+from mmcontext.modules.adapter_module import AdapterModule
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def tmp_dir():
+    with tempfile.TemporaryDirectory() as d:
+        yield d
+
+
+@pytest.fixture
+def real_safetensors():
+    """Undo the global safetensors.torch.load_model patch for persistence tests.
+
+    The root conftest.py patches ``safetensors.torch.load_model`` to return None,
+    which prevents real weight loading. This fixture restores the original
+    function for tests that need actual save/load roundtrips.
+    """
+    import importlib
+
+    import safetensors.torch
+
+    importlib.reload(safetensors.torch)
+    yield
+    # The session-scoped patch in conftest will reassert on the next test that needs it
+
+
+@pytest.fixture
+def adapter():
+    """Default adapter: text_dim=32, omics_dim=8, shared_dim=16, hidden_dim=64."""
+    return AdapterModule(
+        text_input_dim=32,
+        omics_input_dim=8,
+        shared_dim=16,
+        hidden_dim=64,
+    )
+
+
+@pytest.fixture
+def identity_adapter():
+    """Adapter in identity mode (no projection)."""
+    return AdapterModule(
+        text_input_dim=16,
+        omics_input_dim=16,
+        shared_dim=16,
+        hidden_dim=None,
+        force_identity=True,
+    )
+
+
+def _make_features(
+    token_embeddings: torch.Tensor,
+    attention_mask: torch.Tensor,
+    modality_ids: torch.Tensor,
+) -> dict[str, torch.Tensor]:
+    """Helper to build a features dict matching the pipeline contract."""
+    return {
+        "token_embeddings": token_embeddings,
+        "attention_mask": attention_mask,
+        "modality_ids": modality_ids,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Forward — single modality
+# ---------------------------------------------------------------------------
+class TestForwardSingleModality:
+    """Tests for forward() with all-text or all-omics batches."""
+
+    def test_forward_text_only(self, adapter):
+        """All-text batch produces output with shared_dim."""
+        B, L, D_text = 2, 5, 32
+        features = _make_features(
+            token_embeddings=torch.randn(B, L, D_text),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=torch.zeros(B, L, dtype=torch.long),  # all text
+        )
+        result = adapter(features)
+
+        assert result["token_embeddings"].shape == (B, L, 16)
+
+    def test_forward_omics_only(self, adapter):
+        """All-omics batch produces output with shared_dim."""
+        B, L, D_omics = 3, 1, 8
+        features = _make_features(
+            token_embeddings=torch.randn(B, L, D_omics),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=torch.ones(B, L, dtype=torch.long),  # all omics
+        )
+        result = adapter(features)
+
+        assert result["token_embeddings"].shape == (B, L, 16)
+
+    def test_forward_text_output_dimension(self, adapter):
+        """Text projection maps D_text → D_shared exactly."""
+        features = _make_features(
+            token_embeddings=torch.randn(1, 3, 32),
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=torch.zeros(1, 3, dtype=torch.long),
+        )
+        result = adapter(features)
+        assert result["token_embeddings"].shape[-1] == adapter.get_sentence_embedding_dimension()
+
+    def test_forward_omics_output_dimension(self, adapter):
+        """Omics projection maps D_omics → D_shared exactly."""
+        features = _make_features(
+            token_embeddings=torch.randn(1, 1, 8),
+            attention_mask=torch.ones(1, 1, dtype=torch.long),
+            modality_ids=torch.ones(1, 1, dtype=torch.long),
+        )
+        result = adapter(features)
+        assert result["token_embeddings"].shape[-1] == adapter.get_sentence_embedding_dimension()
+
+
+# ---------------------------------------------------------------------------
+# Forward — mixed batch
+# ---------------------------------------------------------------------------
+class TestForwardMixedBatch:
+    """Tests for forward() with mixed text + omics tokens."""
+
+    def test_forward_mixed_batch(self):
+        """Mixed batch: text and omics tokens get different projections."""
+        # Both modalities have same input dim for simplicity
+        adapter = AdapterModule(text_input_dim=16, omics_input_dim=16, shared_dim=8, hidden_dim=32)
+        B, L = 1, 4
+        features = _make_features(
+            token_embeddings=torch.randn(B, L, 16),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            # First 2 tokens text, last 2 omics
+            modality_ids=torch.tensor([[0, 0, 1, 1]], dtype=torch.long),
+        )
+        result = adapter(features)
+
+        assert result["token_embeddings"].shape == (B, L, 8)
+
+    def test_pad_tokens_stay_zero(self, adapter):
+        """Pad tokens (modality_id=2) remain zero after projection."""
+        B, L = 1, 4
+        embeddings = torch.randn(B, L, 32)
+        embeddings[0, 3, :] = 0.0  # pad token is zero
+
+        features = _make_features(
+            token_embeddings=embeddings,
+            attention_mask=torch.tensor([[1, 1, 1, 0]], dtype=torch.long),
+            modality_ids=torch.tensor([[0, 0, 0, 2]], dtype=torch.long),  # last is pad
+        )
+        result = adapter(features)
+
+        # Pad token should remain zero
+        assert torch.all(result["token_embeddings"][0, 3] == 0)
+
+
+# ---------------------------------------------------------------------------
+# Separate weights
+# ---------------------------------------------------------------------------
+class TestSeparateWeights:
+    """Tests verifying text and omics projections are independent."""
+
+    def test_separate_weights(self, adapter):
+        """text_proj and omics_proj have independent parameter sets."""
+        text_params = {id(p) for p in adapter.text_proj.parameters()}
+        omics_params = {id(p) for p in adapter.omics_proj.parameters()}
+
+        # No overlap
+        assert text_params.isdisjoint(omics_params)
+        # Both have parameters
+        assert len(text_params) > 0
+        assert len(omics_params) > 0
+
+    def test_text_omics_produce_different_outputs(self):
+        """Same input through text vs omics projection gives different results."""
+        adapter = AdapterModule(text_input_dim=16, omics_input_dim=16, shared_dim=8, hidden_dim=32)
+        x = torch.randn(1, 3, 16)
+
+        text_features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=torch.zeros(1, 3, dtype=torch.long),
+        )
+        omics_features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=torch.ones(1, 3, dtype=torch.long),
+        )
+
+        text_result = adapter(text_features)
+        omics_result = adapter(omics_features)
+
+        # Different projections should produce different outputs (with overwhelming probability)
+        assert not torch.allclose(text_result["token_embeddings"], omics_result["token_embeddings"])
+
+
+# ---------------------------------------------------------------------------
+# Gradient flow
+# ---------------------------------------------------------------------------
+class TestGradientFlow:
+    """Tests for gradient computation through adapter projections."""
+
+    def test_gradient_flow_text(self, adapter):
+        """Gradients flow through text projection."""
+        features = _make_features(
+            token_embeddings=torch.randn(1, 3, 32, requires_grad=True),
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=torch.zeros(1, 3, dtype=torch.long),
+        )
+        result = adapter(features)
+        # Use squared loss: plain sum(LayerNorm(x)) has zero gradient
+        # because LayerNorm centers each token and summing centered
+        # values cancels out. Squaring breaks this cancellation.
+        loss = (result["token_embeddings"] ** 2).sum()
+        loss.backward()
+
+        for param in adapter.text_proj.parameters():
+            assert param.grad is not None
+            assert param.grad.abs().sum() > 0
+
+    def test_gradient_flow_omics(self, adapter):
+        """Gradients flow through omics projection."""
+        features = _make_features(
+            token_embeddings=torch.randn(2, 1, 8, requires_grad=True),
+            attention_mask=torch.ones(2, 1, dtype=torch.long),
+            modality_ids=torch.ones(2, 1, dtype=torch.long),
+        )
+        result = adapter(features)
+        loss = (result["token_embeddings"] ** 2).sum()
+        loss.backward()
+
+        for param in adapter.omics_proj.parameters():
+            assert param.grad is not None
+            assert param.grad.abs().sum() > 0
+
+    def test_weights_update(self, adapter):
+        """Optimizer step changes both projection weights."""
+        optimizer = torch.optim.SGD(adapter.parameters(), lr=0.1)
+
+        # Snapshot weights before
+        text_before = {n: p.clone() for n, p in adapter.text_proj.named_parameters()}
+        omics_before = {n: p.clone() for n, p in adapter.omics_proj.named_parameters()}
+
+        # Forward + backward through text
+        features_t = _make_features(
+            token_embeddings=torch.randn(2, 3, 32),
+            attention_mask=torch.ones(2, 3, dtype=torch.long),
+            modality_ids=torch.zeros(2, 3, dtype=torch.long),
+        )
+        result_t = adapter(features_t)
+        loss_t = result_t["token_embeddings"].sum()
+
+        # Forward + backward through omics
+        features_o = _make_features(
+            token_embeddings=torch.randn(2, 1, 8),
+            attention_mask=torch.ones(2, 1, dtype=torch.long),
+            modality_ids=torch.ones(2, 1, dtype=torch.long),
+        )
+        result_o = adapter(features_o)
+        loss_o = result_o["token_embeddings"].sum()
+
+        total_loss = loss_t + loss_o
+        total_loss.backward()
+        optimizer.step()
+
+        # Both projections should have changed (check total param delta,
+        # not per-parameter allclose, since some biases may get tiny gradients)
+        text_delta = sum((p - text_before[n]).abs().sum().item() for n, p in adapter.text_proj.named_parameters())
+        omics_delta = sum((p - omics_before[n]).abs().sum().item() for n, p in adapter.omics_proj.named_parameters())
+        assert text_delta > 0, "text_proj parameters did not change"
+        assert omics_delta > 0, "omics_proj parameters did not change"
+
+
+# ---------------------------------------------------------------------------
+# Identity mode
+# ---------------------------------------------------------------------------
+class TestIdentityMode:
+    """Tests for identity/passthrough mode."""
+
+    def test_identity_passthrough(self, identity_adapter):
+        """When force_identity=True and dims match, input passes through."""
+        x = torch.randn(1, 3, 16)
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=torch.zeros(1, 3, dtype=torch.long),
+        )
+        result = identity_adapter(features)
+        torch.testing.assert_close(result["token_embeddings"], x)
+
+
+# ---------------------------------------------------------------------------
+# Properties
+# ---------------------------------------------------------------------------
+class TestProperties:
+    """Tests for module properties and metadata."""
+
+    def test_get_sentence_embedding_dimension(self, adapter):
+        """Returns D_shared."""
+        assert adapter.get_sentence_embedding_dimension() == 16
+
+    def test_preserves_attention_mask(self, adapter):
+        """Attention mask passes through unchanged."""
+        mask = torch.tensor([[1, 1, 0]], dtype=torch.long)
+        features = _make_features(
+            token_embeddings=torch.randn(1, 3, 32),
+            attention_mask=mask,
+            modality_ids=torch.zeros(1, 3, dtype=torch.long),
+        )
+        result = adapter(features)
+        torch.testing.assert_close(result["attention_mask"], mask)
+
+    def test_preserves_modality_ids(self, adapter):
+        """modality_ids pass through unchanged."""
+        mod_ids = torch.tensor([[0, 0, 1]], dtype=torch.long)
+        features = _make_features(
+            token_embeddings=torch.randn(1, 3, 32),
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=mod_ids,
+        )
+        # Need same input dim for this test
+        adapter2 = AdapterModule(text_input_dim=32, omics_input_dim=32, shared_dim=16, hidden_dim=64)
+        result = adapter2(features)
+        torch.testing.assert_close(result["modality_ids"], mod_ids)
+
+
+# ---------------------------------------------------------------------------
+# Persistence
+# ---------------------------------------------------------------------------
+class TestPersistence:
+    """Tests for save/load roundtrip."""
+
+    def test_save_creates_files(self, adapter, tmp_dir, real_safetensors):
+        """save() creates config and weight files on disk."""
+        adapter.save(tmp_dir)
+
+        config_path = os.path.join(tmp_dir, adapter.config_file_name)
+        assert os.path.isfile(config_path)
+        assert os.path.isfile(os.path.join(tmp_dir, "model.safetensors"))
+
+    def test_save_load_roundtrip(self, adapter, tmp_dir, real_safetensors):
+        """Config and weights survive save/load cycle."""
+        adapter.save(tmp_dir)
+        loaded = AdapterModule.load(tmp_dir)
+
+        assert loaded.get_sentence_embedding_dimension() == adapter.get_sentence_embedding_dimension()
+        assert loaded.text_input_dim == adapter.text_input_dim
+        assert loaded.omics_input_dim == adapter.omics_input_dim
+
+    def test_save_load_produces_same_output(self, adapter, tmp_dir, real_safetensors):
+        """Loaded adapter produces identical output to original."""
+        adapter.eval()
+        x = torch.randn(1, 3, 32)
+        features = _make_features(
+            token_embeddings=x,
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=torch.zeros(1, 3, dtype=torch.long),
+        )
+        original_out = adapter(features)["token_embeddings"].detach()
+
+        adapter.save(tmp_dir)
+        loaded = AdapterModule.load(tmp_dir)
+        loaded.eval()
+
+        features2 = _make_features(
+            token_embeddings=x,
+            attention_mask=torch.ones(1, 3, dtype=torch.long),
+            modality_ids=torch.zeros(1, 3, dtype=torch.long),
+        )
+        loaded_out = loaded(features2)["token_embeddings"].detach()
+
+        torch.testing.assert_close(original_out, loaded_out)
diff --git a/tests/test_dataset_prep.py b/tests/test_dataset_prep.py
new file mode 100644
index 0000000..62a1eee
--- /dev/null
+++ b/tests/test_dataset_prep.py
@@ -0,0 +1,217 @@
+"""Tests for mmcontext.embed.dataset_prep — unified train/inference prep.
+
+The pure ``prepare_dataset`` tests run entirely in-memory. The
+``prepare_inference`` integration test uses a local AnnData zarr store on disk
+(no network) and a lightweight fake model that only needs ``set_vector_store``.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import pytest
+from datasets import Dataset, DatasetDict
+
+from mmcontext.embed import InferenceData, prepare_dataset, prepare_inference
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def tmp_dir():
+    with tempfile.TemporaryDirectory() as d:
+        yield d
+
+
+def _train_ds(n: int = 6) -> Dataset:
+    """Training-shaped dataset with positive + negative_*_idx columns."""
+    return Dataset.from_dict(
+        {
+            "sample_idx": [f"cell_{i}" for i in range(n)],
+            "cell_sentence_1": [f"GENE{i} CALM1 GNAS" for i in range(n)],
+            "cell_sentence_2": [f"gene desc {i}" for i in range(n)],
+            "positive": [f"a description of cell {i}" for i in range(n)],
+            "negative_1_idx": [f"cell_{(i + 1) % n}" for i in range(n)],
+            "negative_2_idx": [f"cell_{(i + 2) % n}" for i in range(n)],
+            "adata_link": ["local://chunk"] * n,
+        }
+    )
+
+
+def _test_ds(n: int = 6) -> Dataset:
+    """Inference-shaped dataset: no positive/negative columns."""
+    return Dataset.from_dict(
+        {
+            "sample_idx": [f"cell_{i}" for i in range(n)],
+            "cell_sentence_1": [f"GENE{i} CALM1 GNAS" for i in range(n)],
+            "cell_sentence_2": [f"gene desc {i}" for i in range(n)],
+            "adata_link": ["local://chunk"] * n,
+        }
+    )
+
+
+@pytest.fixture
+def adata_zarr(tmp_dir):
+    """A real AnnData written to a .zarr directory (readable by anndata)."""
+    n_obs, d = 6, 8
+    rng = np.random.default_rng(0)
+    obs = pd.DataFrame(index=[f"cell_{i}" for i in range(n_obs)])
+    adata = ad.AnnData(X=rng.standard_normal((n_obs, 4)).astype(np.float32), obs=obs)
+    adata.obsm["X_test"] = rng.standard_normal((n_obs, d)).astype(np.float32)
+    zarr_path = os.path.join(tmp_dir, "chunk.zarr")
+    adata.write_zarr(zarr_path)
+    return zarr_path, n_obs, d
+
+
+class _FakeModule:
+    def __init__(self):
+        self.store = None
+
+    def set_vector_store(self, store):
+        self.store = store
+
+
+class _FakeModel:
+    """Stand-in for a SentenceTransformer: only ``model[0]`` is exercised."""
+
+    def __init__(self):
+        self._module = _FakeModule()
+
+    def __getitem__(self, idx):
+        return self._module
+
+
+# ---------------------------------------------------------------------------
+# prepare_dataset — training
+# ---------------------------------------------------------------------------
+class TestPrepareTrain:
+    def test_bimodal_columns_and_omics_anchor(self):
+        out = prepare_dataset(_train_ds(), purpose="train", modality="bimodal")
+        assert out.column_names == ["anchor", "positive", "negative_1", "negative_2"]
+        assert all(a.startswith("omics:") for a in out["anchor"])
+        # omics ids carry the original sample_idx
+        assert out["anchor"][0] == "omics:cell_0"
+
+    def test_text_anchor_is_cell_sentence(self):
+        raw = _train_ds()
+        out = prepare_dataset(raw, purpose="train", modality="text")
+        assert out.column_names == ["anchor", "positive", "negative_1", "negative_2"]
+        # anchor equals the (unprefixed) cell sentence text
+        assert out["anchor"] == raw["cell_sentence_1"]
+
+    def test_without_hard_negatives(self):
+        out = prepare_dataset(_train_ds(), purpose="train", modality="bimodal", use_hard_negatives=False)
+        assert out.column_names == ["anchor", "positive"]
+        assert all(a.startswith("omics:") for a in out["anchor"])
+
+    def test_missing_positive_raises(self):
+        ds = _test_ds()  # no positive column
+        with pytest.raises(KeyError, match="positive"):
+            prepare_dataset(ds, purpose="train", modality="bimodal")
+
+
+# ---------------------------------------------------------------------------
+# prepare_dataset — inference
+# ---------------------------------------------------------------------------
+class TestPrepareInference:
+    def test_bimodal_keeps_identifiers(self):
+        out = prepare_dataset(_test_ds(), purpose="inference", modality="bimodal")
+        assert "anchor" in out.column_names
+        assert "sample_idx" in out.column_names
+        assert "adata_link" in out.column_names
+        assert "positive" not in out.column_names
+        assert all(a.startswith("omics:") for a in out["anchor"])
+
+    def test_text_anchor(self):
+        raw = _test_ds()
+        out = prepare_dataset(raw, purpose="inference", modality="text")
+        assert out["anchor"] == raw["cell_sentence_1"]
+
+    def test_no_positive_or_negative_required(self):
+        # The whole point: a test dataset without positive/negative must not error.
+        out = prepare_dataset(_test_ds(), purpose="inference", modality="bimodal")
+        assert len(out) == 6
+
+    def test_datasetdict(self):
+        dd = DatasetDict({"test": _test_ds()})
+        out = prepare_dataset(dd, purpose="inference", modality="bimodal")
+        assert isinstance(out, DatasetDict)
+        assert all(a.startswith("omics:") for a in out["test"]["anchor"])
+
+
+# ---------------------------------------------------------------------------
+# prepare_dataset — validation
+# ---------------------------------------------------------------------------
+class TestValidation:
+    def test_bad_modality(self):
+        with pytest.raises(ValueError, match="modality"):
+            prepare_dataset(_test_ds(), purpose="inference", modality="nope")
+
+    def test_bad_purpose(self):
+        with pytest.raises(ValueError, match="purpose"):
+            prepare_dataset(_test_ds(), purpose="nope", modality="text")
+
+    def test_bimodal_requires_sample_id(self):
+        ds = Dataset.from_dict({"cell_sentence_1": ["a", "b"]})
+        with pytest.raises(KeyError, match="sample"):
+            prepare_dataset(ds, purpose="inference", modality="bimodal")
+
+
+# ---------------------------------------------------------------------------
+# prepare_inference — end-to-end with local zarr
+# ---------------------------------------------------------------------------
+class TestPrepareInferenceOrchestrator:
+    def test_bimodal_builds_store_and_subsets(self, adata_zarr, tmp_dir):
+        zarr_path, n_obs, d = adata_zarr
+        ds = _test_ds(n_obs)
+        ds = ds.remove_columns("adata_link")
+        ds = ds.add_column("adata_link", [zarr_path] * n_obs)
+
+        model = _FakeModel()
+        bundle = prepare_inference(
+            model,
+            ds,
+            modality="bimodal",
+            obsm_key="X_test",
+            cache_dir=tmp_dir,
+            store_path=os.path.join(tmp_dir, "store.mmap"),
+        )
+
+        assert isinstance(bundle, InferenceData)
+        assert bundle.vector_store is not None
+        assert bundle.vector_store.dim == d
+        # store was attached to the model's first module
+        assert model[0].store is bundle.vector_store
+        # dataset aligns with the loaded chunk
+        assert len(bundle.dataset) == bundle.adata.n_obs == n_obs
+        # every omics anchor resolves in the store
+        for anchor in bundle.dataset["anchor"]:
+            sid = anchor[len("omics:") :]
+            assert bundle.vector_store[sid].shape == (d,)
+
+    def test_text_skips_store(self, adata_zarr, tmp_dir):
+        zarr_path, n_obs, _ = adata_zarr
+        ds = _test_ds(n_obs)
+        ds = ds.remove_columns("adata_link")
+        ds = ds.add_column("adata_link", [zarr_path] * n_obs)
+
+        model = _FakeModel()
+        bundle = prepare_inference(
+            model,
+            ds,
+            modality="text",
+            cache_dir=tmp_dir,
+        )
+        assert bundle.vector_store is None
+        assert model[0].store is None
+        # text modality picks cell_sentence_2 internally
+        assert bundle.dataset["anchor"] == ds["cell_sentence_2"]
+
+    def test_bimodal_requires_obsm_key(self, tmp_dir):
+        with pytest.raises(ValueError, match="obsm_key"):
+            prepare_inference(_FakeModel(), _test_ds(), modality="bimodal", cache_dir=tmp_dir)
diff --git a/tests/test_io/conftest.py b/tests/test_io/conftest.py
new file mode 100644
index 0000000..38947fb
--- /dev/null
+++ b/tests/test_io/conftest.py
@@ -0,0 +1 @@
+"""Minimal conftest for io tests — no torch/sentence-transformers dependencies."""
diff --git a/tests/test_io/test_lean_loaders.py b/tests/test_io/test_lean_loaders.py
new file mode 100644
index 0000000..0418955
--- /dev/null
+++ b/tests/test_io/test_lean_loaders.py
@@ -0,0 +1,125 @@
+"""Tests for the opt-in memory-lean paths in the AnnData loaders.
+
+Covers:
+* :func:`mmcontext.io.collect_adata_subset` with ``obsm_key``
+* :func:`mmcontext.file_utils.load_test_adata_from_hf_dataset` with ``lean=True``
+
+Both should read only the obs table and the requested ``obsm`` layer, return a
+minimal AnnData (no ``X``), and never call ``anndata.read_zarr``.
+"""
+
+from __future__ import annotations
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import pytest
+
+from mmcontext import file_utils
+from mmcontext.file_utils import load_test_adata_from_hf_dataset
+from mmcontext.io._adata_subset import collect_adata_subset
+
+D = 8
+
+
+def _make_adata(ids: list[str], seed: int) -> ad.AnnData:
+    rng = np.random.default_rng(seed)
+    n = len(ids)
+    obs = pd.DataFrame(
+        {
+            "cell_type": pd.Categorical(rng.choice(["T", "B", "NK"], n)),
+            "score": rng.random(n).astype(np.float32),
+        },
+        index=ids,
+    )
+    adata = ad.AnnData(X=rng.random((n, 5)).astype(np.float32), obs=obs)
+    adata.obsm["X_scvi"] = rng.random((n, D)).astype(np.float32)
+    adata.obsm["X_pca"] = rng.random((n, 3)).astype(np.float32)
+    return adata
+
+
+@pytest.fixture
+def two_chunks(tmp_path):
+    """Two disjoint chunks (zarr + h5ad) plus an id→vector reference map."""
+    a1 = _make_adata([f"cell_{i}" for i in range(30)], seed=1)
+    a2 = _make_adata([f"cell_{i}" for i in range(30, 60)], seed=2)
+    p1 = tmp_path / "chunk1.zarr"
+    p2 = tmp_path / "chunk2.h5ad"
+    a1.write_zarr(p1)
+    a2.write_h5ad(p2)
+    ref = {}
+    for a in (a1, a2):
+        for i, name in enumerate(a.obs_names):
+            ref[name] = a.obsm["X_scvi"][i]
+    return [p1, p2], ref
+
+
+def test_collect_subset_lean_order_and_values(two_chunks):
+    paths, ref = two_chunks
+    ids = ["cell_45", "cell_2", "cell_58", "cell_10", "cell_31"]
+    adata = collect_adata_subset(file_paths=paths, sample_ids=ids, obsm_key="X_scvi", obs_columns=["cell_type"])
+
+    assert list(adata.obs_names) == ids
+    assert adata.n_vars == 0  # no X materialised
+    assert list(adata.obs.columns) == ["cell_type"]
+    expected = np.vstack([ref[i] for i in ids])
+    np.testing.assert_allclose(adata.obsm["X_scvi"], expected)
+
+
+def test_collect_subset_lean_matches_legacy(two_chunks):
+    paths, _ = two_chunks
+    ids = ["cell_5", "cell_50", "cell_29", "cell_30"]
+    lean = collect_adata_subset(file_paths=paths, sample_ids=ids, obsm_key="X_scvi")
+    legacy = collect_adata_subset(file_paths=paths, sample_ids=ids)
+    np.testing.assert_allclose(lean.obsm["X_scvi"], legacy.obsm["X_scvi"])
+    assert list(lean.obs_names) == list(legacy.obs_names) == ids
+
+
+def test_collect_subset_lean_does_not_read_full_zarr(two_chunks, monkeypatch):
+    paths, _ = two_chunks
+
+    def _boom(*args, **kwargs):
+        raise AssertionError("lean collect_adata_subset must not call ad.read_zarr")
+
+    from mmcontext.io import _adata_subset
+
+    monkeypatch.setattr(_adata_subset.ad, "read_zarr", _boom)
+    ids = ["cell_1", "cell_31"]
+    adata = collect_adata_subset(file_paths=paths, sample_ids=ids, obsm_key="X_scvi")
+    assert list(adata.obs_names) == ids
+
+
+@pytest.fixture
+def hf_split(tmp_path, monkeypatch):
+    """A single-chunk HF split whose download step is stubbed to a local zarr."""
+    from datasets import Dataset
+
+    adata = _make_adata([f"cell_{i}" for i in range(40)], seed=7)
+    zpath = tmp_path / "test_chunk.zarr"
+    adata.write_zarr(zpath)
+
+    url = "https://example.org/fake.zarr"
+    split = Dataset.from_dict({"adata_link": [url] * adata.n_obs, "sample_idx": list(adata.obs_names)})
+
+    def _fake_download(links, target_dir, **kwargs):
+        return {links[0]: zpath}
+
+    monkeypatch.setattr(file_utils, "download_and_extract_links", _fake_download)
+    return split, adata, tmp_path
+
+
+def test_load_test_adata_lean(hf_split):
+    split, adata, save_dir = hf_split
+    lean, _ = load_test_adata_from_hf_dataset(split, save_dir, layer_key="X_scvi", lean=True, obs_columns=["cell_type"])
+    assert lean.n_vars == 0
+    assert list(lean.obs.columns) == ["cell_type"]
+    assert list(lean.obs_names) == list(adata.obs_names)
+    np.testing.assert_allclose(lean.obsm["X_scvi"], adata.obsm["X_scvi"])
+
+
+def test_load_test_adata_lean_matches_full(hf_split):
+    split, adata, save_dir = hf_split
+    lean, _ = load_test_adata_from_hf_dataset(split, save_dir, layer_key="X_scvi", lean=True)
+    full, _ = load_test_adata_from_hf_dataset(split, save_dir, layer_key="X_scvi", lean=False)
+    np.testing.assert_allclose(lean.obsm["X_scvi"], full.obsm["X_scvi"])
+    assert list(lean.obs_names) == list(full.obs_names)
diff --git a/tests/test_io/test_vector_store.py b/tests/test_io/test_vector_store.py
new file mode 100644
index 0000000..f8f5275
--- /dev/null
+++ b/tests/test_io/test_vector_store.py
@@ -0,0 +1,318 @@
+"""Tests for VectorStore — memory-mapped vector lookup.
+
+These tests define the contract that VectorStore must satisfy:
+construction from multiple data sources, lookup, persistence, and edge cases.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import pytest
+
+from mmcontext.io import VectorStore
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def sample_matrix():
+    """A small (5, 8) float32 matrix with known values."""
+    rng = np.random.default_rng(42)
+    return rng.standard_normal((5, 8)).astype(np.float32)
+
+
+@pytest.fixture
+def sample_ids():
+    """Token IDs matching sample_matrix rows."""
+    return ["cell_A", "cell_B", "cell_C", "cell_D", "cell_E"]
+
+
+@pytest.fixture
+def sample_df(sample_matrix, sample_ids):
+    """DataFrame with 'token' and 'embedding' columns."""
+    return pd.DataFrame(
+        {
+            "token": sample_ids,
+            "embedding": [sample_matrix[i] for i in range(len(sample_ids))],
+        }
+    )
+
+
+@pytest.fixture
+def sample_adata_obs(sample_matrix, sample_ids):
+    """AnnData with obs-level embeddings in .obsm['X_scvi']."""
+    adata = ad.AnnData(
+        X=np.zeros((5, 3)),  # dummy expression matrix
+        obs=pd.DataFrame(index=sample_ids),
+    )
+    adata.obsm["X_scvi"] = sample_matrix
+    return adata
+
+
+@pytest.fixture
+def sample_adata_var():
+    """AnnData with var-level embeddings in .varm['gene_emb']."""
+    gene_ids = ["EGFR", "KRAS", "TP53", "BRCA1"]
+    rng = np.random.default_rng(99)
+    gene_matrix = rng.standard_normal((4, 16)).astype(np.float32)
+
+    adata = ad.AnnData(
+        X=np.zeros((2, 4)),
+        var=pd.DataFrame(index=gene_ids),
+    )
+    adata.varm["gene_emb"] = gene_matrix
+    return adata
+
+
+@pytest.fixture
+def tmp_dir():
+    """Temporary directory for memmap files."""
+    with tempfile.TemporaryDirectory() as d:
+        yield d
+
+
+# ---------------------------------------------------------------------------
+# Construction tests
+# ---------------------------------------------------------------------------
+class TestVectorStoreConstruction:
+    """Tests for creating VectorStore from different data sources."""
+
+    def test_from_numpy(self, sample_matrix, sample_ids, tmp_dir):
+        """Round-trip: create from numpy, values match on lookup."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+
+        for i, sid in enumerate(sample_ids):
+            result = store[sid]
+            np.testing.assert_array_almost_equal(result, sample_matrix[i])
+
+    def test_from_dataframe(self, sample_df, tmp_dir):
+        """Create from DataFrame with token/embedding columns."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_dataframe(sample_df, path=path)
+
+        for _, row in sample_df.iterrows():
+            result = store[row["token"]]
+            np.testing.assert_array_almost_equal(result, row["embedding"])
+
+    def test_from_dataframe_custom_columns(self, sample_matrix, sample_ids, tmp_dir):
+        """Create from DataFrame with custom column names."""
+        df = pd.DataFrame(
+            {
+                "my_id": sample_ids,
+                "my_vec": [sample_matrix[i] for i in range(len(sample_ids))],
+            }
+        )
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_dataframe(df, path=path, id_col="my_id", embedding_col="my_vec")
+        result = store["cell_A"]
+        np.testing.assert_array_almost_equal(result, sample_matrix[0])
+
+    def test_from_adata_obs(self, sample_adata_obs, tmp_dir):
+        """Create from adata.obsm, lookup by obs index."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_adata(sample_adata_obs, layer_key="X_scvi", axis="obs", path=path)
+        expected = sample_adata_obs.obsm["X_scvi"]
+
+        for i, sid in enumerate(sample_adata_obs.obs.index):
+            result = store[sid]
+            np.testing.assert_array_almost_equal(result, expected[i])
+
+    def test_from_adata_var(self, sample_adata_var, tmp_dir):
+        """Create from adata.varm, lookup by var index."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_adata(sample_adata_var, layer_key="gene_emb", axis="var", path=path)
+        expected = sample_adata_var.varm["gene_emb"]
+
+        for i, gid in enumerate(sample_adata_var.var.index):
+            result = store[gid]
+            np.testing.assert_array_almost_equal(result, expected[i])
+
+    def test_from_dict(self, sample_matrix, sample_ids, tmp_dir):
+        """Create from {id: vector} mapping."""
+        mapping = {sid: sample_matrix[i] for i, sid in enumerate(sample_ids)}
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_dict(mapping, path=path)
+
+        for sid, vec in mapping.items():
+            np.testing.assert_array_almost_equal(store[sid], vec)
+
+
+# ---------------------------------------------------------------------------
+# Lookup tests
+# ---------------------------------------------------------------------------
+class TestVectorStoreLookup:
+    """Tests for looking up vectors by ID."""
+
+    def test_single_lookup(self, sample_matrix, sample_ids, tmp_dir):
+        """Single ID lookup returns 1-D array."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+
+        result = store["cell_C"]
+        assert result.ndim == 1
+        assert result.shape == (8,)
+        np.testing.assert_array_almost_equal(result, sample_matrix[2])
+
+    def test_batch_lookup(self, sample_matrix, sample_ids, tmp_dir):
+        """Batch of IDs returns correct (N, D) array."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+
+        ids = ["cell_E", "cell_A", "cell_C"]
+        result = store.batch_lookup(ids)
+        assert result.shape == (3, 8)
+        np.testing.assert_array_almost_equal(result[0], sample_matrix[4])
+        np.testing.assert_array_almost_equal(result[1], sample_matrix[0])
+        np.testing.assert_array_almost_equal(result[2], sample_matrix[2])
+
+    def test_unknown_id_raises(self, sample_matrix, sample_ids, tmp_dir):
+        """KeyError for missing IDs."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+
+        with pytest.raises(KeyError, match="unknown_cell"):
+            store["unknown_cell"]
+
+    def test_batch_lookup_unknown_raises(self, sample_matrix, sample_ids, tmp_dir):
+        """KeyError in batch lookup for missing IDs."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+
+        with pytest.raises(KeyError):
+            store.batch_lookup(["cell_A", "nonexistent"])
+
+    def test_duplicate_ids_rejected(self, sample_matrix, tmp_dir):
+        """Duplicate token IDs are rejected at construction time."""
+        ids_with_dup = ["cell_A", "cell_B", "cell_A", "cell_D", "cell_E"]
+        path = os.path.join(tmp_dir, "test.mmap")
+
+        with pytest.raises(ValueError, match="[Dd]uplicate"):
+            VectorStore.from_numpy(sample_matrix, ids_with_dup, path=path)
+
+
+# ---------------------------------------------------------------------------
+# Properties tests
+# ---------------------------------------------------------------------------
+class TestVectorStoreProperties:
+    """Tests for VectorStore metadata properties."""
+
+    def test_dim(self, sample_matrix, sample_ids, tmp_dir):
+        """Reports correct embedding dimension."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+        assert store.dim == 8
+
+    def test_dtype(self, sample_matrix, sample_ids, tmp_dir):
+        """Reports correct dtype."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+        assert store.dtype == np.float32
+
+    def test_len(self, sample_matrix, sample_ids, tmp_dir):
+        """Reports correct number of stored vectors."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+        assert len(store) == 5
+
+    def test_contains(self, sample_matrix, sample_ids, tmp_dir):
+        """Supports 'in' operator for checking IDs."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+        assert "cell_A" in store
+        assert "nonexistent" not in store
+
+
+# ---------------------------------------------------------------------------
+# Persistence tests
+# ---------------------------------------------------------------------------
+class TestVectorStorePersistence:
+    """Tests for save/load and memmap persistence."""
+
+    def test_persistence_across_reopen(self, sample_matrix, sample_ids, tmp_dir):
+        """Store survives close/reopen cycle."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+
+        # Delete the Python object
+        del store
+
+        # Reopen from the same path
+        store2 = VectorStore.load(path)
+        assert len(store2) == 5
+        assert store2.dim == 8
+        np.testing.assert_array_almost_equal(store2["cell_A"], sample_matrix[0])
+        np.testing.assert_array_almost_equal(store2["cell_E"], sample_matrix[4])
+
+    def test_load_nonexistent_raises(self):
+        """Loading from non-existent path raises FileNotFoundError."""
+        with pytest.raises(FileNotFoundError):
+            VectorStore.load("/nonexistent/path/store.mmap")
+
+    def test_files_created_on_disk(self, sample_matrix, sample_ids, tmp_dir):
+        """Construction creates memmap file and index file on disk."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        VectorStore.from_numpy(sample_matrix, sample_ids, path=path)
+
+        # Memmap data file should exist
+        assert os.path.isfile(path)
+        # Index file should exist alongside
+        index_path = path + ".index.json"
+        assert os.path.isfile(index_path)
+
+
+# ---------------------------------------------------------------------------
+# Edge cases
+# ---------------------------------------------------------------------------
+class TestVectorStoreEdgeCases:
+    """Tests for edge cases and special inputs."""
+
+    def test_single_vector(self, tmp_dir):
+        """Store with a single vector works correctly."""
+        matrix = np.array([[1.0, 2.0, 3.0]], dtype=np.float32)
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(matrix, ["only_one"], path=path)
+
+        assert len(store) == 1
+        assert store.dim == 3
+        np.testing.assert_array_almost_equal(store["only_one"], [1.0, 2.0, 3.0])
+
+    def test_large_dimension(self, tmp_dir):
+        """Handles high-dimensional vectors (gs10k-like)."""
+        rng = np.random.default_rng(7)
+        matrix = rng.standard_normal((10, 10_000)).astype(np.float32)
+        ids = [f"cell_{i}" for i in range(10)]
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(matrix, ids, path=path)
+
+        assert store.dim == 10_000
+        np.testing.assert_array_almost_equal(store["cell_5"], matrix[5])
+
+    def test_float16_dtype(self, tmp_dir):
+        """Supports float16 for memory efficiency."""
+        matrix = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float16)
+        path = os.path.join(tmp_dir, "test.mmap")
+        store = VectorStore.from_numpy(matrix, ["a", "b"], path=path)
+
+        assert store.dtype == np.float16
+        np.testing.assert_array_almost_equal(store["a"], [1.0, 2.0], decimal=3)
+
+    def test_empty_ids_raises(self, tmp_dir):
+        """Empty ID list is rejected."""
+        matrix = np.empty((0, 8), dtype=np.float32)
+        path = os.path.join(tmp_dir, "test.mmap")
+
+        with pytest.raises(ValueError, match="[Ee]mpty"):
+            VectorStore.from_numpy(matrix, [], path=path)
+
+    def test_mismatched_ids_matrix_raises(self, sample_matrix, tmp_dir):
+        """Mismatched number of IDs and matrix rows raises ValueError."""
+        path = os.path.join(tmp_dir, "test.mmap")
+        with pytest.raises(ValueError, match="[Mm]ismatch|[Ll]ength"):
+            VectorStore.from_numpy(sample_matrix, ["a", "b"], path=path)  # 5 rows, 2 IDs
diff --git a/tests/test_io/test_zarr_read.py b/tests/test_io/test_zarr_read.py
new file mode 100644
index 0000000..c673276
--- /dev/null
+++ b/tests/test_io/test_zarr_read.py
@@ -0,0 +1,110 @@
+"""Tests for the memory-lean zarr/h5ad readers in ``mmcontext.io._zarr_read``.
+
+These verify that :func:`read_obs_and_obsm` reads only the obs table and the
+requested ``obsm`` layer (optionally row-subset), preserves categoricals and
+ordering, and never materialises the full store via ``anndata.read_zarr``.
+"""
+
+from __future__ import annotations
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import pytest
+
+from mmcontext.io._zarr_read import read_obs_and_obsm
+
+N, G, D = 60, 15, 8
+
+
+@pytest.fixture
+def synthetic_adata() -> ad.AnnData:
+    rng = np.random.default_rng(0)
+    obs = pd.DataFrame(
+        {
+            "cell_type": pd.Categorical(rng.choice(["T", "B", "NK"], N)),
+            "tissue": rng.choice(["lung", "liver"], N),
+            "score": rng.random(N).astype(np.float32),
+        },
+        index=[f"cell_{i}" for i in range(N)],
+    )
+    adata = ad.AnnData(X=rng.random((N, G)).astype(np.float32), obs=obs)
+    adata.obsm["X_scvi"] = rng.random((N, D)).astype(np.float32)
+    adata.obsm["X_pca"] = rng.random((N, 4)).astype(np.float32)
+    return adata
+
+
+@pytest.fixture(params=["zarr", "h5ad"])
+def store(request, synthetic_adata, tmp_path) -> tuple[str, ad.AnnData]:
+    """Write the synthetic AnnData as zarr or h5ad and return its path."""
+    if request.param == "zarr":
+        path = tmp_path / "chunk.zarr"
+        synthetic_adata.write_zarr(path)
+    else:
+        path = tmp_path / "chunk.h5ad"
+        synthetic_adata.write_h5ad(path)
+    return str(path), synthetic_adata
+
+
+def test_full_read_matches_source(store):
+    path, adata = store
+    obs, obsm = read_obs_and_obsm(path, obsm_key="X_scvi")
+    assert list(obs.index) == list(adata.obs.index)
+    assert obsm.dtype == np.float32
+    np.testing.assert_allclose(obsm, adata.obsm["X_scvi"])
+
+
+def test_row_subset_preserves_order(store):
+    path, adata = store
+    rows = np.array([7, 3, 40, 0, 25])
+    obs, obsm = read_obs_and_obsm(path, obsm_key="X_scvi", rows=rows)
+    assert list(obs.index) == [f"cell_{i}" for i in rows]
+    np.testing.assert_allclose(obsm, adata.obsm["X_scvi"][rows])
+
+
+def test_obs_columns_subset_and_missing_ignored(store):
+    path, _ = store
+    obs, _ = read_obs_and_obsm(path, obsm_key="X_scvi", obs_columns=["cell_type", "score", "absent"])
+    assert list(obs.columns) == ["cell_type", "score"]
+
+
+def test_categorical_preserved(store):
+    path, adata = store
+    rows = np.array([5, 1, 30])
+    obs, _ = read_obs_and_obsm(path, obsm_key="X_scvi", rows=rows, obs_columns=["cell_type"])
+    assert isinstance(obs["cell_type"].dtype, pd.CategoricalDtype)
+    assert list(obs["cell_type"]) == list(adata.obs["cell_type"].iloc[rows])
+
+
+def test_no_obsm_key_returns_none(store):
+    path, _ = store
+    obs, obsm = read_obs_and_obsm(path, obsm_key=None)
+    assert obsm is None
+    assert len(obs) == N
+
+
+def test_empty_rows(store):
+    path, _ = store
+    obs, obsm = read_obs_and_obsm(path, obsm_key="X_scvi", rows=np.array([], dtype=int))
+    assert obsm.shape == (0, D)
+    assert len(obs) == 0
+
+
+def test_missing_obsm_key_raises(store):
+    path, _ = store
+    with pytest.raises(KeyError):
+        read_obs_and_obsm(path, obsm_key="does_not_exist")
+
+
+def test_zarr_path_is_lazy(synthetic_adata, tmp_path, monkeypatch):
+    """The zarr lean read must never call ``anndata.read_zarr`` (which would
+    materialise the whole store)."""
+    path = tmp_path / "chunk.zarr"
+    synthetic_adata.write_zarr(path)
+
+    def _boom(*args, **kwargs):
+        raise AssertionError("read_obs_and_obsm must not call anndata.read_zarr")
+
+    monkeypatch.setattr(ad, "read_zarr", _boom)
+    obs, obsm = read_obs_and_obsm(str(path), obsm_key="X_scvi", rows=np.array([1, 2, 3]))
+    assert obsm.shape == (3, D)
diff --git a/tests/test_miniOmicsEncoder.py b/tests/test_miniOmicsEncoder.py
index cd1dc8a..e4caba9 100644
--- a/tests/test_miniOmicsEncoder.py
+++ b/tests/test_miniOmicsEncoder.py
@@ -4,7 +4,7 @@
 import torch
 from torch import nn
 
-from mmcontext.omicsencoder import MiniOmicsModel
+from mmcontext._legacy.omicsencoder import MiniOmicsModel
 
 
 # ---------------------------------------------------------------------- #
diff --git a/tests/test_mmcontext_encoder.py b/tests/test_mmcontext_encoder.py
index c8510c2..430433d 100644
--- a/tests/test_mmcontext_encoder.py
+++ b/tests/test_mmcontext_encoder.py
@@ -5,7 +5,7 @@
 import pytest
 import torch
 
-from mmcontext.mmcontextencoder import AdapterModule, MMContextEncoder
+from mmcontext._legacy.mmcontextencoder import AdapterModule, MMContextEncoder
 
 
 # --------------------------------------------------------------------- #
@@ -791,7 +791,7 @@ def _rand_matrix(tokens, hidden=8, seed=0):
         )
     # test what happens if dimensions are inconsistent
     bad_df3 = pd.DataFrame({"sample_idx": list(mapping.keys()), "embedding": list(mapping.values())})
-    bad_df3["embedding"][0] = np.ones(3, dtype=np.float32)
+    bad_df3.at[0, "embedding"] = np.ones(3, dtype=np.float32)
     # Should raise ValueError due to inconsistent dimensions
     with pytest.raises(ValueError, match="dimensions"):
         text_only_encoder.register_initial_embeddings(
diff --git a/tests/test_mmcontext_module.py b/tests/test_mmcontext_module.py
new file mode 100644
index 0000000..8a15fc6
--- /dev/null
+++ b/tests/test_mmcontext_module.py
@@ -0,0 +1,422 @@
+"""Tests for MMContextModule — the core InputModule for multimodal encoding.
+
+These tests define the contract that MMContextModule must satisfy.
+They are written FIRST (TDD) and drive the implementation.
+
+MMContextModule extends sentence-transformers InputModule (v5.4+) and serves
+as the first module in the ST pipeline. It handles:
+  - Text preprocessing via AutoTokenizer
+  - Omics vector pass-through (direct or via VectorStore lookup)
+  - Producing a unified features dict with token_embeddings, attention_mask,
+    and modality_ids for downstream modules (AdapterModule, Pooling).
+
+The v5.4 API uses ``preprocess()`` instead of the deprecated ``tokenize()``.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+
+from mmcontext.io import VectorStore
+from mmcontext.modules import MMContextModule
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def tmp_dir():
+    """Temporary directory for test artifacts."""
+    with tempfile.TemporaryDirectory() as d:
+        yield d
+
+
+@pytest.fixture
+def omics_store(tmp_dir):
+    """A small VectorStore with 5 cells, dim=8."""
+    rng = np.random.default_rng(42)
+    matrix = rng.standard_normal((5, 8)).astype(np.float32)
+    ids = ["cell_A", "cell_B", "cell_C", "cell_D", "cell_E"]
+    path = os.path.join(tmp_dir, "test_store.mmap")
+    return VectorStore.from_numpy(matrix, ids, path=path)
+
+
+@pytest.fixture
+def module():
+    """MMContextModule backed by stub text encoder (patched globally in conftest).
+
+    The session-scoped ``patch_model_loading`` fixture in conftest.py patches
+    ``AutoModel.from_pretrained`` and ``AutoTokenizer.from_pretrained`` to
+    return lightweight stubs, so no real HF download occurs.
+    """
+    return MMContextModule(model_name_or_path="bert-base-uncased")
+
+
+@pytest.fixture
+def module_with_store(module, omics_store):
+    """Module with a VectorStore attached."""
+    module.set_vector_store(omics_store)
+    return module
+
+
+# ---------------------------------------------------------------------------
+# Preprocess tests — text
+# ---------------------------------------------------------------------------
+class TestPreprocessText:
+    """Tests for preprocess() with text-only inputs."""
+
+    def test_preprocess_text_returns_required_keys(self, module):
+        """Text preprocessing produces input_ids, attention_mask, and modality marker."""
+        features = module.preprocess(["This is a test sentence."])
+
+        assert "input_ids" in features
+        assert "attention_mask" in features
+        assert features["modality"] == "text"
+
+    def test_preprocess_text_batch(self, module):
+        """Batch of text strings produces correct batch dimension."""
+        texts = ["First sentence.", "Second sentence.", "Third one."]
+        features = module.preprocess(texts)
+
+        assert features["input_ids"].shape[0] == 3
+        assert features["attention_mask"].shape[0] == 3
+
+    def test_preprocess_text_tensor_types(self, module):
+        """Preprocessed text features are torch tensors."""
+        features = module.preprocess(["Hello world."])
+
+        assert isinstance(features["input_ids"], torch.Tensor)
+        assert isinstance(features["attention_mask"], torch.Tensor)
+
+    def test_preprocess_text_with_prompt(self, module):
+        """Optional prompt is prepended to text inputs."""
+        module.preprocess(["Test."])
+        features_with_prompt = module.preprocess(["Test."], prompt="Query: ")
+
+        # Both should produce valid features; the prompt version processes
+        # "Query: Test." instead of "Test."
+        assert "input_ids" in features_with_prompt
+        assert features_with_prompt["modality"] == "text"
+
+
+# ---------------------------------------------------------------------------
+# Preprocess tests — omics via VectorStore
+# ---------------------------------------------------------------------------
+class TestPreprocessOmicsViaStore:
+    """Tests for preprocess() with omics IDs resolved through VectorStore."""
+
+    def test_preprocess_omics_ids_returns_embeddings(self, module_with_store):
+        """Prefixed omics IDs are resolved to vectors via VectorStore."""
+        inputs = ["omics:cell_A", "omics:cell_B"]
+        features = module_with_store.preprocess(inputs)
+
+        assert "input_values" in features
+        assert features["modality"] == "omics"
+        # Each cell maps to a single vector → (B, 1, D)
+        assert features["input_values"].shape == (2, 1, 8)
+
+    def test_preprocess_omics_ids_values_match_store(self, module_with_store, omics_store):
+        """Resolved vectors match the VectorStore content."""
+        inputs = ["omics:cell_C"]
+        features = module_with_store.preprocess(inputs)
+
+        expected = omics_store["cell_C"]
+        actual = features["input_values"][0, 0].numpy()
+        np.testing.assert_array_almost_equal(actual, expected)
+
+    def test_preprocess_omics_ids_attention_mask(self, module_with_store):
+        """Omics tokens get attention_mask=1."""
+        inputs = ["omics:cell_A", "omics:cell_B"]
+        features = module_with_store.preprocess(inputs)
+
+        assert "attention_mask" in features
+        assert features["attention_mask"].shape == (2, 1)
+        assert (features["attention_mask"] == 1).all()
+
+    def test_preprocess_omics_unknown_id_raises(self, module_with_store):
+        """Unknown omics ID raises KeyError."""
+        with pytest.raises(KeyError, match="unknown_cell"):
+            module_with_store.preprocess(["omics:unknown_cell"])
+
+
+# ---------------------------------------------------------------------------
+# Preprocess tests — omics direct vectors
+# ---------------------------------------------------------------------------
+class TestPreprocessOmicsDirect:
+    """Tests for preprocess() with direct omics vectors (via dict input)."""
+
+    def test_preprocess_direct_single_vector(self, module):
+        """Dict with omics_values (1-D array) → single omics token per sample."""
+        inputs = [
+            {"omics_values": np.array([1.0, 2.0, 3.0], dtype=np.float32)},
+            {"omics_values": np.array([4.0, 5.0, 6.0], dtype=np.float32)},
+        ]
+        features = module.preprocess(inputs)
+
+        assert features["modality"] == "omics"
+        assert features["input_values"].shape == (2, 1, 3)
+        assert features["attention_mask"].shape == (2, 1)
+
+    def test_preprocess_direct_values_preserved(self, module):
+        """Direct vectors appear unchanged in input_values."""
+        vec = np.array([1.5, -2.5, 3.5], dtype=np.float32)
+        features = module.preprocess([{"omics_values": vec}])
+
+        actual = features["input_values"][0, 0].numpy()
+        np.testing.assert_array_almost_equal(actual, vec)
+
+    def test_preprocess_direct_var_multiple_vectors(self, module):
+        """List of gene vectors → variable-length omics sequence (var case)."""
+        gene_vecs_1 = [
+            np.array([1.0, 2.0], dtype=np.float32),
+            np.array([3.0, 4.0], dtype=np.float32),
+            np.array([5.0, 6.0], dtype=np.float32),
+        ]
+        gene_vecs_2 = [
+            np.array([7.0, 8.0], dtype=np.float32),
+            np.array([9.0, 10.0], dtype=np.float32),
+        ]
+        inputs = [
+            {"omics_values": gene_vecs_1},  # 3 genes
+            {"omics_values": gene_vecs_2},  # 2 genes
+        ]
+        features = module.preprocess(inputs)
+
+        # Padded to max length in batch → (2, 3, 2)
+        assert features["input_values"].shape == (2, 3, 2)
+        # Attention mask reflects real vs padded tokens
+        assert features["attention_mask"][0].tolist() == [1, 1, 1]
+        assert features["attention_mask"][1].tolist() == [1, 1, 0]
+
+
+# ---------------------------------------------------------------------------
+# Preprocess tests — error cases
+# ---------------------------------------------------------------------------
+class TestPreprocessErrors:
+    """Tests for preprocess() error handling."""
+
+    def test_no_store_omics_id_raises(self, module):
+        """Prefixed omics IDs without a VectorStore → clear error."""
+        with pytest.raises(ValueError, match="[Vv]ector[Ss]tore|[Nn]o.*store"):
+            module.preprocess(["omics:cell_A"])
+
+    def test_empty_input_raises(self, module):
+        """Empty input list raises ValueError."""
+        with pytest.raises((ValueError, RuntimeError)):
+            module.preprocess([])
+
+
+# ---------------------------------------------------------------------------
+# Forward tests
+# ---------------------------------------------------------------------------
+class TestForward:
+    """Tests for forward() producing the features dict contract."""
+
+    def test_forward_text_produces_token_embeddings(self, module):
+        """Text forward: input_ids → token_embeddings via text encoder."""
+        features = module.preprocess(["A test sentence."])
+        result = module.forward(features)
+
+        assert "token_embeddings" in result
+        assert "attention_mask" in result
+        assert "modality_ids" in result
+
+    def test_forward_text_shapes(self, module):
+        """Text forward produces correct shapes: (B, L, D)."""
+        features = module.preprocess(["Hello.", "World."])
+        result = module.forward(features)
+
+        B = 2
+        L = features["input_ids"].shape[1]  # sequence length from tokenizer
+        D = module.get_word_embedding_dimension()
+
+        assert result["token_embeddings"].shape == (B, L, D)
+        assert result["attention_mask"].shape == (B, L)
+        assert result["modality_ids"].shape == (B, L)
+
+    def test_forward_text_modality_ids_zero(self, module):
+        """Text tokens get modality_id=0."""
+        features = module.preprocess(["Test."])
+        result = module.forward(features)
+
+        assert (result["modality_ids"] == 0).all()
+
+    def test_forward_omics_passthrough(self, module_with_store):
+        """Omics forward: input_values pass through to token_embeddings unchanged."""
+        features = module_with_store.preprocess(["omics:cell_A"])
+        input_embeddings = features["input_values"].clone()
+        result = module_with_store.forward(features)
+
+        assert "token_embeddings" in result
+        torch.testing.assert_close(result["token_embeddings"], input_embeddings)
+
+    def test_forward_omics_modality_ids_one(self, module_with_store):
+        """Omics tokens get modality_id=1."""
+        features = module_with_store.preprocess(["omics:cell_A"])
+        result = module_with_store.forward(features)
+
+        assert (result["modality_ids"] == 1).all()
+
+    def test_forward_omics_shapes(self, module):
+        """Omics forward produces correct shapes."""
+        vec = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
+        features = module.preprocess([{"omics_values": vec}])
+        result = module.forward(features)
+
+        assert result["token_embeddings"].shape == (1, 1, 4)
+        assert result["attention_mask"].shape == (1, 1)
+        assert result["modality_ids"].shape == (1, 1)
+
+    def test_forward_returns_dict(self, module):
+        """Forward returns a dict (pipeline contract)."""
+        features = module.preprocess(["Test sentence."])
+        result = module.forward(features)
+
+        assert isinstance(result, dict)
+
+
+# ---------------------------------------------------------------------------
+# Properties tests
+# ---------------------------------------------------------------------------
+class TestProperties:
+    """Tests for module properties and metadata."""
+
+    def test_modalities_property(self, module):
+        """Module declares text and omics modalities."""
+        assert hasattr(module, "modalities")
+        assert "text" in module.modalities
+        assert "omics" in module.modalities
+
+    def test_max_seq_length(self, module):
+        """max_seq_length is accessible and returns a positive int."""
+        assert hasattr(module, "max_seq_length")
+        assert isinstance(module.max_seq_length, int)
+        assert module.max_seq_length > 0
+
+    def test_max_seq_length_settable(self, module):
+        """max_seq_length can be set to a new value."""
+        module.max_seq_length = 256
+        assert module.max_seq_length == 256
+
+    def test_get_word_embedding_dimension(self, module):
+        """Reports the text encoder's hidden size."""
+        dim = module.get_word_embedding_dimension()
+        assert isinstance(dim, int)
+        assert dim > 0
+
+    def test_vector_store_management(self, module, omics_store):
+        """set_vector_store and remove_vector_store work correctly."""
+        assert module._vector_store is None
+
+        module.set_vector_store(omics_store)
+        assert module._vector_store is omics_store
+
+        module.remove_vector_store()
+        assert module._vector_store is None
+
+
+# ---------------------------------------------------------------------------
+# Text encoder freezing tests
+# ---------------------------------------------------------------------------
+class TestFreezing:
+    """Tests for freezing/unfreezing the text encoder."""
+
+    def test_unfreeze_text_encoder(self, module):
+        """Unfreezing restores gradient computation after a full freeze."""
+        module.freeze_all_but_top_layers(0)
+        module.unfreeze_text_encoder()
+        for param in module.auto_model.parameters():
+            assert param.requires_grad
+
+    def test_freeze_all_but_top_layers_keeps_only_top(self, module):
+        """freeze_all_but_top_layers(N) trains only the top N layers."""
+        layers = module._get_encoder_layers()
+        assert layers is not None and len(layers) >= 2, "stub encoder needs >=2 layers"
+
+        module.freeze_all_but_top_layers(1)
+
+        # The last layer is trainable; all earlier layers are frozen.
+        for param in layers[-1].parameters():
+            assert param.requires_grad
+        for layer in layers[:-1]:
+            for param in layer.parameters():
+                assert not param.requires_grad
+
+    def test_freeze_all_but_top_zero_freezes_everything(self, module):
+        """num_trainable_layers=0 freezes the whole encoder."""
+        module.freeze_all_but_top_layers(0)
+        for param in module.auto_model.parameters():
+            assert not param.requires_grad
+
+    def test_freeze_all_but_top_n_exceeds_total_keeps_all_layers(self, module):
+        """Asking for more layers than exist keeps every encoder layer trainable."""
+        layers = module._get_encoder_layers()
+        module.freeze_all_but_top_layers(len(layers) + 5)
+        for layer in layers:
+            for param in layer.parameters():
+                assert param.requires_grad
+
+
+# ---------------------------------------------------------------------------
+# Persistence tests
+# ---------------------------------------------------------------------------
+class TestPersistence:
+    """Tests for save/load roundtrip."""
+
+    def test_save_creates_files(self, module, tmp_dir):
+        """save() creates config and weight files on disk."""
+        module.save(tmp_dir)
+
+        # Config file should exist
+        config_path = os.path.join(tmp_dir, module.config_file_name)
+        assert os.path.isfile(config_path)
+
+    def test_save_load_roundtrip(self, module, tmp_dir):
+        """Config values survive save/load cycle."""
+        original_dim = module.get_word_embedding_dimension()
+        original_seq_len = module.max_seq_length
+
+        module.save(tmp_dir)
+        loaded = MMContextModule.load(tmp_dir)
+
+        assert loaded.get_word_embedding_dimension() == original_dim
+        assert loaded.max_seq_length == original_seq_len
+
+    def test_save_load_forward_text(self, module, tmp_dir):
+        """Loaded module can run forward on text."""
+        module.save(tmp_dir)
+        loaded = MMContextModule.load(tmp_dir)
+
+        features = loaded.preprocess(["Test after reload."])
+        result = loaded.forward(features)
+
+        assert "token_embeddings" in result
+        assert result["token_embeddings"].shape[0] == 1
+
+    def test_save_excludes_vector_store(self, module_with_store):
+        """VectorStore data is NOT saved with the model weights."""
+        # Use a separate directory for saving (not the one containing the store)
+        with tempfile.TemporaryDirectory() as save_dir:
+            module_with_store.save(save_dir)
+
+            # No memmap files should appear in the save directory
+            saved_files = os.listdir(save_dir)
+            mmap_files = [f for f in saved_files if f.endswith(".mmap")]
+            assert len(mmap_files) == 0
+
+            # Loaded module should not have a vector store attached
+            loaded = MMContextModule.load(save_dir)
+            assert loaded._vector_store is None
+
+    def test_load_preserves_modalities(self, module, tmp_dir):
+        """Loaded module still reports correct modalities."""
+        module.save(tmp_dir)
+        loaded = MMContextModule.load(tmp_dir)
+
+        assert "text" in loaded.modalities
+        assert "omics" in loaded.modalities
diff --git a/tests/test_mmcontext_tokenizer.py b/tests/test_mmcontext_tokenizer.py
index 5d16eef..aa86987 100644
--- a/tests/test_mmcontext_tokenizer.py
+++ b/tests/test_mmcontext_tokenizer.py
@@ -6,7 +6,7 @@
 import pytest
 import torch
 
-from mmcontext.mmcontextencoder import MMContextEncoder, MMContextProcessor
+from mmcontext._legacy.mmcontextencoder import MMContextEncoder, MMContextProcessor
 
 
 # --- fixtures ---------------------------------------------------------
diff --git a/tests/test_omics_attention_module.py b/tests/test_omics_attention_module.py
new file mode 100644
index 0000000..ebb5212
--- /dev/null
+++ b/tests/test_omics_attention_module.py
@@ -0,0 +1,462 @@
+"""Tests for OmicsAttentionModule — optional self-attention for omics tokens.
+
+The OmicsAttentionModule sits between MMContextModule and AdapterModule in the
+sentence-transformers pipeline. It applies multi-head self-attention ONLY to
+omics tokens (modality_id=1), leaving text tokens (modality_id=0) unchanged.
+This is useful for var-level models where each sample has multiple gene
+embeddings that benefit from contextual mixing before projection.
+
+These tests define the contract and are written FIRST (TDD).
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import pytest
+import torch
+
+from mmcontext.modules.omics_attention_module import OmicsAttentionModule
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def tmp_dir():
+    with tempfile.TemporaryDirectory() as d:
+        yield d
+
+
+@pytest.fixture
+def real_safetensors():
+    """Undo the global safetensors.torch.load_model patch for persistence tests."""
+    import importlib
+
+    import safetensors.torch
+
+    importlib.reload(safetensors.torch)
+    yield
+
+
+@pytest.fixture
+def module():
+    """Default module: input_dim=16, num_heads=2, num_layers=1."""
+    return OmicsAttentionModule(
+        input_dim=16,
+        num_heads=2,
+        num_layers=1,
+        dropout=0.0,
+    )
+
+
+@pytest.fixture
+def deep_module():
+    """Deeper module for testing multi-layer stacking."""
+    return OmicsAttentionModule(
+        input_dim=16,
+        num_heads=2,
+        num_layers=3,
+        dropout=0.0,
+    )
+
+
+def _make_features(
+    token_embeddings: torch.Tensor,
+    attention_mask: torch.Tensor,
+    modality_ids: torch.Tensor,
+) -> dict[str, torch.Tensor]:
+    """Helper to build a features dict matching the pipeline contract."""
+    return {
+        "token_embeddings": token_embeddings,
+        "attention_mask": attention_mask,
+        "modality_ids": modality_ids,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Text passthrough
+# ---------------------------------------------------------------------------
+class TestTextPassthrough:
+    """Text tokens must pass through completely unchanged."""
+
+    def test_text_passthrough(self, module):
+        """Text tokens (modality_id=0) are identical before and after the module."""
+        B, L = 2, 5
+        x = torch.randn(B, L, 16)
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=torch.zeros(B, L, dtype=torch.long),  # all text
+        )
+        module.eval()
+        result = module(features)
+
+        torch.testing.assert_close(result["token_embeddings"], x)
+
+    def test_text_passthrough_in_mixed_batch(self, module):
+        """Text tokens in a mixed batch remain unchanged."""
+        B, L = 1, 6
+        x = torch.randn(B, L, 16)
+        modality_ids = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.long)
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=modality_ids,
+        )
+        module.eval()
+        result = module(features)
+
+        # Text tokens (first 3) should be unchanged
+        torch.testing.assert_close(result["token_embeddings"][:, :3, :], x[:, :3, :])
+
+
+# ---------------------------------------------------------------------------
+# Omics transformation
+# ---------------------------------------------------------------------------
+class TestOmicsTransformed:
+    """Omics tokens should be modified by self-attention."""
+
+    def test_omics_transformed(self, module):
+        """Omics tokens (modality_id=1) are modified by self-attention."""
+        B, L = 1, 4
+        x = torch.randn(B, L, 16)
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=torch.ones(B, L, dtype=torch.long),  # all omics
+        )
+        module.eval()
+        result = module(features)
+
+        # Output should differ from input (attention mixes information)
+        assert not torch.allclose(result["token_embeddings"], x, atol=1e-6), (
+            "Omics tokens should be transformed by self-attention"
+        )
+
+    def test_omics_transformed_in_mixed_batch(self, module):
+        """Omics tokens in a mixed batch are modified."""
+        B, L = 1, 6
+        x = torch.randn(B, L, 16)
+        modality_ids = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.long)
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=modality_ids,
+        )
+        module.eval()
+        result = module(features)
+
+        # Omics tokens (last 3) should be modified
+        assert not torch.allclose(result["token_embeddings"][:, 3:, :], x[:, 3:, :], atol=1e-6), (
+            "Omics tokens should be transformed by self-attention"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Attention mask respected
+# ---------------------------------------------------------------------------
+class TestAttentionMask:
+    """Padded positions must not influence real tokens."""
+
+    def test_attention_mask_respected(self, module):
+        """Padded omics tokens don't affect real omics tokens.
+
+        We verify by comparing outputs with and without a pad token — the
+        real tokens' representations should differ when the pad is replaced
+        by a real token (showing attention sees it), and should NOT change
+        when the pad is properly masked.
+        """
+        module.eval()
+        D = 16
+
+        # Batch element with 3 real omics tokens + 1 pad
+        x = torch.randn(1, 4, D)
+        mask_with_pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.long)
+        modality_ids = torch.tensor([[1, 1, 1, 2]], dtype=torch.long)
+
+        features_padded = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=mask_with_pad,
+            modality_ids=modality_ids,
+        )
+        out_padded = module(features_padded)["token_embeddings"][:, :3, :]
+
+        # Same real tokens but with a different value in the pad position
+        x2 = x.clone()
+        x2[0, 3, :] = torch.randn(D) * 100  # very different pad content
+
+        features_padded2 = _make_features(
+            token_embeddings=x2,
+            attention_mask=mask_with_pad,
+            modality_ids=modality_ids,
+        )
+        out_padded2 = module(features_padded2)["token_embeddings"][:, :3, :]
+
+        # Real tokens should produce same output regardless of pad content
+        torch.testing.assert_close(out_padded, out_padded2, atol=1e-5, rtol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# Variable-length sequences
+# ---------------------------------------------------------------------------
+class TestVariableLengthSequences:
+    """Different-length gene sequences in the same batch."""
+
+    def test_variable_length_sequences(self, module):
+        """Batch with different numbers of omics tokens per sample."""
+        module.eval()
+        B, L, D = 2, 5, 16
+        x = torch.randn(B, L, D)
+
+        # Sample 0: 3 omics tokens + 2 pad
+        # Sample 1: 5 omics tokens + 0 pad
+        modality_ids = torch.tensor(
+            [
+                [1, 1, 1, 2, 2],
+                [1, 1, 1, 1, 1],
+            ],
+            dtype=torch.long,
+        )
+        attention_mask = torch.tensor(
+            [
+                [1, 1, 1, 0, 0],
+                [1, 1, 1, 1, 1],
+            ],
+            dtype=torch.long,
+        )
+
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=attention_mask,
+            modality_ids=modality_ids,
+        )
+        result = module(features)
+
+        # Output shape preserved
+        assert result["token_embeddings"].shape == (B, L, D)
+
+        # Pad positions should remain zero (or unchanged)
+        # The module should not produce non-zero values for pad positions
+        assert torch.all(result["token_embeddings"][0, 3:, :] == 0) or torch.allclose(
+            result["token_embeddings"][0, 3:, :], x[0, 3:, :]
+        )
+
+
+# ---------------------------------------------------------------------------
+# Output shape preserved
+# ---------------------------------------------------------------------------
+class TestOutputShape:
+    """Output shape must be identical to input shape."""
+
+    def test_output_shape_preserved(self, module):
+        """(B, L, D) shape is unchanged after the module."""
+        B, L, D = 3, 7, 16
+        features = _make_features(
+            token_embeddings=torch.randn(B, L, D),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=torch.ones(B, L, dtype=torch.long),
+        )
+        result = module(features)
+        assert result["token_embeddings"].shape == (B, L, D)
+
+    def test_output_shape_with_deep_module(self, deep_module):
+        """Shape preserved through multiple attention layers."""
+        B, L, D = 2, 4, 16
+        features = _make_features(
+            token_embeddings=torch.randn(B, L, D),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=torch.ones(B, L, dtype=torch.long),
+        )
+        result = deep_module(features)
+        assert result["token_embeddings"].shape == (B, L, D)
+
+    def test_preserves_attention_mask(self, module):
+        """attention_mask passes through unchanged."""
+        mask = torch.tensor([[1, 1, 0]], dtype=torch.long)
+        features = _make_features(
+            token_embeddings=torch.randn(1, 3, 16),
+            attention_mask=mask,
+            modality_ids=torch.ones(1, 3, dtype=torch.long),
+        )
+        result = module(features)
+        torch.testing.assert_close(result["attention_mask"], mask)
+
+    def test_preserves_modality_ids(self, module):
+        """modality_ids pass through unchanged."""
+        mod_ids = torch.tensor([[0, 1, 1, 2]], dtype=torch.long)
+        features = _make_features(
+            token_embeddings=torch.randn(1, 4, 16),
+            attention_mask=torch.tensor([[1, 1, 1, 0]], dtype=torch.long),
+            modality_ids=mod_ids,
+        )
+        result = module(features)
+        torch.testing.assert_close(result["modality_ids"], mod_ids)
+
+
+# ---------------------------------------------------------------------------
+# Gradient flow
+# ---------------------------------------------------------------------------
+class TestGradientFlow:
+    """Gradients must flow through the attention layers."""
+
+    def test_gradient_flow(self, module):
+        """Gradients flow through attention to omics tokens."""
+        x = torch.randn(2, 4, 16, requires_grad=True)
+        features = _make_features(
+            token_embeddings=x,
+            attention_mask=torch.ones(2, 4, dtype=torch.long),
+            modality_ids=torch.ones(2, 4, dtype=torch.long),
+        )
+        result = module(features)
+        loss = (result["token_embeddings"] ** 2).sum()
+        loss.backward()
+
+        # Input should have gradients
+        assert x.grad is not None
+        assert x.grad.abs().sum() > 0
+
+        # Module parameters should have gradients
+        has_grad = False
+        for p in module.parameters():
+            if p.grad is not None and p.grad.abs().sum() > 0:
+                has_grad = True
+                break
+        assert has_grad, "At least one module parameter should have non-zero gradient"
+
+    def test_gradient_only_through_omics(self, module):
+        """In a mixed batch, text token gradients are zero (passthrough)."""
+        x = torch.randn(1, 4, 16, requires_grad=True)
+        modality_ids = torch.tensor([[0, 0, 1, 1]], dtype=torch.long)
+        features = _make_features(
+            token_embeddings=x,
+            attention_mask=torch.ones(1, 4, dtype=torch.long),
+            modality_ids=modality_ids,
+        )
+        result = module(features)
+        # Only compute loss on omics tokens
+        loss = (result["token_embeddings"][:, 2:, :] ** 2).sum()
+        loss.backward()
+
+        # Omics positions should have gradients
+        assert x.grad[:, 2:, :].abs().sum() > 0
+
+    def test_weights_update(self, module):
+        """Optimizer step changes attention weights."""
+        optimizer = torch.optim.SGD(module.parameters(), lr=0.1)
+
+        params_before = {n: p.clone() for n, p in module.named_parameters()}
+
+        features = _make_features(
+            token_embeddings=torch.randn(2, 4, 16),
+            attention_mask=torch.ones(2, 4, dtype=torch.long),
+            modality_ids=torch.ones(2, 4, dtype=torch.long),
+        )
+        result = module(features)
+        loss = (result["token_embeddings"] ** 2).sum()
+        loss.backward()
+        optimizer.step()
+
+        total_delta = sum((p - params_before[n]).abs().sum().item() for n, p in module.named_parameters())
+        assert total_delta > 0, "Parameters did not change after optimizer step"
+
+
+# ---------------------------------------------------------------------------
+# Single token sequence (obs-like)
+# ---------------------------------------------------------------------------
+class TestSingleTokenSequence:
+    """Obs-like input with L=1 must work without error."""
+
+    def test_single_token_sequence(self, module):
+        """Single omics token (obs-level) passes through without error."""
+        module.eval()
+        B, L, D = 3, 1, 16
+        x = torch.randn(B, L, D)
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(B, L, dtype=torch.long),
+            modality_ids=torch.ones(B, L, dtype=torch.long),
+        )
+        result = module(features)
+
+        # Shape preserved
+        assert result["token_embeddings"].shape == (B, L, D)
+
+    def test_single_token_no_nan(self, module):
+        """Single token doesn't produce NaN (self-attention on length-1 is trivial)."""
+        module.eval()
+        x = torch.randn(1, 1, 16)
+        features = _make_features(
+            token_embeddings=x.clone(),
+            attention_mask=torch.ones(1, 1, dtype=torch.long),
+            modality_ids=torch.ones(1, 1, dtype=torch.long),
+        )
+        result = module(features)
+        assert not torch.isnan(result["token_embeddings"]).any()
+
+
+# ---------------------------------------------------------------------------
+# Persistence
+# ---------------------------------------------------------------------------
+class TestPersistence:
+    """Save/load roundtrip tests."""
+
+    def test_save_creates_files(self, module, tmp_dir, real_safetensors):
+        """save() creates config and weight files on disk."""
+        module.save(tmp_dir)
+
+        config_path = os.path.join(tmp_dir, module.config_file_name)
+        assert os.path.isfile(config_path)
+        assert os.path.isfile(os.path.join(tmp_dir, "model.safetensors"))
+
+    def test_save_load_roundtrip(self, module, tmp_dir, real_safetensors):
+        """Config and weights survive save/load cycle."""
+        module.save(tmp_dir)
+        loaded = OmicsAttentionModule.load(tmp_dir)
+
+        assert loaded.get_sentence_embedding_dimension() == module.get_sentence_embedding_dimension()
+        assert loaded.input_dim == module.input_dim
+        assert loaded.num_heads == module.num_heads
+        assert loaded.num_layers == module.num_layers
+
+    def test_save_load_produces_same_output(self, module, tmp_dir, real_safetensors):
+        """Loaded module produces identical output to original."""
+        module.eval()
+        x = torch.randn(2, 4, 16)
+        features = _make_features(
+            token_embeddings=x,
+            attention_mask=torch.ones(2, 4, dtype=torch.long),
+            modality_ids=torch.ones(2, 4, dtype=torch.long),
+        )
+        original_out = module(features)["token_embeddings"].detach()
+
+        module.save(tmp_dir)
+        loaded = OmicsAttentionModule.load(tmp_dir)
+        loaded.eval()
+
+        features2 = _make_features(
+            token_embeddings=x,
+            attention_mask=torch.ones(2, 4, dtype=torch.long),
+            modality_ids=torch.ones(2, 4, dtype=torch.long),
+        )
+        loaded_out = loaded(features2)["token_embeddings"].detach()
+
+        torch.testing.assert_close(original_out, loaded_out)
+
+
+# ---------------------------------------------------------------------------
+# Properties
+# ---------------------------------------------------------------------------
+class TestProperties:
+    """Module metadata and properties."""
+
+    def test_get_sentence_embedding_dimension(self, module):
+        """Returns input_dim (self-attention doesn't change dimensionality)."""
+        assert module.get_sentence_embedding_dimension() == 16
+
+    def test_repr(self, module):
+        """repr contains key config info."""
+        r = repr(module)
+        assert "16" in r  # input_dim
+        assert "2" in r  # num_heads
diff --git a/tests/test_prepare_store.py b/tests/test_prepare_store.py
new file mode 100644
index 0000000..e3ff284
--- /dev/null
+++ b/tests/test_prepare_store.py
@@ -0,0 +1,279 @@
+"""Tests for mmcontext.io.prepare_store — VectorStore preparation from zarr.
+
+These tests use synthetic zarr stores on disk (no network calls).
+The Zenodo URL fixup is tested via the internal helper.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import zarr
+
+from mmcontext.io.prepare_store import (
+    _get_obsm_zarr_array,
+    _open_zarr,
+    _read_obs_names_zarr,
+    _url_to_cache_name,
+    prepare_vector_store,
+)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def tmp_dir():
+    with tempfile.TemporaryDirectory() as d:
+        yield d
+
+
+@pytest.fixture
+def synthetic_zarr(tmp_dir):
+    """Create a minimal anndata-like zarr store on disk.
+
+    Layout mirrors what anndata writes:
+        obs/_index  — string array of obs names
+        obsm/X_pca  — (N, D) float32 array
+        obsm/X_scvi — (N, D2) float32 array
+    """
+    zarr_path = os.path.join(tmp_dir, "test.zarr")
+    root = zarr.open(zarr_path, mode="w")
+
+    n_obs, d_pca, d_scvi = 10, 4, 8
+    rng = np.random.default_rng(42)
+
+    def _write(group, name, data):
+        arr = group.create_array(name, shape=data.shape, dtype=data.dtype)
+        arr[:] = data
+        return arr
+
+    # obs group with _index
+    obs = root.create_group("obs")
+    obs_names = [f"cell_{i}" for i in range(n_obs)]
+    _write(obs, "_index", np.array(obs_names, dtype="U"))
+    obs.attrs["_index"] = "_index"
+
+    # obsm group
+    obsm = root.create_group("obsm")
+    _write(obsm, "X_pca", rng.standard_normal((n_obs, d_pca)).astype(np.float32))
+    _write(obsm, "X_scvi", rng.standard_normal((n_obs, d_scvi)).astype(np.float32))
+
+    return zarr_path, obs_names, n_obs, d_pca, d_scvi
+
+
+# ---------------------------------------------------------------------------
+# URL fixup tests
+# ---------------------------------------------------------------------------
+class TestZenodoUrlFixup:
+    """Zenodo draft→published URL conversion."""
+
+    def test_draft_url_is_converted(self):
+        """_download_zarr should strip /draft/ from Zenodo URLs.
+
+        We can't test the actual download without network, but we verify
+        the URL transformation logic by checking the function's behavior
+        with a non-existent cache.
+        """
+        draft_url = "https://zenodo.org/api/records/12345/draft/files/data.zarr.zip/content"
+        expected_published = "https://zenodo.org/api/records/12345/files/data.zarr.zip/content"
+
+        # The fix is inside _download_zarr — we test it indirectly by
+        # verifying the URL would be transformed. Since we can't mock
+        # requests easily here, we just test the string operation.
+        assert "/draft/" in draft_url
+        fixed = draft_url.replace("/draft/", "/")
+        assert fixed == expected_published
+
+    def test_non_zenodo_url_unchanged(self):
+        url = "https://example.com/data/draft/file.zip"
+        # Only zenodo.org URLs get the fixup
+        assert "zenodo.org" not in url
+
+    def test_cache_name_deterministic(self):
+        url = "https://zenodo.org/api/records/12345/files/data.zarr.zip/content"
+        assert _url_to_cache_name(url) == _url_to_cache_name(url)
+
+    def test_cache_name_differs_for_different_urls(self):
+        url_a = "https://zenodo.org/api/records/111/files/a.zip/content"
+        url_b = "https://zenodo.org/api/records/222/files/b.zip/content"
+        assert _url_to_cache_name(url_a) != _url_to_cache_name(url_b)
+
+
+# ---------------------------------------------------------------------------
+# Zarr reading tests
+# ---------------------------------------------------------------------------
+class TestReadObsNames:
+    def test_reads_obs_names(self, synthetic_zarr):
+        zarr_path, expected_names, *_ = synthetic_zarr
+        root = zarr.open(zarr_path, mode="r")
+        names = _read_obs_names_zarr(root)
+        assert names == expected_names
+
+    def test_returns_strings(self, synthetic_zarr):
+        zarr_path, *_ = synthetic_zarr
+        root = zarr.open(zarr_path, mode="r")
+        names = _read_obs_names_zarr(root)
+        assert all(isinstance(n, str) for n in names)
+
+
+class TestReadObsm:
+    def test_reads_obsm_array(self, synthetic_zarr):
+        zarr_path, _, n_obs, d_pca, _ = synthetic_zarr
+        root = zarr.open(zarr_path, mode="r")
+        arr = _get_obsm_zarr_array(root, "X_pca")
+        assert arr.shape == (n_obs, d_pca)
+
+    def test_missing_key_raises(self, synthetic_zarr):
+        zarr_path, *_ = synthetic_zarr
+        root = zarr.open(zarr_path, mode="r")
+        with pytest.raises(KeyError, match="X_nonexistent"):
+            _get_obsm_zarr_array(root, "X_nonexistent")
+
+    def test_no_obsm_group_raises(self, tmp_dir):
+        zarr_path = os.path.join(tmp_dir, "empty.zarr")
+        root = zarr.open(zarr_path, mode="w")
+        root.create_group("obs")
+        root = zarr.open(zarr_path, mode="r")
+        with pytest.raises(KeyError, match="obsm"):
+            _get_obsm_zarr_array(root, "X_pca")
+
+
+class TestOpenZarr:
+    def test_open_directory(self, synthetic_zarr):
+        from pathlib import Path
+
+        zarr_path, *_ = synthetic_zarr
+        root = _open_zarr(Path(zarr_path))
+        assert "obs" in root
+        assert "obsm" in root
+
+    def test_nonexistent_raises(self, tmp_dir):
+        from pathlib import Path
+
+        with pytest.raises(FileNotFoundError):
+            _open_zarr(Path(tmp_dir) / "nope.zarr")
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: prepare_vector_store with local zarr paths
+# ---------------------------------------------------------------------------
+class TestPrepareVectorStore:
+    """Integration test using local zarr paths (no network)."""
+
+    def test_builds_store_from_local_zarr(self, synthetic_zarr, tmp_dir):
+        """prepare_vector_store with local adata_link paths."""
+        from datasets import Dataset
+
+        zarr_path, obs_names, n_obs, d_pca, _ = synthetic_zarr
+
+        # Simulate a dataset where all samples come from one zarr file
+        ds = Dataset.from_dict(
+            {
+                "sample_idx": obs_names[:5],  # use first 5
+                "adata_link": [zarr_path] * 5,
+            }
+        )
+
+        output_path = os.path.join(tmp_dir, "test_store.mmap")
+        store = prepare_vector_store(
+            ds,
+            obsm_key="X_pca",
+            output_path=output_path,
+        )
+
+        assert len(store) == 5
+        assert store.dim == d_pca
+        # All sample IDs should be present
+        for name in obs_names[:5]:
+            assert name in store
+
+    def test_values_match_zarr_source(self, synthetic_zarr, tmp_dir):
+        """Vectors in VectorStore match the zarr source."""
+        from datasets import Dataset
+
+        zarr_path, obs_names, *_ = synthetic_zarr
+
+        ds = Dataset.from_dict(
+            {
+                "sample_idx": [obs_names[3]],
+                "adata_link": [zarr_path],
+            }
+        )
+
+        output_path = os.path.join(tmp_dir, "val_store.mmap")
+        store = prepare_vector_store(ds, obsm_key="X_pca", output_path=output_path)
+
+        # Compare against direct zarr read
+        root = zarr.open(zarr_path, mode="r")
+        expected = np.asarray(root["obsm"]["X_pca"][3], dtype=np.float32)
+        np.testing.assert_array_almost_equal(store[obs_names[3]], expected)
+
+    def test_missing_sample_id_raises(self, synthetic_zarr, tmp_dir):
+        """Unknown sample_idx raises KeyError."""
+        from datasets import Dataset
+
+        zarr_path, *_ = synthetic_zarr
+
+        ds = Dataset.from_dict(
+            {
+                "sample_idx": ["nonexistent_cell"],
+                "adata_link": [zarr_path],
+            }
+        )
+
+        output_path = os.path.join(tmp_dir, "err_store.mmap")
+        with pytest.raises(KeyError, match="nonexistent_cell"):
+            prepare_vector_store(ds, obsm_key="X_pca", output_path=output_path)
+
+    def test_skips_existing_store(self, synthetic_zarr, tmp_dir):
+        """If output_path exists and overwrite=False, loads existing store."""
+        from datasets import Dataset
+
+        zarr_path, obs_names, *_ = synthetic_zarr
+
+        ds = Dataset.from_dict(
+            {
+                "sample_idx": obs_names[:3],
+                "adata_link": [zarr_path] * 3,
+            }
+        )
+
+        output_path = os.path.join(tmp_dir, "cached_store.mmap")
+
+        # Build once
+        store1 = prepare_vector_store(ds, obsm_key="X_pca", output_path=output_path)
+        # Build again — should load from disk
+        store2 = prepare_vector_store(ds, obsm_key="X_pca", output_path=output_path)
+
+        assert len(store2) == len(store1)
+
+    def test_different_obsm_keys(self, synthetic_zarr, tmp_dir):
+        """Can build stores from different obsm keys."""
+        from datasets import Dataset
+
+        zarr_path, obs_names, _, d_pca, d_scvi = synthetic_zarr
+
+        ds = Dataset.from_dict(
+            {
+                "sample_idx": obs_names[:2],
+                "adata_link": [zarr_path] * 2,
+            }
+        )
+
+        pca_store = prepare_vector_store(
+            ds,
+            obsm_key="X_pca",
+            output_path=os.path.join(tmp_dir, "pca.mmap"),
+        )
+        scvi_store = prepare_vector_store(
+            ds,
+            obsm_key="X_scvi",
+            output_path=os.path.join(tmp_dir, "scvi.mmap"),
+        )
+
+        assert pca_store.dim == d_pca
+        assert scvi_store.dim == d_scvi
diff --git a/tests/test_query_annotate.py b/tests/test_query_annotate.py
index 3d456d5..a48c136 100644
--- a/tests/test_query_annotate.py
+++ b/tests/test_query_annotate.py
@@ -240,12 +240,12 @@ def encode(self, list_of_strings):
 
     # Verify best_label matches highest score in inferred_labels
     for i in range(2):
-        label_scores = inferred_labels[i]
-        found_best = best_labels[i]
+        label_scores = inferred_labels.iloc[i]
+        found_best = best_labels.iloc[i]
         assert found_best in label_scores
         max_label = max(label_scores, key=label_scores.get)
         assert found_best == max_label
 
     # Verify expected behavior based on our encoding scheme
-    assert best_labels[0] == "label_0"
-    assert best_labels[1] == "label_1"
+    assert best_labels.iloc[0] == "label_0"
+    assert best_labels.iloc[1] == "label_1"
diff --git a/tests/test_register_initial_embeddings.py b/tests/test_register_initial_embeddings.py
index 5ce21af..42dd272 100644
--- a/tests/test_register_initial_embeddings.py
+++ b/tests/test_register_initial_embeddings.py
@@ -5,7 +5,7 @@
 import pandas as pd
 import pytest
 
-from mmcontext.mmcontextencoder import MMContextEncoder
+from mmcontext._legacy.mmcontextencoder import MMContextEncoder
 
 
 # --- helper that builds a text-only encoder cheaply (no real BERT) --------
diff --git a/tests/test_sentence_transformer_integration.py b/tests/test_sentence_transformer_integration.py
index 1f1b697..e1718e0 100644
--- a/tests/test_sentence_transformer_integration.py
+++ b/tests/test_sentence_transformer_integration.py
@@ -17,7 +17,7 @@
 )
 from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
 
-from mmcontext.mmcontextencoder import MMContextEncoder, MMContextProcessor
+from mmcontext._legacy.mmcontextencoder import MMContextEncoder
 
 logger = logging.getLogger(__name__)
 
diff --git a/tests/test_st_integration.py b/tests/test_st_integration.py
new file mode 100644
index 0000000..3216be8
--- /dev/null
+++ b/tests/test_st_integration.py
@@ -0,0 +1,604 @@
+"""Integration tests for the full SentenceTransformer pipeline.
+
+These tests verify that the mmcontext modules compose correctly into a
+sentence-transformers pipeline and that encode(), save/load, and training
+work end-to-end.
+
+Two pipeline configurations are tested:
+
+1. **Obs pipeline** (samples have a single omics vector each):
+   ``[MMContextModule, AdapterModule, Pooling, Normalize]``
+
+2. **Var pipeline** (samples have multiple gene vectors each):
+   ``[MMContextModule, OmicsAttentionModule, AdapterModule, Pooling, Normalize]``
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.sentence_transformer.modules import Normalize, Pooling
+
+from mmcontext.io import VectorStore
+from mmcontext.modules import AdapterModule, MMContextModule, OmicsAttentionModule
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def tmp_dir():
+    with tempfile.TemporaryDirectory() as d:
+        yield d
+
+
+@pytest.fixture
+def real_safetensors():
+    """Undo the global safetensors.torch.load_model patch for persistence tests."""
+    import safetensors.torch
+
+    importlib.reload(safetensors.torch)
+    yield
+
+
+# -- Shared dims used across all fixtures -----------------------------------
+TEXT_DIM = 32  # matches _TextEncStub hidden_size
+OMICS_DIM = 8
+SHARED_DIM = 16
+
+
+@pytest.fixture
+def obs_store(tmp_dir):
+    """VectorStore with 5 obs-level samples, dim=OMICS_DIM."""
+    ids = [f"cell_{i}" for i in range(5)]
+    rng = np.random.default_rng(42)
+    data = rng.standard_normal((5, OMICS_DIM)).astype(np.float32)
+    store = VectorStore.from_numpy(data, ids, path=os.path.join(tmp_dir, "obs_store.mmap"))
+    return store
+
+
+@pytest.fixture
+def mmcontext_module():
+    """MMContextModule backed by the conftest stubs (patched AutoModel/Tokenizer)."""
+    return MMContextModule(
+        model_name_or_path="bert-base-uncased",  # intercepted by conftest patch
+        max_seq_length=8,
+    )
+
+
+@pytest.fixture
+def adapter_module():
+    """AdapterModule: TEXT_DIM → SHARED_DIM, OMICS_DIM → SHARED_DIM."""
+    return AdapterModule(
+        text_input_dim=TEXT_DIM,
+        omics_input_dim=OMICS_DIM,
+        shared_dim=SHARED_DIM,
+        hidden_dim=32,
+    )
+
+
+@pytest.fixture
+def attention_module():
+    """OmicsAttentionModule matching OMICS_DIM."""
+    return OmicsAttentionModule(
+        input_dim=OMICS_DIM,
+        num_heads=2,
+        num_layers=1,
+        dropout=0.0,
+    )
+
+
+@pytest.fixture
+def obs_pipeline(mmcontext_module, adapter_module):
+    """Obs pipeline: MMContext → Adapter → Pooling → Normalize."""
+    return SentenceTransformer(
+        modules=[
+            mmcontext_module,
+            adapter_module,
+            Pooling(embedding_dimension=SHARED_DIM, pooling_mode="mean"),
+            Normalize(),
+        ]
+    )
+
+
+@pytest.fixture
+def var_pipeline(mmcontext_module, attention_module, adapter_module):
+    """Var pipeline: MMContext → OmicsAttention → Adapter → Pooling → Normalize."""
+    return SentenceTransformer(
+        modules=[
+            mmcontext_module,
+            attention_module,
+            adapter_module,
+            Pooling(embedding_dimension=SHARED_DIM, pooling_mode="mean"),
+            Normalize(),
+        ]
+    )
+
+
+# ---------------------------------------------------------------------------
+# Pipeline construction
+# ---------------------------------------------------------------------------
+class TestPipelineConstruction:
+    """Modules compose into a valid SentenceTransformer."""
+
+    def test_obs_pipeline_construction(self, obs_pipeline):
+        """Obs pipeline constructs without error."""
+        assert isinstance(obs_pipeline, SentenceTransformer)
+        # Should have 4 modules: MMContext, Adapter, Pooling, Normalize
+        modules = list(obs_pipeline.children())
+        assert len(modules) == 4
+
+    def test_var_pipeline_construction(self, var_pipeline):
+        """Var pipeline constructs without error."""
+        assert isinstance(var_pipeline, SentenceTransformer)
+        # Should have 5 modules: MMContext, OmicsAttention, Adapter, Pooling, Normalize
+        modules = list(var_pipeline.children())
+        assert len(modules) == 5
+
+
+# ---------------------------------------------------------------------------
+# Encode — text
+# ---------------------------------------------------------------------------
+class TestEncodeText:
+    """model.encode() with text inputs."""
+
+    def test_encode_text_shape(self, obs_pipeline):
+        """Encoding text produces embeddings of shape (N, SHARED_DIM)."""
+        embeddings = obs_pipeline.encode(["Hello world", "Another sentence"])
+        assert isinstance(embeddings, np.ndarray)
+        assert embeddings.shape == (2, SHARED_DIM)
+
+    def test_encode_text_single(self, obs_pipeline):
+        """Single text input returns (SHARED_DIM,) array."""
+        embedding = obs_pipeline.encode("Hello world")
+        assert isinstance(embedding, np.ndarray)
+        assert embedding.shape == (SHARED_DIM,)
+
+
+# ---------------------------------------------------------------------------
+# Encode — omics direct
+# ---------------------------------------------------------------------------
+class TestEncodeOmicsDirect:
+    """model.encode() with direct omics vectors."""
+
+    def test_encode_omics_direct_obs(self, obs_pipeline):
+        """Encoding direct omics vector produces correct shape."""
+        vec = np.random.randn(OMICS_DIM).astype(np.float32)
+        embedding = obs_pipeline.encode([{"omics_values": vec}])
+        assert isinstance(embedding, np.ndarray)
+        assert embedding.shape == (1, SHARED_DIM)
+
+    def test_encode_omics_direct_var(self, var_pipeline):
+        """Encoding var-level (multiple gene vectors) produces correct shape."""
+        genes = [np.random.randn(OMICS_DIM).astype(np.float32) for _ in range(5)]
+        embedding = var_pipeline.encode([{"omics_values": genes}])
+        assert isinstance(embedding, np.ndarray)
+        assert embedding.shape == (1, SHARED_DIM)
+
+
+# ---------------------------------------------------------------------------
+# Encode — omics via VectorStore
+# ---------------------------------------------------------------------------
+class TestEncodeOmicsViaStore:
+    """model.encode() with prefixed IDs resolved through VectorStore."""
+
+    def test_encode_omics_via_store(self, obs_pipeline, obs_store):
+        """Prefixed IDs resolved through VectorStore produce correct shape."""
+        # Attach store to the MMContextModule (first module)
+        first_module = list(obs_pipeline.children())[0]
+        first_module.set_vector_store(obs_store)
+
+        embeddings = obs_pipeline.encode(["omics:cell_0", "omics:cell_1"])
+        assert isinstance(embeddings, np.ndarray)
+        assert embeddings.shape == (2, SHARED_DIM)
+
+        first_module.remove_vector_store()
+
+
+# ---------------------------------------------------------------------------
+# Encode — normalization
+# ---------------------------------------------------------------------------
+class TestEncodeNormalize:
+    """Output vectors should be L2-normalized (Normalize module is last)."""
+
+    def test_encode_normalize(self, obs_pipeline):
+        """Encoded vectors have unit L2 norm."""
+        embeddings = obs_pipeline.encode(["A sample text", "Another text"])
+        norms = np.linalg.norm(embeddings, axis=1)
+        np.testing.assert_allclose(norms, 1.0, atol=1e-5)
+
+    def test_encode_omics_normalize(self, obs_pipeline):
+        """Omics embeddings also have unit L2 norm."""
+        vec = np.random.randn(OMICS_DIM).astype(np.float32)
+        embedding = obs_pipeline.encode([{"omics_values": vec}])
+        norm = np.linalg.norm(embedding)
+        np.testing.assert_allclose(norm, 1.0, atol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# max_seq_length propagation
+# ---------------------------------------------------------------------------
+class TestMaxSeqLength:
+    """max_seq_length accessible from SentenceTransformer level."""
+
+    def test_max_seq_length_propagation(self, obs_pipeline):
+        """max_seq_length from MMContextModule is accessible on the ST model."""
+        assert obs_pipeline.max_seq_length == 8
+
+
+# ---------------------------------------------------------------------------
+# Obs pipeline end-to-end
+# ---------------------------------------------------------------------------
+class TestObsPipeline:
+    """Full obs pipeline: [MMContext, Adapter, Pooling, Normalize]."""
+
+    def test_obs_text_and_omics_same_output_dim(self, obs_pipeline):
+        """Text and omics inputs produce same dimensionality."""
+        text_emb = obs_pipeline.encode(["Some text"])
+        omics_emb = obs_pipeline.encode([{"omics_values": np.random.randn(OMICS_DIM).astype(np.float32)}])
+        assert text_emb.shape[-1] == omics_emb.shape[-1] == SHARED_DIM
+
+    def test_obs_deterministic(self, obs_pipeline):
+        """Same input produces same output (model in eval mode)."""
+        emb1 = obs_pipeline.encode(["Determinism test"])
+        emb2 = obs_pipeline.encode(["Determinism test"])
+        np.testing.assert_allclose(emb1, emb2)
+
+
+# ---------------------------------------------------------------------------
+# Var pipeline end-to-end
+# ---------------------------------------------------------------------------
+class TestVarPipeline:
+    """Full var pipeline: [MMContext, OmicsAttn, Adapter, Pooling, Normalize]."""
+
+    def test_var_single_gene(self, var_pipeline):
+        """Single gene vector (like obs) works in var pipeline too."""
+        genes = [np.random.randn(OMICS_DIM).astype(np.float32)]
+        embedding = var_pipeline.encode([{"omics_values": genes}])
+        assert embedding.shape == (1, SHARED_DIM)
+
+    def test_var_multiple_genes(self, var_pipeline):
+        """Multiple gene vectors are attended and pooled."""
+        genes = [np.random.randn(OMICS_DIM).astype(np.float32) for _ in range(10)]
+        embedding = var_pipeline.encode([{"omics_values": genes}])
+        assert embedding.shape == (1, SHARED_DIM)
+
+    def test_var_text_passthrough(self, var_pipeline):
+        """Text inputs work through var pipeline (attention is no-op on text)."""
+        embedding = var_pipeline.encode("Some text through var pipeline")
+        assert embedding.shape == (SHARED_DIM,)
+
+
+# ---------------------------------------------------------------------------
+# Save / Load full pipeline
+# ---------------------------------------------------------------------------
+class TestSaveLoad:
+    """Save and load full pipeline roundtrip."""
+
+    def test_save_load_full_pipeline(self, obs_pipeline, tmp_dir, real_safetensors):
+        """Saved pipeline can be loaded and produces same output."""
+        obs_pipeline.eval()
+
+        # Encode before save
+        text_input = ["Test save load roundtrip"]
+        original_emb = obs_pipeline.encode(text_input)
+
+        # Save
+        save_path = os.path.join(tmp_dir, "test_model")
+        obs_pipeline.save(save_path)
+
+        # Verify modules.json exists
+        modules_json_path = os.path.join(save_path, "modules.json")
+        assert os.path.isfile(modules_json_path)
+
+        # Load and re-encode
+        loaded = SentenceTransformer(save_path, trust_remote_code=True)
+        loaded.eval()
+        loaded_emb = loaded.encode(text_input)
+
+        np.testing.assert_allclose(original_emb, loaded_emb, atol=1e-5)
+
+    def test_modules_json_structure(self, obs_pipeline, tmp_dir, real_safetensors):
+        """modules.json contains correct module chain."""
+        save_path = os.path.join(tmp_dir, "test_model")
+        obs_pipeline.save(save_path)
+
+        with open(os.path.join(save_path, "modules.json")) as f:
+            modules_config = json.load(f)
+
+        assert len(modules_config) == 4
+
+        # Check types contain our custom module classes
+        types = [m["type"] for m in modules_config]
+        assert any("MMContextModule" in t for t in types)
+        assert any("AdapterModule" in t for t in types)
+        assert any("Pooling" in t for t in types)
+        assert any("Normalize" in t for t in types)
+
+    def test_save_load_var_pipeline(self, var_pipeline, tmp_dir, real_safetensors):
+        """Var pipeline save/load roundtrip."""
+        var_pipeline.eval()
+
+        save_path = os.path.join(tmp_dir, "test_var_model")
+        var_pipeline.save(save_path)
+
+        with open(os.path.join(save_path, "modules.json")) as f:
+            modules_config = json.load(f)
+
+        assert len(modules_config) == 5
+        types = [m["type"] for m in modules_config]
+        assert any("OmicsAttentionModule" in t for t in types)
+
+
+# ---------------------------------------------------------------------------
+# Precision conversion
+# ---------------------------------------------------------------------------
+class TestPrecisionConversion:
+    """fp32 → fp16 works across all modules."""
+
+    def test_precision_conversion_parameters(self, mmcontext_module, adapter_module):
+        """Pipeline modules can be converted to fp16."""
+        pipeline = SentenceTransformer(
+            modules=[
+                mmcontext_module,
+                adapter_module,
+                Pooling(embedding_dimension=SHARED_DIM, pooling_mode="mean"),
+                Normalize(),
+            ]
+        )
+        pipeline.half()
+
+        # All learnable parameters should now be fp16
+        for name, param in pipeline.named_parameters():
+            assert param.dtype == torch.float16, f"Parameter {name} is {param.dtype}, expected float16"
+
+        # Restore to fp32 — the session-scoped _TextEncStub (from conftest)
+        # is shared across all tests; leaving it in fp16 would pollute
+        # subsequent tests that expect fp32 weights.
+        pipeline.float()
+
+    def test_precision_conversion_omics_encode(self):
+        """fp16 pipeline produces valid omics embeddings.
+
+        Skipped: the stub text encoder always returns fp32 tensors
+        regardless of module dtype, causing a dtype mismatch when the
+        fp32 input hits fp16 Linear weights. Real encoders handle
+        autocasting properly, but stubs don't.
+        """
+        pytest.skip("Stub encoder produces fp32 regardless of model dtype — not testable without real encoder")
+
+
+# ---------------------------------------------------------------------------
+# Training with SentenceTransformerTrainer + MNR loss
+# ---------------------------------------------------------------------------
+class TestTraining:
+    """Training integration with SentenceTransformerTrainer.
+
+    Uses MultipleNegativesRankingLoss with (anchor, positive) pairs.
+    Three modes mirror real usage:
+
+    1. Text-only — both columns are plain text.
+    2. Bimodal — anchors are omics (via VectorStore), positives are text.
+    3. Gene-list — anchors are gene-name strings (treated as text), positives are text.
+    """
+
+    def test_training_text_only(self, obs_pipeline, tmp_dir):
+        """One training step with text-only (anchor, positive) pairs."""
+        from datasets import Dataset
+        from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
+        from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
+
+        ds = Dataset.from_dict(
+            {
+                "anchor": [
+                    "A neuron from the thalamus.",
+                    "An epithelial cell from the lung.",
+                    "A B cell from peripheral blood.",
+                    "A fibroblast from skin tissue.",
+                ],
+                "positive": [
+                    "Thalamic neuron expressing SYT1 and GNAS.",
+                    "Lung epithelial cell with high EPCAM.",
+                    "CD19-positive B lymphocyte.",
+                    "Dermal fibroblast with COL1A1 expression.",
+                ],
+            }
+        )
+
+        loss = MultipleNegativesRankingLoss(obs_pipeline)
+        args = SentenceTransformerTrainingArguments(
+            output_dir=os.path.join(tmp_dir, "train_text"),
+            num_train_epochs=1,
+            per_device_train_batch_size=2,
+            learning_rate=1e-3,
+            use_cpu=True,
+            report_to="none",
+        )
+        trainer = SentenceTransformerTrainer(
+            model=obs_pipeline,
+            args=args,
+            train_dataset=ds,
+            loss=loss,
+        )
+
+        # Snapshot parameters before training
+        params_before = {n: p.clone() for n, p in obs_pipeline.named_parameters() if p.requires_grad}
+
+        trainer.train()
+
+        # At least some parameters should have changed
+        total_delta = sum(
+            (p - params_before[n]).abs().sum().item() for n, p in obs_pipeline.named_parameters() if n in params_before
+        )
+        assert total_delta > 0, "No parameters changed during text-only training"
+
+    def test_training_bimodal(self, obs_pipeline, obs_store, tmp_dir):
+        """One training step with omics anchors + text positives.
+
+        Mimics the real dataset structure: sample_idx column is prefixed
+        with the omics prefix, positive column is plain text.
+        """
+        from datasets import Dataset
+        from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
+        from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
+
+        # Attach VectorStore
+        first_module = list(obs_pipeline.children())[0]
+        first_module.set_vector_store(obs_store)
+
+        ds = Dataset.from_dict(
+            {
+                "anchor": [
+                    "omics:cell_0",
+                    "omics:cell_1",
+                    "omics:cell_2",
+                    "omics:cell_3",
+                ],
+                "positive": [
+                    "Thalamic neuron expressing SYT1.",
+                    "Lung epithelial cell with EPCAM.",
+                    "CD19-positive B lymphocyte.",
+                    "Dermal fibroblast with COL1A1.",
+                ],
+            }
+        )
+
+        loss = MultipleNegativesRankingLoss(obs_pipeline)
+        args = SentenceTransformerTrainingArguments(
+            output_dir=os.path.join(tmp_dir, "train_bimodal"),
+            num_train_epochs=1,
+            per_device_train_batch_size=2,
+            learning_rate=1e-3,
+            use_cpu=True,
+            report_to="none",
+        )
+        trainer = SentenceTransformerTrainer(
+            model=obs_pipeline,
+            args=args,
+            train_dataset=ds,
+            loss=loss,
+        )
+
+        params_before = {n: p.clone() for n, p in obs_pipeline.named_parameters() if p.requires_grad}
+
+        trainer.train()
+
+        total_delta = sum(
+            (p - params_before[n]).abs().sum().item() for n, p in obs_pipeline.named_parameters() if n in params_before
+        )
+        assert total_delta > 0, "No parameters changed during bimodal training"
+
+        first_module.remove_vector_store()
+
+    def test_training_gene_list(self, obs_pipeline, tmp_dir):
+        """One training step with gene-name strings as anchors.
+
+        Gene-name lists (like cell_sentence_1 in the real dataset) are
+        plain text — they go through the text tokenization path.
+        """
+        from datasets import Dataset
+        from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
+        from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
+
+        ds = Dataset.from_dict(
+            {
+                "anchor": [
+                    "MALAT1 MT-CO3 GNAS SYT1 CALM1",
+                    "EPCAM KRT8 KRT18 MUC1",
+                    "CD19 MS4A1 CD79A PAX5",
+                    "COL1A1 COL3A1 FN1 VIM",
+                ],
+                "positive": [
+                    "Thalamic neuron expressing SYT1.",
+                    "Lung epithelial cell with EPCAM.",
+                    "CD19-positive B lymphocyte.",
+                    "Dermal fibroblast with COL1A1.",
+                ],
+            }
+        )
+
+        loss = MultipleNegativesRankingLoss(obs_pipeline)
+        args = SentenceTransformerTrainingArguments(
+            output_dir=os.path.join(tmp_dir, "train_genelist"),
+            num_train_epochs=1,
+            per_device_train_batch_size=2,
+            learning_rate=1e-3,
+            use_cpu=True,
+            report_to="none",
+        )
+        trainer = SentenceTransformerTrainer(
+            model=obs_pipeline,
+            args=args,
+            train_dataset=ds,
+            loss=loss,
+        )
+
+        params_before = {n: p.clone() for n, p in obs_pipeline.named_parameters() if p.requires_grad}
+
+        trainer.train()
+
+        total_delta = sum(
+            (p - params_before[n]).abs().sum().item() for n, p in obs_pipeline.named_parameters() if n in params_before
+        )
+        assert total_delta > 0, "No parameters changed during gene-list training"
+
+    def test_training_save_load_roundtrip(self, obs_pipeline, tmp_dir, real_safetensors):
+        """Trained model can be saved, reloaded, and produces identical encodings."""
+        from datasets import Dataset
+        from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
+        from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
+
+        ds = Dataset.from_dict(
+            {
+                "anchor": [
+                    "MALAT1 MT-CO3 GNAS SYT1",
+                    "EPCAM KRT8 KRT18 MUC1",
+                    "CD19 MS4A1 CD79A PAX5",
+                    "COL1A1 COL3A1 FN1 VIM",
+                ],
+                "positive": [
+                    "Thalamic neuron expressing SYT1.",
+                    "Lung epithelial cell with EPCAM.",
+                    "CD19-positive B lymphocyte.",
+                    "Dermal fibroblast with COL1A1.",
+                ],
+            }
+        )
+
+        loss = MultipleNegativesRankingLoss(obs_pipeline)
+        save_path = os.path.join(tmp_dir, "trained_model")
+        args = SentenceTransformerTrainingArguments(
+            output_dir=save_path,
+            num_train_epochs=1,
+            per_device_train_batch_size=2,
+            learning_rate=1e-3,
+            use_cpu=True,
+            report_to="none",
+        )
+        trainer = SentenceTransformerTrainer(
+            model=obs_pipeline,
+            args=args,
+            train_dataset=ds,
+            loss=loss,
+        )
+        trainer.train()
+
+        # Encode with trained model
+        obs_pipeline.eval()
+        test_inputs = ["Test sentence after training."]
+        trained_emb = obs_pipeline.encode(test_inputs)
+
+        # Save and reload
+        obs_pipeline.save(save_path)
+        loaded = SentenceTransformer(save_path)
+        loaded.eval()
+        loaded_emb = loaded.encode(test_inputs)
+
+        np.testing.assert_allclose(trained_emb, loaded_emb, atol=1e-5)
diff --git a/tests/test_training.py b/tests/test_training.py
new file mode 100644
index 0000000..f952ded
--- /dev/null
+++ b/tests/test_training.py
@@ -0,0 +1,374 @@
+"""Tests for config-driven multi-dataset training (src/mmcontext/training.py).
+
+These cover the parts of the new pipeline that the single-dataset tests in
+``test_st_integration.py`` do not exercise:
+
+1. ``build_namespaced_vector_store`` — merging per-dataset stores under
+   ``{name}:{id}`` keys without collisions.
+2. ``assemble_training_data`` + a one-step multi-dataset
+   ``SentenceTransformerTrainer`` run (dict of datasets, dict of losses,
+   ``SequentialEvaluator``) mixing a bimodal omics dataset and a bio dataset.
+3. End-to-end namespacing: ``omics:{name}:{id}`` anchors resolve in the merged
+   store and ``model.encode`` succeeds.
+
+All tests run offline: the conftest session fixture patches AutoModel /
+AutoTokenizer with stubs (text hidden_size=32), and VectorStores are built from
+synthetic numpy arrays.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import pytest
+from datasets import Dataset, DatasetDict
+from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
+from sentence_transformers.evaluation import SequentialEvaluator
+
+from mmcontext.embed import build_pipeline
+from mmcontext.io import VectorStore, build_namespaced_vector_store
+from mmcontext.training import (
+    DatasetConfig,
+    MiningConfig,
+    TrainConfig,
+    TrainerConfig,
+    assemble_training_data,
+    config_from_dict,
+    format_training_details,
+)
+
+OMICS_DIM = 8
+SHARED_DIM = 16
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def pipeline():
+    """Obs pipeline built via the shared builder (stubbed text encoder)."""
+    return build_pipeline(
+        "bert-base-uncased",
+        omics_dim=OMICS_DIM,
+        shared_dim=SHARED_DIM,
+        adapter_hidden_dim=SHARED_DIM,
+        max_seq_length=8,
+    )
+
+
+def _make_store(ids, tmp_path, name="store", seed=0):
+    rng = np.random.default_rng(seed)
+    data = rng.standard_normal((len(ids), OMICS_DIM)).astype(np.float32)
+    return VectorStore.from_numpy(data, ids, path=os.path.join(tmp_path, f"{name}.mmap")), data
+
+
+def _omics_split(sample_ids):
+    """A raw omics split with hard-negative idx columns referencing sample_ids."""
+    n = len(sample_ids)
+    return Dataset.from_dict(
+        {
+            "sample_idx": list(sample_ids),
+            "cell_sentence_1": [f"GENE{i} GENE{(i + 1) % n}" for i in range(n)],
+            "positive": [f"description of cell {i}" for i in range(n)],
+            # odd negative -> resolves to positive of the referenced idx;
+            # even negative -> resolves to cell_sentence_1 of the referenced idx.
+            "negative_1_idx": [sample_ids[(i + 1) % n] for i in range(n)],
+            "negative_2_idx": [sample_ids[(i + 2) % n] for i in range(n)],
+            "adata_link": ["http://example/chunk.zarr"] * n,
+        }
+    )
+
+
+# ---------------------------------------------------------------------------
+# 1. build_namespaced_vector_store
+# ---------------------------------------------------------------------------
+class TestNamespacedStore:
+    def test_merge_namespaces_and_preserves_vectors(self, tmp_path):
+        store_a, data_a = _make_store(["0", "1"], tmp_path, name="a", seed=1)
+        store_b, data_b = _make_store(["0", "1"], tmp_path, name="b", seed=2)  # same orig ids
+
+        merged = build_namespaced_vector_store(
+            {"ds_a": (store_a, ["0", "1"]), "ds_b": (store_b, ["0", "1"])},
+            output_path=os.path.join(tmp_path, "merged.mmap"),
+        )
+
+        # No collisions despite overlapping original ids.
+        assert len(merged) == 4
+        for key in ("ds_a:0", "ds_a:1", "ds_b:0", "ds_b:1"):
+            assert key in merged
+
+        np.testing.assert_allclose(merged["ds_a:0"], data_a[0])
+        np.testing.assert_allclose(merged["ds_b:1"], data_b[1])
+
+    def test_dedup_preserves_order(self, tmp_path):
+        store, data = _make_store(["0", "1", "2"], tmp_path, name="c", seed=3)
+        merged = build_namespaced_vector_store(
+            {"ds": (store, ["0", "0", "1", "2", "1"])},  # duplicates
+            output_path=os.path.join(tmp_path, "merged_dedup.mmap"),
+        )
+        assert len(merged) == 3
+        np.testing.assert_allclose(merged["ds:2"], data[2])
+
+    def test_dim_mismatch_raises(self, tmp_path):
+        store_a, _ = _make_store(["0"], tmp_path, name="d1", seed=4)
+        other = np.random.default_rng(5).standard_normal((1, OMICS_DIM + 3)).astype(np.float32)
+        store_b = VectorStore.from_numpy(other, ["0"], path=os.path.join(tmp_path, "d2.mmap"))
+        with pytest.raises(ValueError, match="dim"):
+            build_namespaced_vector_store(
+                {"a": (store_a, ["0"]), "b": (store_b, ["0"])},
+                output_path=os.path.join(tmp_path, "bad.mmap"),
+            )
+
+    def test_empty_raises(self, tmp_path):
+        with pytest.raises(ValueError, match="empty"):
+            build_namespaced_vector_store({}, output_path=os.path.join(tmp_path, "x.mmap"))
+
+
+# ---------------------------------------------------------------------------
+# 2. assemble_training_data + multi-dataset training
+# ---------------------------------------------------------------------------
+class TestAssembleAndTrain:
+    def _cfg(self):
+        return TrainConfig(
+            obsm_key="X_scvi_fm",
+            shared_dim=SHARED_DIM,
+            omics_datasets=[
+                DatasetConfig(id="dummy/a", name="omics_a", type="multiplets", modality="bimodal"),
+            ],
+            bio_datasets=[
+                DatasetConfig(id="dummy/bio", name="bio_x", type="multiplets"),
+            ],
+            trainer=TrainerConfig(per_device_eval_batch_size=2, logging_steps=1),
+        )
+
+    def _omics_raw(self):
+        return {
+            "omics_a": DatasetDict(
+                train=_omics_split(["0", "1", "2", "3"]),
+                val=_omics_split(["0", "1"]),
+            )
+        }
+
+    def _bio_raw(self):
+        return {
+            "bio_x": DatasetDict(
+                train=Dataset.from_dict(
+                    {
+                        "anchor": ["a neuron", "a B cell", "a fibroblast"],
+                        "positive": ["SYT1+ neuron", "CD19+ B cell", "COL1A1+ fibroblast"],
+                        "negative": ["unrelated 1", "unrelated 2", "unrelated 3"],  # singular name
+                        "sample_id": ["x", "y", "z"],  # stray column that must be dropped
+                    }
+                )
+            )
+        }
+
+    def test_assemble_shapes(self, pipeline, tmp_path):
+        merged = build_namespaced_vector_store(
+            {"omics_a": (_make_store(["0", "1", "2", "3"], tmp_path, name="oa")[0], ["0", "1", "2", "3"])},
+            output_path=os.path.join(tmp_path, "merged.mmap"),
+        )
+        pipeline[0].set_vector_store(merged)
+
+        assembled = assemble_training_data(
+            self._cfg(), pipeline, self._omics_raw(), self._bio_raw(), log_backend="none"
+        )
+
+        assert set(assembled.train_datasets) == {"omics_a", "bio_x"}
+        assert set(assembled.losses) == {"omics_a", "bio_x"}
+        # Both datasets contribute an evaluator (omics from val, bio from train slice).
+        assert len(assembled.evaluators) == 2
+
+        # Omics anchors are namespaced omics ids; bio anchors are plain text.
+        assert assembled.train_datasets["omics_a"]["anchor"][0] == "omics:omics_a:0"
+        assert assembled.train_datasets["bio_x"]["anchor"][0] == "a neuron"
+
+        # Bio dataset keeps anchor/positive/negative and drops the stray column.
+        assert set(assembled.train_datasets["bio_x"].column_names) == {"anchor", "positive", "negative"}
+
+    def test_multi_dataset_training_step(self, pipeline, tmp_path):
+        merged = build_namespaced_vector_store(
+            {"omics_a": (_make_store(["0", "1", "2", "3"], tmp_path, name="oa2")[0], ["0", "1", "2", "3"])},
+            output_path=os.path.join(tmp_path, "merged2.mmap"),
+        )
+        pipeline[0].set_vector_store(merged)
+
+        assembled = assemble_training_data(
+            self._cfg(), pipeline, self._omics_raw(), self._bio_raw(), log_backend="none"
+        )
+        evaluator = SequentialEvaluator(assembled.evaluators)
+
+        args = SentenceTransformerTrainingArguments(
+            output_dir=os.path.join(tmp_path, "run"),
+            num_train_epochs=1,
+            per_device_train_batch_size=2,
+            learning_rate=1e-3,
+            use_cpu=True,
+            report_to="none",
+        )
+        trainer = SentenceTransformerTrainer(
+            model=pipeline,
+            args=args,
+            train_dataset=assembled.train_datasets,
+            loss=assembled.losses,
+            evaluator=evaluator,
+        )
+
+        before = {n: p.clone() for n, p in pipeline.named_parameters() if p.requires_grad}
+        trainer.train()
+        delta = sum((p - before[n]).abs().sum().item() for n, p in pipeline.named_parameters() if n in before)
+        assert delta > 0, "No parameters changed during multi-dataset training"
+
+
+# ---------------------------------------------------------------------------
+# 3. Namespacing resolves end-to-end
+# ---------------------------------------------------------------------------
+class TestNamespacingResolves:
+    def test_anchors_resolve_in_store(self, pipeline, tmp_path):
+        merged = build_namespaced_vector_store(
+            {"omics_a": (_make_store(["0", "1", "2", "3"], tmp_path, name="oa3")[0], ["0", "1", "2", "3"])},
+            output_path=os.path.join(tmp_path, "merged3.mmap"),
+        )
+        pipeline[0].set_vector_store(merged)
+
+        cfg = TrainConfig(
+            omics_datasets=[DatasetConfig(id="dummy/a", name="omics_a", modality="bimodal")],
+            trainer=TrainerConfig(per_device_eval_batch_size=2, logging_steps=1),
+        )
+        omics_raw = {"omics_a": DatasetDict(train=_omics_split(["0", "1", "2", "3"]))}
+        assembled = assemble_training_data(cfg, pipeline, omics_raw, log_backend="none")
+
+        anchors = assembled.train_datasets["omics_a"]["anchor"]
+        for anchor in anchors:
+            assert anchor.startswith("omics:omics_a:")
+            assert anchor[len("omics:") :] in merged  # store key present
+
+        emb = pipeline.encode(anchors)
+        assert emb.shape == (len(anchors), SHARED_DIM)
+
+
+# ---------------------------------------------------------------------------
+# 4. Config parsing
+# ---------------------------------------------------------------------------
+class TestConfig:
+    def test_config_from_dict(self):
+        cfg = config_from_dict(
+            {
+                "obsm_key": "X_pca",
+                "shared_dim": 128,
+                "text_encoder": {"name": "some/model", "max_seq_length": 256},
+                "omics_datasets": [{"id": "a/b", "name": "d1", "modality": "bimodal"}],
+                "bio_datasets": [{"id": "c/d", "name": "bio", "revision": "main"}],
+                "trainer": {"num_train_epochs": 5},
+            }
+        )
+        assert cfg.obsm_key == "X_pca"
+        assert cfg.shared_dim == 128
+        assert cfg.text_encoder.max_seq_length == 256
+        assert cfg.omics_datasets[0].name == "d1"
+        assert cfg.bio_datasets[0].revision == "main"
+        assert cfg.trainer.num_train_epochs == 5
+
+    def test_config_ignores_unknown_keys(self):
+        cfg = config_from_dict({"obsm_key": "X_pca", "bogus_key": 123})
+        assert cfg.obsm_key == "X_pca"
+        assert not hasattr(cfg, "bogus_key")
+
+    def test_mining_defaults_to_none(self):
+        cfg = config_from_dict({"obsm_key": "X_pca"})
+        assert cfg.mining is None
+
+    def test_mining_parsed_with_overrides(self):
+        cfg = config_from_dict(
+            {
+                "mining": {
+                    "enabled": True,
+                    "num_negatives": 5,
+                    "range_max": 50,
+                    "max_score": 0.9,
+                    "sampling_strategy": "random",
+                    "use_faiss": True,
+                }
+            }
+        )
+        assert isinstance(cfg.mining, MiningConfig)
+        assert cfg.mining.enabled is True
+        assert cfg.mining.num_negatives == 5
+        assert cfg.mining.range_max == 50
+        assert cfg.mining.max_score == 0.9
+        assert cfg.mining.sampling_strategy == "random"
+        assert cfg.mining.use_faiss is True
+        # Untouched fields keep their defaults.
+        assert cfg.mining.range_min == 1
+        assert cfg.mining.relative_margin is None
+
+
+# ---------------------------------------------------------------------------
+# 5. Model-card training details
+# ---------------------------------------------------------------------------
+class TestTrainingDetails:
+    def test_details_without_mining(self):
+        cfg = TrainConfig(
+            model="jo-mengr/pretrained",
+            obsm_key="X_scvi_fm",
+            omics_datasets=[DatasetConfig(id="a/b", name="d1")],
+        )
+        md = format_training_details(cfg)
+        assert "Continued from" in md and "jo-mengr/pretrained" in md
+        assert "not used" in md  # mining disabled
+        assert "X_scvi_fm" in md
+
+    def test_details_with_mining(self):
+        cfg = TrainConfig(
+            model="outputs/pretrained/final",
+            mining=MiningConfig(enabled=True, num_negatives=4, use_faiss=True),
+            omics_datasets=[DatasetConfig(id="a/b", name="d1")],
+            bio_datasets=[DatasetConfig(id="c/d", name="bio")],
+        )
+        md = format_training_details(cfg)
+        assert "enabled" in md
+        assert "num_negatives`: 4" in md
+        assert "use_faiss`: True" in md
+        assert "mine_hard_negatives" in md
+        assert "d1, bio" in md  # training datasets list
+
+
+# ---------------------------------------------------------------------------
+# 6. Mining-enabled assembly
+# ---------------------------------------------------------------------------
+class TestMiningAssembly:
+    def test_mining_replaces_dataset_negatives(self, pipeline, tmp_path):
+        merged = build_namespaced_vector_store(
+            {"omics_a": (_make_store(["0", "1", "2", "3"], tmp_path, name="om")[0], ["0", "1", "2", "3"])},
+            output_path=os.path.join(tmp_path, "merged_mine.mmap"),
+        )
+        pipeline[0].set_vector_store(merged)
+
+        cfg = TrainConfig(
+            obsm_key="X_scvi_fm",
+            shared_dim=SHARED_DIM,
+            # use_hard_negatives=True in the dataset must be overridden by mining.
+            omics_datasets=[
+                DatasetConfig(id="dummy/a", name="omics_a", modality="bimodal", use_hard_negatives=True),
+            ],
+            mining=MiningConfig(enabled=True, num_negatives=1, range_min=0, max_score=None),
+            trainer=TrainerConfig(per_device_train_batch_size=2, per_device_eval_batch_size=2, logging_steps=1),
+        )
+        omics_raw = {
+            "omics_a": DatasetDict(
+                train=_omics_split(["0", "1", "2", "3"]),
+                val=_omics_split(["0", "1", "2", "3"]),
+            )
+        }
+        assembled = assemble_training_data(cfg, pipeline, omics_raw, log_backend="none")
+
+        train = assembled.train_datasets["omics_a"]
+        # Mining produces an n-tuple with mined negative_1 (not the dataset's negative_1/2).
+        assert "negative_1" in train.column_names
+        assert "anchor" in train.column_names and "positive" in train.column_names
+        assert len(train) > 0
+        # Omics anchors are still resolvable namespaced ids.
+        assert train["anchor"][0].startswith("omics:omics_a:")
+        # Eval set was mined too, so the evaluator has a usable negative column.
+        assert len(assembled.evaluators) == 1
diff --git a/tutorials/evaluate_model_2.0.ipynb b/tutorials/evaluate_model_2.0.ipynb
new file mode 100644
index 0000000..ceb014d
--- /dev/null
+++ b/tutorials/evaluate_model_2.0.ipynb
@@ -0,0 +1,672 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This tutorial shows how to evaluate an MMContext Sentence Transformers model on one dataset. It assumes you created a huggingface dataset, which contains the cell representations (either cell ids for numerical embeddings or cell sentences for text_only usage). Such datasets can be created with a pipeline available through the https://github.com/mengerj/adata_hf_datasets repo. If you instead want to start from an adata object, see the tutorial pretrained_inference.ipynb"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Figure 1D in the publication was created with this notebook"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "repo_name = \"jo-mengr\"\n",
+    "dataset_name = \"hiha_100k\"\n",
+    "split_name = \"test\"\n",
+    "label_key = \"AIFI_L2\"  # \"AIFI_L2\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = load_dataset(f\"{repo_name}/{dataset_name}\")\n",
+    "test_dataset = dataset[split_name]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sentence_transformers import SentenceTransformer\n",
+    "\n",
+    "# model_name = \"jo-mengr/mmcontext-pubmedbert-gs10k\"\n",
+    "model_name = \"../outputs/multi_gs10k/final\"\n",
+    "model = SentenceTransformer(model_name, trust_remote_code=True)\n",
+    "data_type = \"gs10k\"\n",
+    "layer_key = f\"X_{data_type}\"\n",
+    "text_only = False  # set True for text-based models (uses cell_sentence_2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "47464e182bd94f42a6ab503b16a8a92c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Processing:   0%|          | 0/1 [00:00<?, ?file/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "aa04d1cf7c7b4797993a6a5c025863ff",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Processing adata files:   0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "from mmcontext.embed import prepare_inference\n",
+    "\n",
+    "# Cache + store locations are derived from the knobs set at the top of the\n",
+    "# notebook — nothing to edit here:\n",
+    "#   * adata cache depends only on the dataset (same raw data across models/keys)\n",
+    "#   * the vector store is keyed by model + dataset + embedding key, so switching\n",
+    "#     the model or layer_key never silently reuses a stale store.\n",
+    "model_slug = model_name.strip(\"/\").replace(\"/\", \"_\")\n",
+    "adata_cache_dir = Path(\"../data/test_adata\") / dataset_name\n",
+    "store_path = Path(\"../outputs/inference_stores\") / model_slug / f\"{dataset_name}__{layer_key}.mmap\"\n",
+    "\n",
+    "# One call replaces: load_test_adata_from_hf_dataset + subset_dataset_by_chunk\n",
+    "# + prepare_vector_store + set_vector_store (bimodal) / truncation (text).\n",
+    "# The cell sentence is chosen internally from the modality\n",
+    "# (cell_sentence_1 for bimodal omics, cell_sentence_2 for text).\n",
+    "modality = \"text\" if text_only else \"bimodal\"\n",
+    "bundle = prepare_inference(\n",
+    "    model,\n",
+    "    test_dataset,\n",
+    "    modality=modality,\n",
+    "    obsm_key=layer_key,\n",
+    "    cache_dir=str(adata_cache_dir),\n",
+    "    store_path=str(store_path),\n",
+    "    zenodo_token=os.getenv(\"ZENODO_TOKEN\"),\n",
+    "    truncate=text_only,\n",
+    "    truncate_kwargs={\"max_length\": 64, \"filter_strings\": [\"RPS\", \"RPL\", \"MT\"]},\n",
+    ")\n",
+    "adata = bundle.adata\n",
+    "dataset_ready = bundle.dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_to_use = dataset_ready  # [split_name]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_encoder_name = model[0].model_name_or_path\n",
+    "text_encoder = SentenceTransformer(text_encoder_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'anchor': 'omics:cfe6b50ea2e611eb91855e3d154f7c3f',\n",
+       " 'sample_idx': 'cfe6b50ea2e611eb91855e3d154f7c3f',\n",
+       " 'adata_link': 'https://zenodo.org/api/records/17715306/draft/files/all_chunk_0.zarr.zip/content'}"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dataset_to_use[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "omics_embeddings = model.encode(dataset_to_use[\"anchor\"])\n",
+    "adata.obsm[\"mmcontext_emb\"] = omics_embeddings"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_colours = len(adata.obs[\"AIFI_L1\"].unique())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import seaborn as sns\n",
+    "\n",
+    "auto_colors = sns.color_palette(\"tab10\", n_colours)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<svg  width=\"495\" height=\"55\"><rect x=\"0\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#1f77b4;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"55\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#ff7f0e;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"110\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#2ca02c;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"165\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#d62728;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"220\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#9467bd;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"275\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#8c564b;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"330\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#e377c2;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"385\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#7f7f7f;stroke-width:2;stroke:rgb(255,255,255)\"/><rect x=\"440\" y=\"0\" width=\"55\" height=\"55\" style=\"fill:#bcbd22;stroke-width:2;stroke:rgb(255,255,255)\"/></svg>"
+      ],
+      "text/plain": [
+       "[(0.12156862745098039, 0.4666666666666667, 0.7058823529411765),\n",
+       " (1.0, 0.4980392156862745, 0.054901960784313725),\n",
+       " (0.17254901960784313, 0.6274509803921569, 0.17254901960784313),\n",
+       " (0.8392156862745098, 0.15294117647058825, 0.1568627450980392),\n",
+       " (0.5803921568627451, 0.403921568627451, 0.7411764705882353),\n",
+       " (0.5490196078431373, 0.33725490196078434, 0.29411764705882354),\n",
+       " (0.8901960784313725, 0.4666666666666667, 0.7607843137254902),\n",
+       " (0.4980392156862745, 0.4980392156862745, 0.4980392156862745),\n",
+       " (0.7372549019607844, 0.7411764705882353, 0.13333333333333333)]"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "auto_colors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "label_colors = {\n",
+    "    \"T cell\": auto_colors[7],\n",
+    "    \"B cell\": auto_colors[0],\n",
+    "    \"NK cell\": auto_colors[1],\n",
+    "    \"Monocyte\": auto_colors[2],\n",
+    "    \"DC\": auto_colors[3],\n",
+    "    \"Platelet\": auto_colors[4],\n",
+    "    \"Progenitor cell\": auto_colors[5],\n",
+    "    \"ILC\": auto_colors[6],\n",
+    "    \"Erythrocyte\": auto_colors[8],\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mmcontext.eval import get\n",
+    "\n",
+    "EvClass = get(\"LabelSimilarity\")\n",
+    "ev = EvClass(\n",
+    "    auto_filter_labels=False,\n",
+    "    umap_n_neighbors=10,\n",
+    "    umap_min_dist=0.4,\n",
+    "    similarity=\"cosine\",\n",
+    "    logit_scale=1,\n",
+    "    score_norm_method=None,\n",
+    "    label_colors=None,\n",
+    "    annotation_fontsize=16,\n",
+    "    font_family=\"Arial\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# precompute umap coordinates to reuse on subset\n",
+    "full_omics_embeddings = adata.obsm[\"mmcontext_emb\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# full_cell_umap = ev._compute_umap(full_omics_embeddings)\n",
+    "# add umap coordinates to adata\n",
+    "# adata.obsm[\"cell_umap\"] = full_cell_umap"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "full_query_labels = adata.obs[label_key].unique().tolist()\n",
+    "full_label_embeddings = model.encode(full_query_labels)\n",
+    "full_true_labels = adata.obs[label_key]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/mengerj/repos/mmcontext/src/mmcontext/eval/label_similarity.py:76: RuntimeWarning: divide by zero encountered in matmul\n",
+      "  return (emb1_norm @ emb2_norm.T).astype(np.float32)\n",
+      "/Users/mengerj/repos/mmcontext/src/mmcontext/eval/label_similarity.py:76: RuntimeWarning: overflow encountered in matmul\n",
+      "  return (emb1_norm @ emb2_norm.T).astype(np.float32)\n",
+      "/Users/mengerj/repos/mmcontext/src/mmcontext/eval/label_similarity.py:76: RuntimeWarning: invalid value encountered in matmul\n",
+      "  return (emb1_norm @ emb2_norm.T).astype(np.float32)\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "result = ev.compute(\n",
+    "    omics_embeddings=full_omics_embeddings,\n",
+    "    label_embeddings=full_label_embeddings,\n",
+    "    query_labels=full_query_labels,\n",
+    "    true_labels=full_true_labels,\n",
+    "    label_key=label_key,\n",
+    "    out_dir=Path(f\"LabelSimilarity/{model_name}/{dataset_name}\"),  # Pass output directory for caching\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ASDC/auc: 0.4956\n",
+       "ASDC/accuracy: 0.0000\n",
+       "CD14 monocyte/auc: 0.9785\n",
+       "CD14 monocyte/accuracy: 0.9417\n",
+       "CD16 monocyte/auc: 0.9181\n",
+       "CD16 monocyte/accuracy: 0.0009\n",
+       "CD56bright NK cell/auc: 0.9277\n",
+       "CD56bright NK cell/accuracy: 0.6163\n",
+       "CD56dim NK cell/auc: 0.9499\n",
+       "CD56dim NK cell/accuracy: 0.0829\n",
+       "CD8aa/auc: 0.6865\n",
+       "CD8aa/accuracy: 0.0000\n",
+       "DN T cell/auc: 0.7996\n",
+       "DN T cell/accuracy: 0.0089\n",
+       "Effector B cell/auc: 0.9421\n",
+       "Effector B cell/accuracy: 0.0516\n",
+       "Erythrocyte/auc: 0.8086\n",
+       "Erythrocyte/accuracy: 0.6000\n",
+       "ILC/auc: 0.4408\n",
+       "ILC/accuracy: 0.0000\n",
+       "Intermediate monocyte/auc: 0.9244\n",
+       "Intermediate monocyte/accuracy: 0.0549\n",
+       "MAIT/auc: 0.7842\n",
+       "MAIT/accuracy: 0.0000\n",
+       "Memory B cell/auc: 0.9552\n",
+       "Memory B cell/accuracy: 0.0521\n",
+       "Memory CD4 T cell/auc: 0.7963\n",
+       "Memory CD4 T cell/accuracy: 0.7317\n",
+       "Memory CD8 T cell/auc: 0.8231\n",
+       "Memory CD8 T cell/accuracy: 0.1451\n",
+       "Naive B cell/auc: 0.9625\n",
+       "Naive B cell/accuracy: 0.3367\n",
+       "Naive CD4 T cell/auc: 0.7379\n",
+       "Naive CD4 T cell/accuracy: 0.4369\n",
+       "Naive CD8 T cell/auc: 0.6935\n",
+       "Naive CD8 T cell/accuracy: 0.0002\n",
+       "Plasma cell/auc: 0.9708\n",
+       "Plasma cell/accuracy: 1.0000\n",
+       "Platelet/auc: 0.9996\n",
+       "Platelet/accuracy: 1.0000\n",
+       "Progenitor cell/auc: 0.9825\n",
+       "Progenitor cell/accuracy: 0.3919\n",
+       "Proliferating NK cell/auc: 0.9346\n",
+       "Proliferating NK cell/accuracy: 0.7214\n",
+       "Proliferating T cell/auc: 0.7610\n",
+       "Proliferating T cell/accuracy: 0.5238\n",
+       "Transitional B cell/auc: 0.9379\n",
+       "Transitional B cell/accuracy: 0.0126\n",
+       "Treg/auc: 0.6738\n",
+       "Treg/accuracy: 0.0086\n",
+       "cDC1/auc: 0.8383\n",
+       "cDC1/accuracy: 0.0000\n",
+       "cDC2/auc: 0.7170\n",
+       "cDC2/accuracy: 0.0000\n",
+       "gdT/auc: 0.6811\n",
+       "gdT/accuracy: 0.0000\n",
+       "pDC/auc: 0.6160\n",
+       "pDC/accuracy: 0.0000\n",
+       "mean_auc: 0.8185\n",
+       "std_auc: 0.1468\n",
+       "accuracy: 0.4123\n",
+       "balanced_accuracy: 0.2661\n",
+       "random_baseline_accuracy: 0.0345\n",
+       "accuracy_over_random: 11.9570\n",
+       "n_labels: 29\n",
+       "topk_accuracy@1: 0.4123\n",
+       "topk_accuracy@3: 0.7662\n",
+       "topk_accuracy@5: 0.8706\n",
+       "macro_f1: 0.1717\n",
+       "mrr: 0.6091"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ev.plot(\n",
+    "    omics_embeddings=full_omics_embeddings,\n",
+    "    # cell_umap=full_cell_umap,\n",
+    "    out_dir=Path(f\"LabelSimilarity/{model_name}/{dataset_name}/{label_key}_combined\"),\n",
+    "    label_embeddings=full_label_embeddings,\n",
+    "    query_labels=full_query_labels,\n",
+    "    true_labels=full_true_labels,\n",
+    "    label_key=label_key,  # column name (e.g. \"celltype\")\n",
+    "    save_format=\"svg\",\n",
+    "    figsize=(2.5, 2.5),\n",
+    "    dpi=300,\n",
+    "    font_size=12,\n",
+    "    font_style=\"normal\",\n",
+    "    font_weight=\"normal\",\n",
+    "    legend_fontsize=54,\n",
+    "    axis_label_size=20,\n",
+    "    axis_tick_size=12,\n",
+    "    point_size=0.25,\n",
+    "    legend_layout=\"vertical\",\n",
+    "    legend_point_size=16,\n",
+    "    umap_method=\"combined\",\n",
+    "    label_min_distance=0.2,\n",
+    "    label_spring_strength=0.5,\n",
+    "    label_repulsion_strength=1,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Option to subset adata based on one or more label values (e.g., \"Monocyte\" and \"DC\")\n",
+    "subset_label_values = [\"T cell\"]  # Change this list to your desired label values\n",
+    "subset_label_key = \"AIFI_L1\"\n",
+    "annotation_label_key = \"AIFI_L2\"\n",
+    "# Subset the AnnData object for any of the specified label values\n",
+    "adata_subset = adata[adata.obs[subset_label_key].isin(subset_label_values)].copy()\n",
+    "subset_labels = adata_subset.obs[annotation_label_key].values.unique()\n",
+    "label_embeddings_subset = model.encode(subset_labels)\n",
+    "# Create a new LabelSimilarity evaluator instance\n",
+    "# ev_subset = EvClass(auto_filter_labels=False, umap_n_neighbors=15, umap_min_dist=0.5)\n",
+    "subset_label_string = \"_\".join(subset_label_values)\n",
+    "subset_omics_embeddings = adata_subset.obsm[\"mmcontext_emb\"]\n",
+    "# subset_umap_coords = adata_subset.obsm[\"cell_umap\"]\n",
+    "# ev_subset.eb_lfdr_q = 0.01\n",
+    "ev = EvClass(\n",
+    "    auto_filter_labels=False,\n",
+    "    umap_n_neighbors=10,\n",
+    "    umap_min_dist=0.4,\n",
+    "    similarity=\"cosine\",\n",
+    "    logit_scale=1,\n",
+    "    score_norm_method=None,\n",
+    "    font_family=\"Arial\",\n",
+    "    annotation_fontsize=18,\n",
+    ")\n",
+    "# Compute metrics on the subsetted data\n",
+    "result_subset = ev.compute(\n",
+    "    omics_embeddings=subset_omics_embeddings,\n",
+    "    label_embeddings=label_embeddings_subset,\n",
+    "    query_labels=subset_labels,\n",
+    "    true_labels=adata_subset.obs[annotation_label_key],\n",
+    "    label_key=annotation_label_key,\n",
+    "    out_dir=Path(\n",
+    "        f\"LabelSimilarity/{model_name}/{dataset_name}/{annotation_label_key}_subset_{subset_label_string}/results\"\n",
+    "    ),\n",
+    ")\n",
+    "\n",
+    "# Plot results for the subset\n",
+    "ev.plot(\n",
+    "    omics_embeddings=subset_omics_embeddings,\n",
+    "    #    cell_umap=subset_umap_coords,\n",
+    "    out_dir=Path(\n",
+    "        f\"LabelSimilarity/{model_name}/{dataset_name}/{annotation_label_key}_subset_{subset_label_string}_combined\"\n",
+    "    ),\n",
+    "    label_embeddings=label_embeddings_subset,\n",
+    "    query_labels=subset_labels,\n",
+    "    true_labels=adata_subset.obs[annotation_label_key],\n",
+    "    label_key=annotation_label_key,\n",
+    "    save_format=\"svg\",\n",
+    "    figsize=(2.5, 2.5),\n",
+    "    dpi=300,\n",
+    "    font_size=12,\n",
+    "    axis_tick_size=12,\n",
+    "    font_style=\"normal\",\n",
+    "    font_weight=\"normal\",\n",
+    "    axis_label_size=20,\n",
+    "    point_size=0.25,\n",
+    "    legend_layout=\"vertical\",\n",
+    "    legend_point_size=20,\n",
+    "    umap_method=\"combined\",\n",
+    "    label_min_distance=0.15,\n",
+    "    label_spring_strength=0.35,\n",
+    "    label_repulsion_strength=1.4,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_subset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Visualise the embeddings\n",
+    "from mmcontext.pl import plot_umap\n",
+    "from mmcontext.utils import consolidate_low_frequency_categories\n",
+    "\n",
+    "current_key = label_key\n",
+    "adata_cut = consolidate_low_frequency_categories(adata, [current_key], threshold=50, remove=True)\n",
+    "emb_key = \"mmcontext_emb\"\n",
+    "plot_umap(\n",
+    "    adata,\n",
+    "    color_key=label_key,\n",
+    "    embedding_key=emb_key,\n",
+    "    save_format=\"svg\",\n",
+    "    save_dir=f\"figs/{model_name}/{dataset_name}\",\n",
+    "    save_plot=False,\n",
+    "    title=\"\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Visualise the embeddings\n",
+    "from mmcontext.pl import plot_umap\n",
+    "from mmcontext.utils import consolidate_low_frequency_categories\n",
+    "\n",
+    "current_key = label_key\n",
+    "adata_cut = consolidate_low_frequency_categories(adata, [current_key], threshold=1, remove=False)\n",
+    "emb_key = layer_key\n",
+    "plot_umap(\n",
+    "    adata_cut,\n",
+    "    color_key=label_key,\n",
+    "    embedding_key=emb_key,\n",
+    "    save_format=\"svg\",\n",
+    "    nametag=\"\",\n",
+    "    save_dir=f\"figs/{model_name}/{dataset_name}\",\n",
+    "    save_plot=False,\n",
+    "    title=\"\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mmcontext.eval.query_annotate import OmicsQueryAnnotator\n",
+    "\n",
+    "annotator = OmicsQueryAnnotator(model)\n",
+    "annotator.annotate_omics_data(adata, full_query_labels, emb_key=\"mmcontext_emb\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# get accuracy of best label vs true label\n",
+    "from sklearn.metrics import accuracy_score\n",
+    "\n",
+    "accuracy_score(adata.obs[\"best_label\"], adata.obs[label_key])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "queries_csv = \"../../data/queries/additional_combined.csv\"\n",
+    "if dataset_name == \"hiha_100k\" and os.path.exists(queries_csv):\n",
+    "    df = pd.read_csv(queries_csv)\n",
+    "    labels = df[\"Cell Type\"]\n",
+    "    Definition = df[\"Definition\"]\n",
+    "    from mmcontext.eval.query_annotate import OmicsQueryAnnotator\n",
+    "    from mmcontext.pl.plotting import plot_query_scores_with_labels_umap\n",
+    "\n",
+    "    annotator = OmicsQueryAnnotator(model)\n",
+    "    annotator.query_with_text(adata, Definition, emb_key=\"mmcontext_emb\")\n",
+    "    # Call the plotting function\n",
+    "    plot_query_scores_with_labels_umap(\n",
+    "        adata=adata,\n",
+    "        queries=Definition,\n",
+    "        labels=labels,\n",
+    "        label_key=\"AIFI_L2\",\n",
+    "        save_dir=f\"figs/{model_name}/{dataset_name}/umap_with_labels\",\n",
+    "        nametag=\"\",\n",
+    "        figsize=(4, 4),\n",
+    "        point_size=2,\n",
+    "        dpi=300,  # Lower DPI for faster generation\n",
+    "        axis_label_size=18,\n",
+    "        axis_tick_size=18,\n",
+    "    )"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mmcontext (3.11.14)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/tutorials/train_new.ipynb b/tutorials/train_new.ipynb
index b876bc4..029649e 100644
--- a/tutorials/train_new.ipynb
+++ b/tutorials/train_new.ipynb
@@ -1,510 +1,509 @@
 {
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "id": "641d7ea1",
-      "metadata": {},
-      "source": [
-        "## Train a new mmcontext model (demonstrated on proteomics data)\n",
-        "\n",
-        "This tutorial walks through the core idea of **mmcontext**: aligning a numerical representation of a sample (e.g. a vector derived from omics data) with its **biological context** (metadata such as tissue, disease, cell type, perturbations, etc.).\n",
-        "\n",
-        "At a high level:\n",
-        "- We start from a single-cell proteomics dataset.\n",
-        "- For each cell, we build a **numerical representation** (\"data embedding\") and a **textual description** summarising its biological context.\n",
-        "- We then train an `MMContextEncoder` so that samples with similar biology end up close in the shared embedding space, even across different data modalities.\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "11c910e0",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from mmcontext.utils import setup_logging\n",
-        "\n",
-        "setup_logging()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "6eccf0fc",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import anndata as ad\n",
-        "\n",
-        "from mmcontext.file_utils import download_file_from_share_link\n",
-        "\n",
-        "# a protein dataset from figshare (https://plus.figshare.com/articles/dataset/scPerturb_Single-Cell_Perturbation_Data_RNA_and_protein_h5ad_files/24160713)\n",
-        "# Note: use ndownloader.figshare.com (not plus.figshare.com) for programmatic downloads - plus.figshare.com is behind WAF that blocks automated requests\n",
-        "data_link = \"https://ndownloader.figshare.com/files/42428325\"\n",
-        "local_path = \"Frangiehlzar2021_protein.h5ad\"\n",
-        "# download the data\n",
-        "download_file_from_share_link(share_link=data_link, save_path=local_path)\n",
-        "# load the data\n",
-        "adata = ad.read_h5ad(local_path)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "72302131",
-      "metadata": {},
-      "source": [
-        "### Step 1: Describe each cell in words\n",
-        "\n",
-        "The general idea of the framework is to **link biology-rich text to data-rich numeric vectors**:\n",
-        "\n",
-        "- The **biological context** may include perturbations, tissue, disease status, cell type and other metadata.\n",
-        "- The **numerical representation** should summarise the measured data for that sample (here: protein expression; in other settings: scRNA-seq, ATAC, etc.).\n",
-        "\n",
-        "In this simple example we construct a short textual description for each cell using columns from `adata.obs`. For real applications you might want much richer descriptions (e.g. including time points, doses, quality metrics, etc.), but the mechanism is the same."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "95f28a07",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Create a description string for each cell by looping over the rows in adata.obs\n",
-        "# This text will serve as the \"biological context\" that we want to align with the numeric representation.\n",
-        "def make_description(row):\n",
-        "    \"\"\"Make a quick description of a cell based on its metadata\"\"\"\n",
-        "    return (\n",
-        "        f\"The first perturbation is {row['perturbation']} \"\n",
-        "        f\"and the second perturbation is {row['perturbation_2']}.\"\n",
-        "        f\" The tissue is {row['tissue_type']} and it has cancer yes or no: {row['cancer']}.\"\n",
-        "        f\" The disease is {row['disease']}.\"\n",
-        "        f\" The celltype is {row['celltype']}.\"\n",
-        "    )\n",
-        "\n",
-        "\n",
-        "# Add a \"description\" column to adata.obs using the function above\n",
-        "adata.obs[\"description\"] = adata.obs.apply(make_description, axis=1)\n",
-        "# Also add a sample index column that uniquely identifies each cell; this will be used as a token later\n",
-        "adata.obs[\"sample_idx\"] = adata.obs.index"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "8c513a34",
-      "metadata": {},
-      "source": [
-        "### Step 2: Choose a numerical representation\n",
-        "\n",
-        "In this proteomics example, the assay only measures **24 proteins per cell**, so we can directly use the normalised expression vector as the numeric representation:\n",
-        "\n",
-        "- Each cell gets a 24-dimensional vector from `adata.X`.\n",
-        "- We store this in `adata.obsm[\"X_prot\"]` and later tell the model to use this key.\n",
-        "\n",
-        "For single-cell RNA datasets, mmcontext is more flexible:\n",
-        "\n",
-        "- You can plug in **precomputed representations** such as\n",
-        "  - Principal components (PCA),\n",
-        "  - Latent variables from **scVI**,\n",
-        "  - Embeddings from **Geneformer** or other pretrained models.\n",
-        "- In the publication, we explore exactly these kinds of pretrained embeddings as the numeric representation `X_*` for each cell."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "6d999832",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import numpy as np\n",
-        "\n",
-        "# let's split by train and val, randomly 80% train\n",
-        "adata.obs[\"split\"] = np.random.rand(len(adata)) < 0.8\n",
-        "adata_train = adata[adata.obs[\"split\"]].copy()\n",
-        "adata_val = adata[~adata.obs[\"split\"]].copy()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "8de1a9c1",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "adata_val.shape"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "d68558ad",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "adata_train.shape"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "7cbfa52a",
-      "metadata": {},
-      "source": [
-        "### Step 2b: Split into train/validation and preprocess\n",
-        "\n",
-        "Before building embeddings, we:\n",
-        "\n",
-        "- Randomly split cells into an **80% train** and **20% validation** subset using a boolean mask in `adata.obs[\"split\"]`.\n",
-        "- Apply standard **normalisation and log-transformation** with Scanpy (`sc.pp.normalize_total` + `sc.pp.log1p`).\n",
-        "\n",
-        "These steps ensure that the numeric representation we later feed into the model is on a comparable scale across cells, and that we can monitor generalisation on a held-out validation set."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "3ffb5747",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import scanpy as sc\n",
-        "\n",
-        "# normalise and log transform the data\n",
-        "sc.pp.normalize_total(adata_train, inplace=True)\n",
-        "sc.pp.log1p(adata_train)\n",
-        "\n",
-        "sc.pp.normalize_total(adata_val, inplace=True)\n",
-        "sc.pp.log1p(adata_val)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "b699d36c",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# since these datasets only contain 24 proteins, we will just use their expression as embeddings\n",
-        "# we will use the protein names as the embedding keys\n",
-        "adata_train.obsm[\"X_prot\"] = adata_train.X\n",
-        "adata_val.obsm[\"X_prot\"] = adata_val.X\n",
-        "processed_paths = {\"train\": \"Frangiehlzar2021_protein_pp_train.h5ad\", \"val\": \"Frangiehlzar2021_protein_pp_val.h5ad\"}\n",
-        "adata_train.write_h5ad(processed_paths[\"train\"])\n",
-        "adata_val.write_h5ad(processed_paths[\"val\"])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "09af0ef7",
-      "metadata": {},
-      "source": [
-        "### Step 3: Construct a contrastive training dataset\n",
-        "\n",
-        "So far we have:\n",
-        "- A **numeric representation** for each cell (stored in `adata.obsm[\"X_prot\"]`).\n",
-        "- A **text description** for each cell in `adata.obs[\"description\"]`.\n",
-        "\n",
-        "Next, we convert these into a Hugging Face `DatasetDict` suitable for contrastive training:\n",
-        "\n",
-        "- `AnnDataSetConstructor` in `\"multiplets\"` mode creates, for each cell, a\n",
-        "  - **query token** (`sample_idx` / later `cell_sentence_1`),\n",
-        "  - a **matched positive**, and\n",
-        "  - at least one **negative**.\n",
-        "- This gives us a flexible format where the same dataset can later support different training regimens (e.g. using text-only, numeric-only, or mixed representations), simply by choosing which columns to feed into the model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "2bd012cb",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from adata_hf_datasets import AnnDataSetConstructor\n",
-        "from datasets import DatasetDict\n",
-        "\n",
-        "ds_dict = DatasetDict()\n",
-        "# multiplets format for training datasets (with descriptions)\n",
-        "# A sentence key refers to the column in adata.obs that is used to represent the sample.\n",
-        "# For numeric data, we use the sample index, and later register the created embedding linked to their indices in the tokenizer\n",
-        "for split_name, adata_split in {\"train\": adata_train, \"val\": adata_val}.items():\n",
-        "    constructor = AnnDataSetConstructor(dataset_format=\"multiplets\", resolve_negatives=True)\n",
-        "    constructor.add_anndata(\n",
-        "        adata_split,\n",
-        "        caption_key=\"description\",\n",
-        "        sentence_keys=[\"sample_idx\"],\n",
-        "        adata_link=processed_paths[split_name],\n",
-        "        batch_key=\"library_preparation_protocol\",  # In this case all are from the same batch, but providing a batch key can improve batch integration by negative sampling\n",
-        "    )\n",
-        "    ds = constructor.get_dataset()\n",
-        "    ds_dict[split_name] = ds"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "1e4d531a",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "ds_dict"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "70612269",
-      "metadata": {},
-      "source": [
-        "## Configure the Model\n",
-        "\n",
-        "The `MMContextEncoder` is the **main model** in this framework. Conceptually, it is a custom Sentence Transformers module that:\n",
-        "\n",
-        "- Starts from a standard **text encoder** (here: `sentence-transformers/all-MiniLM-L6-v2`).\n",
-        "- Adds **adapters** that map\n",
-        "  - the text encoder output into a desired embedding dimension, and\n",
-        "  - the numeric representation (e.g. `X_prot`, PCA, scVI, Geneformer embeddings) into the **same** space.\n",
-        "- Lets you configure whether the text encoder is **frozen**, partially unfrozen (last *n* layers), or fully trainable.\n",
-        "\n",
-        "To make the numeric data usable by the model, we need to **register initial embeddings**:\n",
-        "\n",
-        "- Tokens in `cell_sentence_1` (e.g. `\"sample_idx:123\"`) should map to a vector from our numeric representation.\n",
-        "- Internally, the model builds a lookup table (`sample_idx` → integer id) and an embedding matrix (id → numeric vector), which can be kept frozen or made trainable.\n",
-        "\n",
-        "After this setup, training looks like standard Sentence Transformers training, but every sample has **both** a text description and a numeric embedding living in a shared space."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "6f04ed8b",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from sentence_transformers import SentenceTransformer\n",
-        "\n",
-        "from mmcontext.mmcontextencoder import MMContextEncoder\n",
-        "\n",
-        "enc = MMContextEncoder(\n",
-        "    text_encoder_name=\"sentence-transformers/all-MiniLM-L6-v2\",\n",
-        "    adapter_hidden_dim=128,\n",
-        "    adapter_output_dim=64,\n",
-        "    freeze_text_encoder=True,\n",
-        "    unfreeze_last_n_layers=2,\n",
-        "    output_token_embeddings=False,\n",
-        "    train_lookup=False,\n",
-        "    joint_adapter_hidden_dim=None,\n",
-        "    text_model_kwargs=None,\n",
-        "    use_text_adapter=True,\n",
-        ")\n",
-        "model = SentenceTransformer(modules=[enc])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "d0926c15",
-      "metadata": {},
-      "source": [
-        "### Register the initial numeric embeddings\n",
-        "\n",
-        "We now connect the numeric representations on disk to the model:\n",
-        "\n",
-        "- `get_initial_embeddings_from_adata_link` reads the preprocessed `.h5ad` files and builds a table with one row per `sample_idx` and its embedding from `X_prot`.\n",
-        "- `register_initial_embeddings` turns this table into an internal embedding matrix and lookup table inside the `MMContextEncoder`.\n",
-        "\n",
-        "After this step, whenever the model sees a token like `\"sample_idx:123\"` in the dataset, it can look up the correct numeric vector and combine it with the text pathway."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "ee9a6e96",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "token_df, _ = model[0].get_initial_embeddings_from_adata_link(\n",
-        "    ds_dict,\n",
-        "    layer_key=\"X_prot\",\n",
-        "    download_dir=\"data_cache\",\n",
-        "    axis=\"obs\",  # since we get embeddings from adata.obsm. We could also use \"varm\" and for example use an embedding for each protein\n",
-        ")\n",
-        "model[0].register_initial_embeddings(token_df, data_origin=\"prot\")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "6975ab72",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# the model expects a certain prefix on the cell tokens.\n",
-        "model[0].processor.prefix"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "52d74f63",
-      "metadata": {},
-      "source": [
-        "### Prepare token prefixes and final training columns\n",
-        "\n",
-        "At this point, the model knows about all numeric embeddings, but the dataset still uses the plain `sample_idx` strings.\n",
-        "\n",
-        "- `model[0].processor.prefix` defines the required prefix for these tokens (e.g. `\"sample_idx:\"`).\n",
-        "- `prefix_ds` rewrites the `cell_sentence_1` column so that each entry matches what the tokenizer and lookup table expect.\n",
-        "\n",
-        "Finally, we simplify the dataset to the three columns we need for contrastive learning:\n",
-        "\n",
-        "- `anchor`: the main cell token (formerly `cell_sentence_1`).\n",
-        "- `positive`: a biologically matched partner.\n",
-        "- `negative_1`: a contrastive negative.\n",
-        "\n",
-        "This might look a bit cumbersome, but it makes the dataset **reusable**: you can easily construct different views (text-only, numeric-only, different negatives) without having to rebuild everything from scratch."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "6fc91911",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# you could add this manually or use the method below\n",
-        "model[0].prefix_ds(ds_dict, columns_to_prefix=[\"cell_sentence_1\"])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "c09bac80",
-      "metadata": {},
-      "source": [
-        "## Training setup and objective\n",
-        "\n",
-        "With the dataset and model ready, we now configure the training loop using the standard Sentence Transformers API:\n",
-        "\n",
-        "- `SentenceTransformerTrainingArguments` controls the **optimisation hyperparameters** (epochs, batch size, learning rate, warmup, logging and checkpointing frequency).\n",
-        "- `SentenceTransformerTrainer` automatically picks an appropriate device (CPU, CUDA GPU, or Apple **MPS**) and handles mixed-precision flags (`fp16`/`bf16`) where supported.\n",
-        "- We use `MultipleNegativesRankingLoss`, which encourages anchors to be close to their positives and far from all other samples in the batch.\n",
-        "- `TripletEvaluator` periodically evaluates the model on the validation set using triplets `(anchor, positive, negative)` and reports an accuracy metric.\n",
-        "\n",
-        "Running `trainer.train()` then performs the actual optimisation and will log progress and evaluation metrics during training."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "f63b0e4b",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# lastly we have to drop some columns from the dataset and\n",
-        "# rename the main column to \"anchor\".\n",
-        "# you might think that this is a bit cumbersome, which it is.\n",
-        "# But this setup allowed for fleixble training,\n",
-        "# using either cell or feature level tokens, using text based cell sentences\n",
-        "# or numeric embeddings and resolving negatives\n",
-        "# to whatever column was chosen for training.\n",
-        "# That means that for a certain training run, the same dataset can be reused, and\n",
-        "# only modified differently. But in the end, it is a bit of work to set up.\n",
-        "ds_final = ds_dict.rename_column(\"cell_sentence_1\", \"anchor\")\n",
-        "ds_final = ds_final.remove_columns([\"sample_idx\", \"adata_link\", \"negative_1_idx\"])\n",
-        "ds_final"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "25c0a4e9",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments\n",
-        "\n",
-        "args = SentenceTransformerTrainingArguments(\n",
-        "    num_train_epochs=1,\n",
-        "    per_device_train_batch_size=128,\n",
-        "    per_device_eval_batch_size=128,\n",
-        "    learning_rate=1e-5,\n",
-        "    warmup_ratio=0.1,\n",
-        "    fp16=False,\n",
-        "    bf16=False,\n",
-        "    eval_strategy=\"steps\",\n",
-        "    eval_steps=100,\n",
-        "    save_strategy=\"steps\",\n",
-        "    save_steps=100,\n",
-        "    save_total_limit=1,\n",
-        "    max_grad_norm=1.0,\n",
-        "    logging_steps=10,\n",
-        "    run_name=\"protein_test\",\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "d44a49b5",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from sentence_transformers.evaluation import TripletEvaluator\n",
-        "from sentence_transformers.losses import MultipleNegativesRankingLoss\n",
-        "\n",
-        "loss = MultipleNegativesRankingLoss(model)\n",
-        "evaluator = TripletEvaluator(\n",
-        "    anchors=ds_final[\"val\"][\"anchor\"],\n",
-        "    positives=ds_final[\"val\"][\"positive\"],\n",
-        "    negatives=ds_final[\"val\"][\"negative_1\"],\n",
-        "    name=\"val\",\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "53d32063",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "trainer = SentenceTransformerTrainer(\n",
-        "    model=model,\n",
-        "    args=args,\n",
-        "    train_dataset=ds_final[\"train\"],\n",
-        "    eval_dataset=ds_final[\"val\"],\n",
-        "    loss=loss,\n",
-        "    evaluator=evaluator,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "84f83d70",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "trainer.train()"
-      ]
-    }
-  ],
-  "metadata": {
-    "kernelspec": {
-      "display_name": ".venv",
-      "language": "python",
-      "name": "python3"
-    },
-    "language_info": {
-      "codemirror_mode": {
-        "name": "ipython",
-        "version": 3
-      },
-      "file_extension": ".py",
-      "mimetype": "text/x-python",
-      "name": "python",
-      "nbconvert_exporter": "python",
-      "pygments_lexer": "ipython3",
-      "version": "3.11.14"
-    }
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "641d7ea1",
+   "metadata": {},
+   "source": [
+    "## Train a new mmcontext model (demonstrated on proteomics data)\n",
+    "\n",
+    "This tutorial walks through the core idea of **mmcontext**: aligning a numerical representation of a sample (e.g. a vector derived from omics data) with its **biological context** (metadata such as tissue, disease, cell type, perturbations, etc.).\n",
+    "\n",
+    "At a high level:\n",
+    "- We start from a single-cell proteomics dataset.\n",
+    "- For each cell, we build a **numerical representation** (\"data embedding\") and a **textual description** summarising its biological context.\n",
+    "- We then train an `MMContextEncoder` so that samples with similar biology end up close in the shared embedding space, even across different data modalities.\n"
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 5
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11c910e0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mmcontext.utils import setup_logging\n",
+    "\n",
+    "setup_logging()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6eccf0fc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import anndata as ad\n",
+    "\n",
+    "from mmcontext.file_utils import download_file_from_share_link\n",
+    "\n",
+    "# a protein dataset from figshare (https://plus.figshare.com/articles/dataset/scPerturb_Single-Cell_Perturbation_Data_RNA_and_protein_h5ad_files/24160713)\n",
+    "# Note: use ndownloader.figshare.com (not plus.figshare.com) for programmatic downloads - plus.figshare.com is behind WAF that blocks automated requests\n",
+    "data_link = \"https://ndownloader.figshare.com/files/42428325\"\n",
+    "local_path = \"Frangiehlzar2021_protein.h5ad\"\n",
+    "# download the data\n",
+    "download_file_from_share_link(share_link=data_link, save_path=local_path)\n",
+    "# load the data\n",
+    "adata = ad.read_h5ad(local_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "72302131",
+   "metadata": {},
+   "source": [
+    "### Step 1: Describe each cell in words\n",
+    "\n",
+    "The general idea of the framework is to **link biology-rich text to data-rich numeric vectors**:\n",
+    "\n",
+    "- The **biological context** may include perturbations, tissue, disease status, cell type and other metadata.\n",
+    "- The **numerical representation** should summarise the measured data for that sample (here: protein expression; in other settings: scRNA-seq, ATAC, etc.).\n",
+    "\n",
+    "In this simple example we construct a short textual description for each cell using columns from `adata.obs`. For real applications you might want much richer descriptions (e.g. including time points, doses, quality metrics, etc.), but the mechanism is the same."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "95f28a07",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a description string for each cell by looping over the rows in adata.obs\n",
+    "# This text will serve as the \"biological context\" that we want to align with the numeric representation.\n",
+    "def make_description(row):\n",
+    "    \"\"\"Make a quick description of a cell based on its metadata\"\"\"\n",
+    "    return (\n",
+    "        f\"The first perturbation is {row['perturbation']} \"\n",
+    "        f\"and the second perturbation is {row['perturbation_2']}.\"\n",
+    "        f\" The tissue is {row['tissue_type']} and it has cancer yes or no: {row['cancer']}.\"\n",
+    "        f\" The disease is {row['disease']}.\"\n",
+    "        f\" The celltype is {row['celltype']}.\"\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "# Add a \"description\" column to adata.obs using the function above\n",
+    "adata.obs[\"description\"] = adata.obs.apply(make_description, axis=1)\n",
+    "# Also add a sample index column that uniquely identifies each cell; this will be used as a token later\n",
+    "adata.obs[\"sample_idx\"] = adata.obs.index"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8c513a34",
+   "metadata": {},
+   "source": [
+    "### Step 2: Choose a numerical representation\n",
+    "\n",
+    "In this proteomics example, the assay only measures **24 proteins per cell**, so we can directly use the normalised expression vector as the numeric representation:\n",
+    "\n",
+    "- Each cell gets a 24-dimensional vector from `adata.X`.\n",
+    "- We store this in `adata.obsm[\"X_prot\"]` and later tell the model to use this key.\n",
+    "\n",
+    "For single-cell RNA datasets, mmcontext is more flexible:\n",
+    "\n",
+    "- You can plug in **precomputed representations** such as\n",
+    "  - Principal components (PCA),\n",
+    "  - Latent variables from **scVI**,\n",
+    "  - Embeddings from **Geneformer** or other pretrained models.\n",
+    "- In the publication, we explore exactly these kinds of pretrained embeddings as the numeric representation `X_*` for each cell."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d999832",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "# let's split by train and val, randomly 80% train\n",
+    "adata.obs[\"split\"] = np.random.rand(len(adata)) < 0.8\n",
+    "adata_train = adata[adata.obs[\"split\"]].copy()\n",
+    "adata_val = adata[~adata.obs[\"split\"]].copy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8de1a9c1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "adata_val.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d68558ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "adata_train.shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7cbfa52a",
+   "metadata": {},
+   "source": [
+    "### Step 2b: Split into train/validation and preprocess\n",
+    "\n",
+    "Before building embeddings, we:\n",
+    "\n",
+    "- Randomly split cells into an **80% train** and **20% validation** subset using a boolean mask in `adata.obs[\"split\"]`.\n",
+    "- Apply standard **normalisation and log-transformation** with Scanpy (`sc.pp.normalize_total` + `sc.pp.log1p`).\n",
+    "\n",
+    "These steps ensure that the numeric representation we later feed into the model is on a comparable scale across cells, and that we can monitor generalisation on a held-out validation set."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3ffb5747",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import scanpy as sc\n",
+    "\n",
+    "# normalise and log transform the data\n",
+    "sc.pp.normalize_total(adata_train, inplace=True)\n",
+    "sc.pp.log1p(adata_train)\n",
+    "\n",
+    "sc.pp.normalize_total(adata_val, inplace=True)\n",
+    "sc.pp.log1p(adata_val)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b699d36c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# since these datasets only contain 24 proteins, we will just use their expression as embeddings\n",
+    "# we will use the protein names as the embedding keys\n",
+    "adata_train.obsm[\"X_prot\"] = adata_train.X\n",
+    "adata_val.obsm[\"X_prot\"] = adata_val.X\n",
+    "processed_paths = {\"train\": \"Frangiehlzar2021_protein_pp_train.h5ad\", \"val\": \"Frangiehlzar2021_protein_pp_val.h5ad\"}\n",
+    "adata_train.write_h5ad(processed_paths[\"train\"])\n",
+    "adata_val.write_h5ad(processed_paths[\"val\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "09af0ef7",
+   "metadata": {},
+   "source": [
+    "### Step 3: Construct a contrastive training dataset\n",
+    "\n",
+    "So far we have:\n",
+    "- A **numeric representation** for each cell (stored in `adata.obsm[\"X_prot\"]`).\n",
+    "- A **text description** for each cell in `adata.obs[\"description\"]`.\n",
+    "\n",
+    "Next, we convert these into a Hugging Face `DatasetDict` suitable for contrastive training:\n",
+    "\n",
+    "- `AnnDataSetConstructor` in `\"multiplets\"` mode creates, for each cell, a\n",
+    "  - **query token** (`sample_idx` / later `cell_sentence_1`),\n",
+    "  - a **matched positive**, and\n",
+    "  - at least one **negative**.\n",
+    "- This gives us a flexible format where the same dataset can later support different training regimens (e.g. using text-only, numeric-only, or mixed representations), simply by choosing which columns to feed into the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2bd012cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from adata_hf_datasets import AnnDataSetConstructor\n",
+    "from datasets import DatasetDict\n",
+    "\n",
+    "ds_dict = DatasetDict()\n",
+    "# multiplets format for training datasets (with descriptions)\n",
+    "# A sentence key refers to the column in adata.obs that is used to represent the sample.\n",
+    "# For numeric data, we use the sample index, and later register the created embedding linked to their indices in the tokenizer\n",
+    "for split_name, adata_split in {\"train\": adata_train, \"val\": adata_val}.items():\n",
+    "    constructor = AnnDataSetConstructor(dataset_format=\"multiplets\", resolve_negatives=True)\n",
+    "    constructor.add_anndata(\n",
+    "        adata_split,\n",
+    "        caption_key=\"description\",\n",
+    "        sentence_keys=[\"sample_idx\"],\n",
+    "        adata_link=processed_paths[split_name],\n",
+    "        batch_key=\"library_preparation_protocol\",  # In this case all are from the same batch, but providing a batch key can improve batch integration by negative sampling\n",
+    "    )\n",
+    "    ds = constructor.get_dataset()\n",
+    "    ds_dict[split_name] = ds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e4d531a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_dict"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "70612269",
+   "metadata": {},
+   "source": [
+    "## Configure the Model\n",
+    "\n",
+    "The `MMContextEncoder` is the **main model** in this framework. Conceptually, it is a custom Sentence Transformers module that:\n",
+    "\n",
+    "- Starts from a standard **text encoder** (here: `sentence-transformers/all-MiniLM-L6-v2`).\n",
+    "- Adds **adapters** that map\n",
+    "  - the text encoder output into a desired embedding dimension, and\n",
+    "  - the numeric representation (e.g. `X_prot`, PCA, scVI, Geneformer embeddings) into the **same** space.\n",
+    "- Lets you configure whether the text encoder is **frozen**, partially unfrozen (last *n* layers), or fully trainable.\n",
+    "\n",
+    "To make the numeric data usable by the model, we need to **register initial embeddings**:\n",
+    "\n",
+    "- Tokens in `cell_sentence_1` (e.g. `\"sample_idx:123\"`) should map to a vector from our numeric representation.\n",
+    "- Internally, the model builds a lookup table (`sample_idx` → integer id) and an embedding matrix (id → numeric vector), which can be kept frozen or made trainable.\n",
+    "\n",
+    "After this setup, training looks like standard Sentence Transformers training, but every sample has **both** a text description and a numeric embedding living in a shared space."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6f04ed8b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mmcontext.mmcontextencoder import MMContextEncoder\n",
+    "from sentence_transformers import SentenceTransformer\n",
+    "\n",
+    "enc = MMContextEncoder(\n",
+    "    text_encoder_name=\"sentence-transformers/all-MiniLM-L6-v2\",\n",
+    "    adapter_hidden_dim=128,\n",
+    "    adapter_output_dim=64,\n",
+    "    freeze_text_encoder=True,\n",
+    "    unfreeze_last_n_layers=2,\n",
+    "    output_token_embeddings=False,\n",
+    "    train_lookup=False,\n",
+    "    joint_adapter_hidden_dim=None,\n",
+    "    text_model_kwargs=None,\n",
+    "    use_text_adapter=True,\n",
+    ")\n",
+    "model = SentenceTransformer(modules=[enc])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d0926c15",
+   "metadata": {},
+   "source": [
+    "### Register the initial numeric embeddings\n",
+    "\n",
+    "We now connect the numeric representations on disk to the model:\n",
+    "\n",
+    "- `get_initial_embeddings_from_adata_link` reads the preprocessed `.h5ad` files and builds a table with one row per `sample_idx` and its embedding from `X_prot`.\n",
+    "- `register_initial_embeddings` turns this table into an internal embedding matrix and lookup table inside the `MMContextEncoder`.\n",
+    "\n",
+    "After this step, whenever the model sees a token like `\"sample_idx:123\"` in the dataset, it can look up the correct numeric vector and combine it with the text pathway."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee9a6e96",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "token_df, _ = model[0].get_initial_embeddings_from_adata_link(\n",
+    "    ds_dict,\n",
+    "    layer_key=\"X_prot\",\n",
+    "    download_dir=\"data_cache\",\n",
+    "    axis=\"obs\",  # since we get embeddings from adata.obsm. We could also use \"varm\" and for example use an embedding for each protein\n",
+    ")\n",
+    "model[0].register_initial_embeddings(token_df, data_origin=\"prot\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6975ab72",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# the model expects a certain prefix on the cell tokens.\n",
+    "model[0].processor.prefix"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "52d74f63",
+   "metadata": {},
+   "source": [
+    "### Prepare token prefixes and final training columns\n",
+    "\n",
+    "At this point, the model knows about all numeric embeddings, but the dataset still uses the plain `sample_idx` strings.\n",
+    "\n",
+    "- `model[0].processor.prefix` defines the required prefix for these tokens (e.g. `\"sample_idx:\"`).\n",
+    "- `prefix_ds` rewrites the `cell_sentence_1` column so that each entry matches what the tokenizer and lookup table expect.\n",
+    "\n",
+    "Finally, we simplify the dataset to the three columns we need for contrastive learning:\n",
+    "\n",
+    "- `anchor`: the main cell token (formerly `cell_sentence_1`).\n",
+    "- `positive`: a biologically matched partner.\n",
+    "- `negative_1`: a contrastive negative.\n",
+    "\n",
+    "This might look a bit cumbersome, but it makes the dataset **reusable**: you can easily construct different views (text-only, numeric-only, different negatives) without having to rebuild everything from scratch."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6fc91911",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# you could add this manually or use the method below\n",
+    "model[0].prefix_ds(ds_dict, columns_to_prefix=[\"cell_sentence_1\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c09bac80",
+   "metadata": {},
+   "source": [
+    "## Training setup and objective\n",
+    "\n",
+    "With the dataset and model ready, we now configure the training loop using the standard Sentence Transformers API:\n",
+    "\n",
+    "- `SentenceTransformerTrainingArguments` controls the **optimisation hyperparameters** (epochs, batch size, learning rate, warmup, logging and checkpointing frequency).\n",
+    "- `SentenceTransformerTrainer` automatically picks an appropriate device (CPU, CUDA GPU, or Apple **MPS**) and handles mixed-precision flags (`fp16`/`bf16`) where supported.\n",
+    "- We use `MultipleNegativesRankingLoss`, which encourages anchors to be close to their positives and far from all other samples in the batch.\n",
+    "- `TripletEvaluator` periodically evaluates the model on the validation set using triplets `(anchor, positive, negative)` and reports an accuracy metric.\n",
+    "\n",
+    "Running `trainer.train()` then performs the actual optimisation and will log progress and evaluation metrics during training."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f63b0e4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# lastly we have to drop some columns from the dataset and\n",
+    "# rename the main column to \"anchor\".\n",
+    "# you might think that this is a bit cumbersome, which it is.\n",
+    "# But this setup allowed for fleixble training,\n",
+    "# using either cell or feature level tokens, using text based cell sentences\n",
+    "# or numeric embeddings and resolving negatives\n",
+    "# to whatever column was chosen for training.\n",
+    "# That means that for a certain training run, the same dataset can be reused, and\n",
+    "# only modified differently. But in the end, it is a bit of work to set up.\n",
+    "ds_final = ds_dict.rename_column(\"cell_sentence_1\", \"anchor\")\n",
+    "ds_final = ds_final.remove_columns([\"sample_idx\", \"adata_link\", \"negative_1_idx\"])\n",
+    "ds_final"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "25c0a4e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments\n",
+    "\n",
+    "args = SentenceTransformerTrainingArguments(\n",
+    "    num_train_epochs=1,\n",
+    "    per_device_train_batch_size=128,\n",
+    "    per_device_eval_batch_size=128,\n",
+    "    learning_rate=1e-5,\n",
+    "    warmup_ratio=0.1,\n",
+    "    fp16=False,\n",
+    "    bf16=False,\n",
+    "    eval_strategy=\"steps\",\n",
+    "    eval_steps=100,\n",
+    "    save_strategy=\"steps\",\n",
+    "    save_steps=100,\n",
+    "    save_total_limit=1,\n",
+    "    max_grad_norm=1.0,\n",
+    "    logging_steps=10,\n",
+    "    run_name=\"protein_test\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d44a49b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sentence_transformers.evaluation import TripletEvaluator\n",
+    "from sentence_transformers.losses import MultipleNegativesRankingLoss\n",
+    "\n",
+    "loss = MultipleNegativesRankingLoss(model)\n",
+    "evaluator = TripletEvaluator(\n",
+    "    anchors=ds_final[\"val\"][\"anchor\"],\n",
+    "    positives=ds_final[\"val\"][\"positive\"],\n",
+    "    negatives=ds_final[\"val\"][\"negative_1\"],\n",
+    "    name=\"val\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53d32063",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = SentenceTransformerTrainer(\n",
+    "    model=model,\n",
+    "    args=args,\n",
+    "    train_dataset=ds_final[\"train\"],\n",
+    "    eval_dataset=ds_final[\"val\"],\n",
+    "    loss=loss,\n",
+    "    evaluator=evaluator,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84f83d70",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer.train()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
 }
diff --git a/uv.lock b/uv.lock
index b11beaf..0397703 100644
--- a/uv.lock
+++ b/uv.lock
@@ -28,38 +28,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/9f/1c/a17fb513aeb684fb83bef5f395910f53103ab30308bbdd77fd66d6698c46/accelerate-1.9.0-py3-none-any.whl", hash = "sha256:c24739a97ade1d54af4549a65f8b6b046adc87e2b3e4d6c66516e32c53d5a8f1", size = 367073, upload-time = "2025-07-16T16:24:52.957Z" },
 ]
 
-[[package]]
-name = "adata-hf-datasets"
-version = "0.1.0"
-source = { git = "https://github.com/mengerj/adata_hf_datasets.git#ca4a4a95195b2461a9231a88a98562661777aa12" }
-dependencies = [
-    { name = "anndata" },
-    { name = "appdirs" },
-    { name = "boto3" },
-    { name = "botocore" },
-    { name = "datasets", version = "3.6.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" },
-    { name = "datasets", version = "4.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "huggingface-hub" },
-    { name = "hydra-core" },
-    { name = "ipykernel" },
-    { name = "ipywidgets" },
-    { name = "jupyter" },
-    { name = "mygene" },
-    { name = "numpy" },
-    { name = "pandas" },
-    { name = "pyarrow", version = "17.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" },
-    { name = "pyarrow", version = "22.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "pybiomart" },
-    { name = "pynvml" },
-    { name = "pysradb" },
-    { name = "python-dotenv" },
-    { name = "requests" },
-    { name = "scanpy" },
-    { name = "torch" },
-    { name = "transformers" },
-    { name = "zarr" },
-]
-
 [[package]]
 name = "aiohappyeyeballs"
 version = "2.6.1"
@@ -143,7 +111,8 @@ version = "1.4.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "frozenlist" },
-    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.12.*'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" }
 wheels = [
@@ -152,20 +121,25 @@ wheels = [
 
 [[package]]
 name = "anndata"
-version = "0.11.3"
+version = "0.12.17"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "array-api-compat" },
     { name = "h5py" },
+    { name = "legacy-api-wrap" },
     { name = "natsort" },
     { name = "numpy" },
     { name = "packaging" },
     { name = "pandas" },
     { name = "scipy" },
+    { name = "scverse-misc", version = "0.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "scverse-misc", version = "0.0.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+    { name = "zarr", version = "3.1.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "zarr", version = "3.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/fd/1b/917c82b44cffc3157e74618fda92a61e9353c3d45a21bbffce90df337bef/anndata-0.11.3.tar.gz", hash = "sha256:986e626b62108e8db7034db2f391fd160c11d4b4e3af9c31183c82862297f7f6", size = 538016, upload-time = "2025-01-10T20:17:54.059Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/04/e4/8cdf4d18c422e6c30edf512d599f9781fb458939d6317cc1e1e7940f2157/anndata-0.12.17.tar.gz", hash = "sha256:86608c28dc78d8a30c03fb3566caf6fa73d9c5aa7be53e4ed6b20888cc5e4337", size = 2257555, upload-time = "2026-06-16T12:30:15.608Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ca/ee/2c4ba5ddd3d5664c02e8b4faefc94c7689601a29af4735b1491eb44e12a6/anndata-0.11.3-py3-none-any.whl", hash = "sha256:2ce220c4c03b7dba31fab0540d3d64f6811bc343cd3f273d7bdeed50a5cd6367", size = 142665, upload-time = "2025-01-10T20:17:51.992Z" },
+    { url = "https://files.pythonhosted.org/packages/2f/f6/bca233b8311dc2217d9d7c1687acdfa347f778f6f6eb622d8e5afa39f014/anndata-0.12.17-py3-none-any.whl", hash = "sha256:a19e48741a935752d11c9d2dc968ac0cdf3900e0c8c3413320fde84ab11dc9e6", size = 175991, upload-time = "2026-06-16T12:30:13.686Z" },
 ]
 
 [[package]]
@@ -177,77 +151,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" },
 ]
 
-[[package]]
-name = "antlr4-python3-runtime"
-version = "4.9.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" }
-
 [[package]]
 name = "anyio"
 version = "4.11.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "idna" },
-    { name = "sniffio" },
-    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
+    { name = "idna", marker = "python_full_version >= '3.13'" },
+    { name = "sniffio", marker = "python_full_version >= '3.13'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/c6/78/7d432127c41b50bccba979505f272c16cbcadcc33645d5fa3a738110ae75/anyio-4.11.0.tar.gz", hash = "sha256:82a8d0b81e318cc5ce71a5f1f8b5c4e63619620b63141ef8c995fa0db95a57c4", size = 219094, upload-time = "2025-09-23T09:19:12.58Z" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/15/b3/9b1a8074496371342ec1e796a96f99c82c945a339cd81a8e73de28b4cf9e/anyio-4.11.0-py3-none-any.whl", hash = "sha256:0287e96f4d26d4149305414d4e3bc32f0dcd0862365a4bddea19d7a1ec38c4fc", size = 109097, upload-time = "2025-09-23T09:19:10.601Z" },
 ]
 
-[[package]]
-name = "appdirs"
-version = "1.4.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d7/d8/05696357e0311f5b5c316d7b95f46c669dd9c15aaeecbb48c7d0aeb88c40/appdirs-1.4.4.tar.gz", hash = "sha256:7d5d0167b2b1ba821647616af46a749d1c653740dd0d2415100fe26e27afdf41", size = 13470, upload-time = "2020-05-11T07:59:51.037Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/3b/00/2344469e2084fb287c2e0b57b72910309874c3245463acd6cf5e3db69324/appdirs-1.4.4-py2.py3-none-any.whl", hash = "sha256:a841dacd6b99318a741b166adb07e19ee71a274450e68237b4650ca1055ab128", size = 9566, upload-time = "2020-05-11T07:59:49.499Z" },
-]
-
-[[package]]
-name = "appnope"
-version = "0.1.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/35/5d/752690df9ef5b76e169e68d6a129fa6d08a7100ca7f754c89495db3c6019/appnope-0.1.4.tar.gz", hash = "sha256:1de3860566df9caf38f01f86f65e0e13e379af54f9e4bee1e66b48f2efffd1ee", size = 4170, upload-time = "2024-02-06T09:43:11.258Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" },
-]
-
-[[package]]
-name = "argon2-cffi"
-version = "25.1.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "argon2-cffi-bindings" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/0e/89/ce5af8a7d472a67cc819d5d998aa8c82c5d860608c4db9f46f1162d7dab9/argon2_cffi-25.1.0.tar.gz", hash = "sha256:694ae5cc8a42f4c4e2bf2ca0e64e51e23a040c6a517a85074683d3959e1346c1", size = 45706, upload-time = "2025-06-03T06:55:32.073Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/4f/d3/a8b22fa575b297cd6e3e3b0155c7e25db170edf1c74783d6a31a2490b8d9/argon2_cffi-25.1.0-py3-none-any.whl", hash = "sha256:fdc8b074db390fccb6eb4a3604ae7231f219aa669a2652e0f20e16ba513d5741", size = 14657, upload-time = "2025-06-03T06:55:30.804Z" },
-]
-
-[[package]]
-name = "argon2-cffi-bindings"
-version = "25.1.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "cffi" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/5c/2d/db8af0df73c1cf454f71b2bbe5e356b8c1f8041c979f505b3d3186e520a9/argon2_cffi_bindings-25.1.0.tar.gz", hash = "sha256:b957f3e6ea4d55d820e40ff76f450952807013d361a65d7f28acc0acbf29229d", size = 1783441, upload-time = "2025-07-30T10:02:05.147Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1d/57/96b8b9f93166147826da5f90376e784a10582dd39a393c99bb62cfcf52f0/argon2_cffi_bindings-25.1.0-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:aecba1723ae35330a008418a91ea6cfcedf6d31e5fbaa056a166462ff066d500", size = 54121, upload-time = "2025-07-30T10:01:50.815Z" },
-    { url = "https://files.pythonhosted.org/packages/0a/08/a9bebdb2e0e602dde230bdde8021b29f71f7841bd54801bcfd514acb5dcf/argon2_cffi_bindings-25.1.0-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:2630b6240b495dfab90aebe159ff784d08ea999aa4b0d17efa734055a07d2f44", size = 29177, upload-time = "2025-07-30T10:01:51.681Z" },
-    { url = "https://files.pythonhosted.org/packages/b6/02/d297943bcacf05e4f2a94ab6f462831dc20158614e5d067c35d4e63b9acb/argon2_cffi_bindings-25.1.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:7aef0c91e2c0fbca6fc68e7555aa60ef7008a739cbe045541e438373bc54d2b0", size = 31090, upload-time = "2025-07-30T10:01:53.184Z" },
-    { url = "https://files.pythonhosted.org/packages/c1/93/44365f3d75053e53893ec6d733e4a5e3147502663554b4d864587c7828a7/argon2_cffi_bindings-25.1.0-cp39-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1e021e87faa76ae0d413b619fe2b65ab9a037f24c60a1e6cc43457ae20de6dc6", size = 81246, upload-time = "2025-07-30T10:01:54.145Z" },
-    { url = "https://files.pythonhosted.org/packages/09/52/94108adfdd6e2ddf58be64f959a0b9c7d4ef2fa71086c38356d22dc501ea/argon2_cffi_bindings-25.1.0-cp39-abi3-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d3e924cfc503018a714f94a49a149fdc0b644eaead5d1f089330399134fa028a", size = 87126, upload-time = "2025-07-30T10:01:55.074Z" },
-    { url = "https://files.pythonhosted.org/packages/72/70/7a2993a12b0ffa2a9271259b79cc616e2389ed1a4d93842fac5a1f923ffd/argon2_cffi_bindings-25.1.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c87b72589133f0346a1cb8d5ecca4b933e3c9b64656c9d175270a000e73b288d", size = 80343, upload-time = "2025-07-30T10:01:56.007Z" },
-    { url = "https://files.pythonhosted.org/packages/78/9a/4e5157d893ffc712b74dbd868c7f62365618266982b64accab26bab01edc/argon2_cffi_bindings-25.1.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:1db89609c06afa1a214a69a462ea741cf735b29a57530478c06eb81dd403de99", size = 86777, upload-time = "2025-07-30T10:01:56.943Z" },
-    { url = "https://files.pythonhosted.org/packages/74/cd/15777dfde1c29d96de7f18edf4cc94c385646852e7c7b0320aa91ccca583/argon2_cffi_bindings-25.1.0-cp39-abi3-win32.whl", hash = "sha256:473bcb5f82924b1becbb637b63303ec8d10e84c8d241119419897a26116515d2", size = 27180, upload-time = "2025-07-30T10:01:57.759Z" },
-    { url = "https://files.pythonhosted.org/packages/e2/c6/a759ece8f1829d1f162261226fbfd2c6832b3ff7657384045286d2afa384/argon2_cffi_bindings-25.1.0-cp39-abi3-win_amd64.whl", hash = "sha256:a98cd7d17e9f7ce244c0803cad3c23a7d379c301ba618a5fa76a67d116618b98", size = 31715, upload-time = "2025-07-30T10:01:58.56Z" },
-    { url = "https://files.pythonhosted.org/packages/42/b9/f8d6fa329ab25128b7e98fd83a3cb34d9db5b059a9847eddb840a0af45dd/argon2_cffi_bindings-25.1.0-cp39-abi3-win_arm64.whl", hash = "sha256:b0fdbcf513833809c882823f98dc2f931cf659d9a1429616ac3adebb49f5db94", size = 27149, upload-time = "2025-07-30T10:01:59.329Z" },
-]
-
 [[package]]
 name = "array-api-compat"
 version = "1.12.0"
@@ -257,43 +173,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e0/b1/0542e0cab6f49f151a2d7a42400f84f706fc0b64e85dc1f56708b2e9fd37/array_api_compat-1.12.0-py3-none-any.whl", hash = "sha256:a0b4795b6944a9507fde54679f9350e2ad2b1e2acf4a2408a098cdc27f890a8b", size = 58156, upload-time = "2025-05-16T08:49:58.129Z" },
 ]
 
-[[package]]
-name = "arrow"
-version = "1.4.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "python-dateutil" },
-    { name = "tzdata" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/b9/33/032cdc44182491aa708d06a68b62434140d8c50820a087fac7af37703357/arrow-1.4.0.tar.gz", hash = "sha256:ed0cc050e98001b8779e84d461b0098c4ac597e88704a655582b21d116e526d7", size = 152931, upload-time = "2025-10-18T17:46:46.761Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ed/c9/d7977eaacb9df673210491da99e6a247e93df98c715fc43fd136ce1d3d33/arrow-1.4.0-py3-none-any.whl", hash = "sha256:749f0769958ebdc79c173ff0b0670d59051a535fa26e8eba02953dc19eb43205", size = 68797, upload-time = "2025-10-18T17:46:45.663Z" },
-]
-
-[[package]]
-name = "asciitree"
-version = "0.3.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/2d/6a/885bc91484e1aa8f618f6f0228d76d0e67000b0fdd6090673b777e311913/asciitree-0.3.3.tar.gz", hash = "sha256:4aa4b9b649f85e3fcb343363d97564aa1fb62e249677f2e18a96765145cc0f6e", size = 3951, upload-time = "2016-09-05T19:10:42.681Z" }
-
-[[package]]
-name = "asttokens"
-version = "3.0.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/4a/e7/82da0a03e7ba5141f05cce0d302e6eed121ae055e0456ca228bf693984bc/asttokens-3.0.0.tar.gz", hash = "sha256:0dcd8baa8d62b0c1d118b399b2ddba3c4aff271d0d7a9e0d4c1681c79035bbc7", size = 61978, upload-time = "2024-11-30T04:30:14.439Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/25/8a/c46dcc25341b5bce5472c718902eb3d38600a903b14fa6aeecef3f21a46f/asttokens-3.0.0-py3-none-any.whl", hash = "sha256:e3078351a059199dd5138cb1c706e6430c05eff2ff136af5eb4790f9d28932e2", size = 26918, upload-time = "2024-11-30T04:30:10.946Z" },
-]
-
-[[package]]
-name = "async-lru"
-version = "2.0.5"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/b2/4d/71ec4d3939dc755264f680f6c2b4906423a304c3d18e96853f0a595dfe97/async_lru-2.0.5.tar.gz", hash = "sha256:481d52ccdd27275f42c43a928b4a50c3bfb2d67af4e78b170e3e0bb39c66e5bb", size = 10380, upload-time = "2025-03-16T17:25:36.919Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/03/49/d10027df9fce941cb8184e78a02857af36360d33e1721df81c5ed2179a1a/async_lru-2.0.5-py3-none-any.whl", hash = "sha256:ab95404d8d2605310d345932697371a5f40def0487c03d6d0ad9138de52c9943", size = 6069, upload-time = "2025-03-16T17:25:35.422Z" },
-]
-
 [[package]]
 name = "attrs"
 version = "25.3.0"
@@ -303,15 +182,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/77/06/bb80f5f86020c4551da315d78b3ab75e8228f89f0162f2c3a819e407941a/attrs-25.3.0-py3-none-any.whl", hash = "sha256:427318ce031701fea540783410126f03899a97ffc6f61596ad581ac2e40e3bc3", size = 63815, upload-time = "2025-03-13T11:10:21.14Z" },
 ]
 
-[[package]]
-name = "babel"
-version = "2.17.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/7d/6b/d52e42361e1aa00709585ecc30b3f9684b3ab62530771402248b1b1d6240/babel-2.17.0.tar.gz", hash = "sha256:0c54cffb19f690cdcc52a3b50bcbf71e07a808d1c80d549f2459b9d2cf0afb9d", size = 9951852, upload-time = "2025-02-01T15:17:41.026Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b7/b8/3fe70c75fe32afc4bb507f75563d39bc5642255d1d94f1f23604725780bf/babel-2.17.0-py3-none-any.whl", hash = "sha256:4d0b53093fdfb4b21c92b5213dba5a1b23885afa8383709427046b21c366e5f2", size = 10182537, upload-time = "2025-02-01T15:17:37.39Z" },
-]
-
 [[package]]
 name = "backports-tarfile"
 version = "1.2.0"
@@ -322,86 +192,17 @@ wheels = [
 ]
 
 [[package]]
-name = "beautifulsoup4"
-version = "4.13.4"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "soupsieve" },
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d8/e4/0c4c39e18fd76d6a628d4dd8da40543d136ce2d1752bd6eeeab0791f4d6b/beautifulsoup4-4.13.4.tar.gz", hash = "sha256:dbb3c4e1ceae6aefebdaf2423247260cd062430a410e38c66f2baa50a8437195", size = 621067, upload-time = "2025-04-15T17:05:13.836Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/50/cd/30110dc0ffcf3b131156077b90e9f60ed75711223f306da4db08eff8403b/beautifulsoup4-4.13.4-py3-none-any.whl", hash = "sha256:9bbbb14bfde9d79f38b8cd5f8c7c85f4b8f2523190ebed90e950a8dea4cb1c4b", size = 187285, upload-time = "2025-04-15T17:05:12.221Z" },
-]
-
-[[package]]
-name = "biothings-client"
-version = "0.4.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "httpx" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/ef/2f/2ef9115e317f4c2acb690b341f40de751b16d14169fdbb8d6eb86964166c/biothings_client-0.4.1.tar.gz", hash = "sha256:5b34e09c905280b5bd2538f1f34b6fc780c53c8da9b4074e3ff304836046f613", size = 56096, upload-time = "2025-01-14T22:55:39.926Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/dc/6d/130477cfbd7294949b919c45cc1ed14a642cec95afba06a54400a4419235/biothings_client-0.4.1-py3-none-any.whl", hash = "sha256:9cbc17461b2bf6af6ed200929b886d6670d450af2034b428cd833f725695265a", size = 46698, upload-time = "2025-01-14T22:55:37.44Z" },
-]
-
-[[package]]
-name = "bleach"
-version = "6.3.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "webencodings" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/07/18/3c8523962314be6bf4c8989c79ad9531c825210dd13a8669f6b84336e8bd/bleach-6.3.0.tar.gz", hash = "sha256:6f3b91b1c0a02bb9a78b5a454c92506aa0fdf197e1d5e114d2e00c6f64306d22", size = 203533, upload-time = "2025-10-27T17:57:39.211Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/cd/3a/577b549de0cc09d95f11087ee63c739bba856cd3952697eec4c4bb91350a/bleach-6.3.0-py3-none-any.whl", hash = "sha256:fe10ec77c93ddf3d13a73b035abaac7a9f5e436513864ccdad516693213c65d6", size = 164437, upload-time = "2025-10-27T17:57:37.538Z" },
-]
-
-[package.optional-dependencies]
-css = [
-    { name = "tinycss2" },
-]
-
-[[package]]
-name = "boto3"
-version = "1.40.70"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "botocore" },
-    { name = "jmespath" },
-    { name = "s3transfer" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/37/12/d5ac34e0536e1914dde28245f014a635056dde0427f6efa09f104d7999f4/boto3-1.40.70.tar.gz", hash = "sha256:191443707b391232ed15676bf6bba7e53caec1e71aafa12ccad2e825c5ee15cc", size = 111638, upload-time = "2025-11-10T20:29:15.199Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f3/cf/e24d08b37cd318754a8e94906c8b34b88676899aad1907ff6942311f13c4/boto3-1.40.70-py3-none-any.whl", hash = "sha256:e8c2f4f4cb36297270f1023ebe5b100333e0e88ab6457a9687d80143d2e15bf9", size = 139358, upload-time = "2025-11-10T20:29:13.512Z" },
-]
-
-[[package]]
-name = "botocore"
-version = "1.40.70"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jmespath" },
-    { name = "python-dateutil" },
-    { name = "urllib3" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/35/c1/8c4c199ae1663feee579a15861e34f10b29da11ae6ea0ad7b6a847ef3823/botocore-1.40.70.tar.gz", hash = "sha256:61b1f2cecd54d1b28a081116fa113b97bf4e17da57c62ae2c2751fe4c528af1f", size = 14444592, upload-time = "2025-11-10T20:29:04.046Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/55/d2/507fd0ee4dd574d2bdbdeac5df83f39d2cae1ffe97d4622cca6f6bab39f1/botocore-1.40.70-py3-none-any.whl", hash = "sha256:4a394ad25f5d9f1ef0bed610365744523eeb5c22de6862ab25d8c93f9f6d295c", size = 14106829, upload-time = "2025-11-10T20:29:01.101Z" },
-]
-
-[[package]]
-name = "cattrs"
-version = "25.2.0"
+name = "build"
+version = "1.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "attrs" },
-    { name = "typing-extensions" },
+    { name = "colorama", marker = "os_name == 'nt' and sys_platform != 'linux'" },
+    { name = "packaging" },
+    { name = "pyproject-hooks" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/e3/42/988b3a667967e9d2d32346e7ed7edee540ef1cee829b53ef80aa8d4a0222/cattrs-25.2.0.tar.gz", hash = "sha256:f46c918e955db0177be6aa559068390f71988e877c603ae2e56c71827165cc06", size = 506531, upload-time = "2025-08-31T20:41:59.301Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/78/e0/df5e171f685f82f37b12e1f208064e24244911079d7b767447d1af7e0d70/build-1.5.0.tar.gz", hash = "sha256:302c22c3ba2a0fd5f3911918651341ebb3896176cbdec15bd421f80b1afc7647", size = 89796, upload-time = "2026-04-30T03:18:25.17Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/20/a5/b3771ac30b590026b9d721187110194ade05bfbea3d98b423a9cafd80959/cattrs-25.2.0-py3-none-any.whl", hash = "sha256:539d7eedee7d2f0706e4e109182ad096d608ba84633c32c75ef3458f1d11e8f1", size = 70040, upload-time = "2025-08-31T20:41:57.543Z" },
+    { url = "https://files.pythonhosted.org/packages/0d/fe/6bea5c9162869c5beba5d9c8abbed835ec85bf1ec1fba05a3822325c45f3/build-1.5.0-py3-none-any.whl", hash = "sha256:13f3eecb844759ab66efec90ca17639bbf14dc06cb2fdf37a9010322d9c50a6f", size = 26018, upload-time = "2026-04-30T03:18:23.644Z" },
 ]
 
 [[package]]
@@ -418,44 +219,26 @@ name = "cffi"
 version = "1.17.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "pycparser" },
+    { name = "pycparser", marker = "sys_platform == 'linux'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/fc/97/c783634659c2920c3fc70419e3af40972dbaf758daa229a7d6ea6135c90d/cffi-1.17.1.tar.gz", hash = "sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824", size = 516621, upload-time = "2024-09-04T20:45:21.852Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/6b/f4/927e3a8899e52a27fa57a48607ff7dc91a9ebe97399b357b85a0c7892e00/cffi-1.17.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a45e3c6913c5b87b3ff120dcdc03f6131fa0065027d0ed7ee6190736a74cd401", size = 182264, upload-time = "2024-09-04T20:43:51.124Z" },
-    { url = "https://files.pythonhosted.org/packages/6c/f5/6c3a8efe5f503175aaddcbea6ad0d2c96dad6f5abb205750d1b3df44ef29/cffi-1.17.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:30c5e0cb5ae493c04c8b42916e52ca38079f1b235c2f8ae5f4527b963c401caf", size = 178651, upload-time = "2024-09-04T20:43:52.872Z" },
     { url = "https://files.pythonhosted.org/packages/94/dd/a3f0118e688d1b1a57553da23b16bdade96d2f9bcda4d32e7d2838047ff7/cffi-1.17.1-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f75c7ab1f9e4aca5414ed4d8e5c0e303a34f4421f8a0d47a4d019ceff0ab6af4", size = 445259, upload-time = "2024-09-04T20:43:56.123Z" },
     { url = "https://files.pythonhosted.org/packages/2e/ea/70ce63780f096e16ce8588efe039d3c4f91deb1dc01e9c73a287939c79a6/cffi-1.17.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a1ed2dd2972641495a3ec98445e09766f077aee98a1c896dcb4ad0d303628e41", size = 469200, upload-time = "2024-09-04T20:43:57.891Z" },
-    { url = "https://files.pythonhosted.org/packages/1c/a0/a4fa9f4f781bda074c3ddd57a572b060fa0df7655d2a4247bbe277200146/cffi-1.17.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:46bf43160c1a35f7ec506d254e5c890f3c03648a4dbac12d624e4490a7046cd1", size = 477235, upload-time = "2024-09-04T20:44:00.18Z" },
-    { url = "https://files.pythonhosted.org/packages/62/12/ce8710b5b8affbcdd5c6e367217c242524ad17a02fe5beec3ee339f69f85/cffi-1.17.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a24ed04c8ffd54b0729c07cee15a81d964e6fee0e3d4d342a27b020d22959dc6", size = 459721, upload-time = "2024-09-04T20:44:01.585Z" },
     { url = "https://files.pythonhosted.org/packages/ff/6b/d45873c5e0242196f042d555526f92aa9e0c32355a1be1ff8c27f077fd37/cffi-1.17.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:610faea79c43e44c71e1ec53a554553fa22321b65fae24889706c0a84d4ad86d", size = 467242, upload-time = "2024-09-04T20:44:03.467Z" },
     { url = "https://files.pythonhosted.org/packages/1a/52/d9a0e523a572fbccf2955f5abe883cfa8bcc570d7faeee06336fbd50c9fc/cffi-1.17.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:a9b15d491f3ad5d692e11f6b71f7857e7835eb677955c00cc0aefcd0669adaf6", size = 477999, upload-time = "2024-09-04T20:44:05.023Z" },
     { url = "https://files.pythonhosted.org/packages/44/74/f2a2460684a1a2d00ca799ad880d54652841a780c4c97b87754f660c7603/cffi-1.17.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:de2ea4b5833625383e464549fec1bc395c1bdeeb5f25c4a3a82b5a8c756ec22f", size = 454242, upload-time = "2024-09-04T20:44:06.444Z" },
     { url = "https://files.pythonhosted.org/packages/f8/4a/34599cac7dfcd888ff54e801afe06a19c17787dfd94495ab0c8d35fe99fb/cffi-1.17.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:fc48c783f9c87e60831201f2cce7f3b2e4846bf4d8728eabe54d60700b318a0b", size = 478604, upload-time = "2024-09-04T20:44:08.206Z" },
-    { url = "https://files.pythonhosted.org/packages/34/33/e1b8a1ba29025adbdcda5fb3a36f94c03d771c1b7b12f726ff7fef2ebe36/cffi-1.17.1-cp311-cp311-win32.whl", hash = "sha256:85a950a4ac9c359340d5963966e3e0a94a676bd6245a4b55bc43949eee26a655", size = 171727, upload-time = "2024-09-04T20:44:09.481Z" },
-    { url = "https://files.pythonhosted.org/packages/3d/97/50228be003bb2802627d28ec0627837ac0bf35c90cf769812056f235b2d1/cffi-1.17.1-cp311-cp311-win_amd64.whl", hash = "sha256:caaf0640ef5f5517f49bc275eca1406b0ffa6aa184892812030f04c2abf589a0", size = 181400, upload-time = "2024-09-04T20:44:10.873Z" },
-    { url = "https://files.pythonhosted.org/packages/5a/84/e94227139ee5fb4d600a7a4927f322e1d4aea6fdc50bd3fca8493caba23f/cffi-1.17.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:805b4371bf7197c329fcb3ead37e710d1bca9da5d583f5073b799d5c5bd1eee4", size = 183178, upload-time = "2024-09-04T20:44:12.232Z" },
-    { url = "https://files.pythonhosted.org/packages/da/ee/fb72c2b48656111c4ef27f0f91da355e130a923473bf5ee75c5643d00cca/cffi-1.17.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:733e99bc2df47476e3848417c5a4540522f234dfd4ef3ab7fafdf555b082ec0c", size = 178840, upload-time = "2024-09-04T20:44:13.739Z" },
     { url = "https://files.pythonhosted.org/packages/cc/b6/db007700f67d151abadf508cbfd6a1884f57eab90b1bb985c4c8c02b0f28/cffi-1.17.1-cp312-cp312-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1257bdabf294dceb59f5e70c64a3e2f462c30c7ad68092d01bbbfb1c16b1ba36", size = 454803, upload-time = "2024-09-04T20:44:15.231Z" },
     { url = "https://files.pythonhosted.org/packages/1a/df/f8d151540d8c200eb1c6fba8cd0dfd40904f1b0682ea705c36e6c2e97ab3/cffi-1.17.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da95af8214998d77a98cc14e3a3bd00aa191526343078b530ceb0bd710fb48a5", size = 478850, upload-time = "2024-09-04T20:44:17.188Z" },
-    { url = "https://files.pythonhosted.org/packages/28/c0/b31116332a547fd2677ae5b78a2ef662dfc8023d67f41b2a83f7c2aa78b1/cffi-1.17.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d63afe322132c194cf832bfec0dc69a99fb9bb6bbd550f161a49e9e855cc78ff", size = 485729, upload-time = "2024-09-04T20:44:18.688Z" },
-    { url = "https://files.pythonhosted.org/packages/91/2b/9a1ddfa5c7f13cab007a2c9cc295b70fbbda7cb10a286aa6810338e60ea1/cffi-1.17.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f79fc4fc25f1c8698ff97788206bb3c2598949bfe0fef03d299eb1b5356ada99", size = 471256, upload-time = "2024-09-04T20:44:20.248Z" },
     { url = "https://files.pythonhosted.org/packages/b2/d5/da47df7004cb17e4955df6a43d14b3b4ae77737dff8bf7f8f333196717bf/cffi-1.17.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b62ce867176a75d03a665bad002af8e6d54644fad99a3c70905c543130e39d93", size = 479424, upload-time = "2024-09-04T20:44:21.673Z" },
     { url = "https://files.pythonhosted.org/packages/0b/ac/2a28bcf513e93a219c8a4e8e125534f4f6db03e3179ba1c45e949b76212c/cffi-1.17.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:386c8bf53c502fff58903061338ce4f4950cbdcb23e2902d86c0f722b786bbe3", size = 484568, upload-time = "2024-09-04T20:44:23.245Z" },
     { url = "https://files.pythonhosted.org/packages/d4/38/ca8a4f639065f14ae0f1d9751e70447a261f1a30fa7547a828ae08142465/cffi-1.17.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4ceb10419a9adf4460ea14cfd6bc43d08701f0835e979bf821052f1805850fe8", size = 488736, upload-time = "2024-09-04T20:44:24.757Z" },
-    { url = "https://files.pythonhosted.org/packages/86/c5/28b2d6f799ec0bdecf44dced2ec5ed43e0eb63097b0f58c293583b406582/cffi-1.17.1-cp312-cp312-win32.whl", hash = "sha256:a08d7e755f8ed21095a310a693525137cfe756ce62d066e53f502a83dc550f65", size = 172448, upload-time = "2024-09-04T20:44:26.208Z" },
-    { url = "https://files.pythonhosted.org/packages/50/b9/db34c4755a7bd1cb2d1603ac3863f22bcecbd1ba29e5ee841a4bc510b294/cffi-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:51392eae71afec0d0c8fb1a53b204dbb3bcabcb3c9b807eedf3e1e6ccf2de903", size = 181976, upload-time = "2024-09-04T20:44:27.578Z" },
-    { url = "https://files.pythonhosted.org/packages/8d/f8/dd6c246b148639254dad4d6803eb6a54e8c85c6e11ec9df2cffa87571dbe/cffi-1.17.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f3a2b4222ce6b60e2e8b337bb9596923045681d71e5a082783484d845390938e", size = 182989, upload-time = "2024-09-04T20:44:28.956Z" },
-    { url = "https://files.pythonhosted.org/packages/8b/f1/672d303ddf17c24fc83afd712316fda78dc6fce1cd53011b839483e1ecc8/cffi-1.17.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0984a4925a435b1da406122d4d7968dd861c1385afe3b45ba82b750f229811e2", size = 178802, upload-time = "2024-09-04T20:44:30.289Z" },
     { url = "https://files.pythonhosted.org/packages/0e/2d/eab2e858a91fdff70533cab61dcff4a1f55ec60425832ddfdc9cd36bc8af/cffi-1.17.1-cp313-cp313-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d01b12eeeb4427d3110de311e1774046ad344f5b1a7403101878976ecd7a10f3", size = 454792, upload-time = "2024-09-04T20:44:32.01Z" },
     { url = "https://files.pythonhosted.org/packages/75/b2/fbaec7c4455c604e29388d55599b99ebcc250a60050610fadde58932b7ee/cffi-1.17.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:706510fe141c86a69c8ddc029c7910003a17353970cff3b904ff0686a5927683", size = 478893, upload-time = "2024-09-04T20:44:33.606Z" },
-    { url = "https://files.pythonhosted.org/packages/4f/b7/6e4a2162178bf1935c336d4da8a9352cccab4d3a5d7914065490f08c0690/cffi-1.17.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:de55b766c7aa2e2a3092c51e0483d700341182f08e67c63630d5b6f200bb28e5", size = 485810, upload-time = "2024-09-04T20:44:35.191Z" },
-    { url = "https://files.pythonhosted.org/packages/c7/8a/1d0e4a9c26e54746dc08c2c6c037889124d4f59dffd853a659fa545f1b40/cffi-1.17.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c59d6e989d07460165cc5ad3c61f9fd8f1b4796eacbd81cee78957842b834af4", size = 471200, upload-time = "2024-09-04T20:44:36.743Z" },
     { url = "https://files.pythonhosted.org/packages/26/9f/1aab65a6c0db35f43c4d1b4f580e8df53914310afc10ae0397d29d697af4/cffi-1.17.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd398dbc6773384a17fe0d3e7eeb8d1a21c2200473ee6806bb5e6a8e62bb73dd", size = 479447, upload-time = "2024-09-04T20:44:38.492Z" },
     { url = "https://files.pythonhosted.org/packages/5f/e4/fb8b3dd8dc0e98edf1135ff067ae070bb32ef9d509d6cb0f538cd6f7483f/cffi-1.17.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3edc8d958eb099c634dace3c7e16560ae474aa3803a5df240542b305d14e14ed", size = 484358, upload-time = "2024-09-04T20:44:40.046Z" },
     { url = "https://files.pythonhosted.org/packages/f1/47/d7145bf2dc04684935d57d67dff9d6d795b2ba2796806bb109864be3a151/cffi-1.17.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:72e72408cad3d5419375fc87d289076ee319835bdfa2caad331e377589aebba9", size = 488469, upload-time = "2024-09-04T20:44:41.616Z" },
-    { url = "https://files.pythonhosted.org/packages/bf/ee/f94057fa6426481d663b88637a9a10e859e492c73d0384514a17d78ee205/cffi-1.17.1-cp313-cp313-win32.whl", hash = "sha256:e03eab0a8677fa80d646b5ddece1cbeaf556c313dcfac435ba11f107ba117b5d", size = 172475, upload-time = "2024-09-04T20:44:43.733Z" },
-    { url = "https://files.pythonhosted.org/packages/7c/fc/6a8cb64e5f0324877d503c854da15d76c1e50eb722e320b15345c4d0c6de/cffi-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:f6a16c31041f09ead72d69f583767292f750d24913dadacf5756b966aacb3f1a", size = 182009, upload-time = "2024-09-04T20:44:45.309Z" },
 ]
 
 [[package]]
@@ -497,15 +280,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" },
 ]
 
-[[package]]
-name = "comm"
-version = "0.2.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/4c/13/7d740c5849255756bc17888787313b61fd38a0a8304fc4f073dfc46122aa/comm-0.2.3.tar.gz", hash = "sha256:2dc8048c10962d55d7ad693be1e7045d891b7ce8d999c97963a5e3e99c055971", size = 6319, upload-time = "2025-07-25T14:02:04.452Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/60/97/891a0971e1e4a8c5d2b20bbe0e524dc04548d2307fee33cdeba148fd4fc7/comm-0.2.3-py3-none-any.whl", hash = "sha256:c615d91d75f7f04f095b30d1c1711babd43bdc6419c1be9886a85f2f4e489417", size = 7294, upload-time = "2025-07-25T14:02:02.896Z" },
-]
-
 [[package]]
 name = "contourpy"
 version = "1.3.3"
@@ -720,45 +494,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/3b/5e/6f8d874366788ad5d549e9ba258037d974dda6e004843be1bda794571701/datasets-4.4.1-py3-none-any.whl", hash = "sha256:c1163de5211e42546079ab355cc0250c7e6db16eb209ac5ac6252f801f596c44", size = 511591, upload-time = "2025-11-05T16:00:36.365Z" },
 ]
 
-[[package]]
-name = "debugpy"
-version = "1.8.15"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/8c/8b/3a9a28ddb750a76eaec445c7f4d3147ea2c579a97dbd9e25d39001b92b21/debugpy-1.8.15.tar.gz", hash = "sha256:58d7a20b7773ab5ee6bdfb2e6cf622fdf1e40c9d5aef2857d85391526719ac00", size = 1643279, upload-time = "2025-07-15T16:43:29.135Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d2/b3/1c44a2ed311199ab11c2299c9474a6c7cd80d19278defd333aeb7c287995/debugpy-1.8.15-cp311-cp311-macosx_14_0_universal2.whl", hash = "sha256:babc4fb1962dd6a37e94d611280e3d0d11a1f5e6c72ac9b3d87a08212c4b6dd3", size = 2183442, upload-time = "2025-07-15T16:43:36.733Z" },
-    { url = "https://files.pythonhosted.org/packages/f6/69/e2dcb721491e1c294d348681227c9b44fb95218f379aa88e12a19d85528d/debugpy-1.8.15-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f778e68f2986a58479d0ac4f643e0b8c82fdd97c2e200d4d61e7c2d13838eb53", size = 3134215, upload-time = "2025-07-15T16:43:38.116Z" },
-    { url = "https://files.pythonhosted.org/packages/17/76/4ce63b95d8294dcf2fd1820860b300a420d077df4e93afcaa25a984c2ca7/debugpy-1.8.15-cp311-cp311-win32.whl", hash = "sha256:f9d1b5abd75cd965e2deabb1a06b0e93a1546f31f9f621d2705e78104377c702", size = 5154037, upload-time = "2025-07-15T16:43:39.471Z" },
-    { url = "https://files.pythonhosted.org/packages/c2/a7/e5a7c784465eb9c976d84408873d597dc7ce74a0fc69ed009548a1a94813/debugpy-1.8.15-cp311-cp311-win_amd64.whl", hash = "sha256:62954fb904bec463e2b5a415777f6d1926c97febb08ef1694da0e5d1463c5c3b", size = 5178133, upload-time = "2025-07-15T16:43:40.969Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/4a/4508d256e52897f5cdfee6a6d7580974811e911c6d01321df3264508a5ac/debugpy-1.8.15-cp312-cp312-macosx_14_0_universal2.whl", hash = "sha256:3dcc7225cb317469721ab5136cda9ff9c8b6e6fb43e87c9e15d5b108b99d01ba", size = 2511197, upload-time = "2025-07-15T16:43:42.343Z" },
-    { url = "https://files.pythonhosted.org/packages/99/8d/7f6ef1097e7fecf26b4ef72338d08e41644a41b7ee958a19f494ffcffc29/debugpy-1.8.15-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:047a493ca93c85ccede1dbbaf4e66816794bdc214213dde41a9a61e42d27f8fc", size = 4229517, upload-time = "2025-07-15T16:43:44.14Z" },
-    { url = "https://files.pythonhosted.org/packages/3f/e8/e8c6a9aa33a9c9c6dacbf31747384f6ed2adde4de2e9693c766bdf323aa3/debugpy-1.8.15-cp312-cp312-win32.whl", hash = "sha256:b08e9b0bc260cf324c890626961dad4ffd973f7568fbf57feb3c3a65ab6b6327", size = 5276132, upload-time = "2025-07-15T16:43:45.529Z" },
-    { url = "https://files.pythonhosted.org/packages/e9/ad/231050c6177b3476b85fcea01e565dac83607b5233d003ff067e2ee44d8f/debugpy-1.8.15-cp312-cp312-win_amd64.whl", hash = "sha256:e2a4fe357c92334272eb2845fcfcdbec3ef9f22c16cf613c388ac0887aed15fa", size = 5317645, upload-time = "2025-07-15T16:43:46.968Z" },
-    { url = "https://files.pythonhosted.org/packages/28/70/2928aad2310726d5920b18ed9f54b9f06df5aa4c10cf9b45fa18ff0ab7e8/debugpy-1.8.15-cp313-cp313-macosx_14_0_universal2.whl", hash = "sha256:f5e01291ad7d6649aed5773256c5bba7a1a556196300232de1474c3c372592bf", size = 2495538, upload-time = "2025-07-15T16:43:48.927Z" },
-    { url = "https://files.pythonhosted.org/packages/9e/c6/9b8ffb4ca91fac8b2877eef63c9cc0e87dd2570b1120054c272815ec4cd0/debugpy-1.8.15-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:94dc0f0d00e528d915e0ce1c78e771475b2335b376c49afcc7382ee0b146bab6", size = 4221874, upload-time = "2025-07-15T16:43:50.282Z" },
-    { url = "https://files.pythonhosted.org/packages/55/8a/9b8d59674b4bf489318c7c46a1aab58e606e583651438084b7e029bf3c43/debugpy-1.8.15-cp313-cp313-win32.whl", hash = "sha256:fcf0748d4f6e25f89dc5e013d1129ca6f26ad4da405e0723a4f704583896a709", size = 5275949, upload-time = "2025-07-15T16:43:52.079Z" },
-    { url = "https://files.pythonhosted.org/packages/72/83/9e58e6fdfa8710a5e6ec06c2401241b9ad48b71c0a7eb99570a1f1edb1d3/debugpy-1.8.15-cp313-cp313-win_amd64.whl", hash = "sha256:73c943776cb83e36baf95e8f7f8da765896fd94b05991e7bc162456d25500683", size = 5317720, upload-time = "2025-07-15T16:43:53.703Z" },
-    { url = "https://files.pythonhosted.org/packages/07/d5/98748d9860e767a1248b5e31ffa7ce8cb7006e97bf8abbf3d891d0a8ba4e/debugpy-1.8.15-py2.py3-none-any.whl", hash = "sha256:bce2e6c5ff4f2e00b98d45e7e01a49c7b489ff6df5f12d881c67d2f1ac635f3d", size = 5282697, upload-time = "2025-07-15T16:44:07.996Z" },
-]
-
-[[package]]
-name = "decorator"
-version = "5.2.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" },
-]
-
-[[package]]
-name = "defusedxml"
-version = "0.7.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/0f/d5/c66da9b79e5bdb124974bfe172b4daf3c984ebd9c2a06e2b8a4dc7331c72/defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69", size = 75520, upload-time = "2021-03-08T10:59:26.269Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl", hash = "sha256:a352e7e428770286cc899e2542b6cdaedb2b4953ff269a210103ec58f6198a61", size = 25604, upload-time = "2021-03-08T10:59:24.45Z" },
-]
-
 [[package]]
 name = "deprecated"
 version = "1.2.18"
@@ -799,30 +534,15 @@ wheels = [
 ]
 
 [[package]]
-name = "executing"
-version = "2.2.0"
+name = "donfig"
+version = "0.8.1.post1"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/91/50/a9d80c47ff289c611ff12e63f7c5d13942c65d68125160cefd768c73e6e4/executing-2.2.0.tar.gz", hash = "sha256:5d108c028108fe2551d1a7b2e8b713341e2cb4fc0aa7dcf966fa4327a5226755", size = 978693, upload-time = "2025-01-22T15:41:29.403Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7b/8f/c4d9bafc34ad7ad5d8dc16dd1347ee0e507a52c3adb6bfa8887e1c6a26ba/executing-2.2.0-py2.py3-none-any.whl", hash = "sha256:11387150cad388d62750327a53d3339fad4888b39a6fe233c3afbb54ecffd3aa", size = 26702, upload-time = "2025-01-22T15:41:25.929Z" },
-]
-
-[[package]]
-name = "fasteners"
-version = "0.19"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/5f/d4/e834d929be54bfadb1f3e3b931c38e956aaa3b235a46a3c764c26c774902/fasteners-0.19.tar.gz", hash = "sha256:b4f37c3ac52d8a445af3a66bce57b33b5e90b97c696b7b984f530cf8f0ded09c", size = 24832, upload-time = "2023-09-19T17:11:20.228Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/61/bf/fd60001b3abc5222d8eaa4a204cd8c0ae78e75adc688f33ce4bf25b7fafa/fasteners-0.19-py3-none-any.whl", hash = "sha256:758819cb5d94cdedf4e836988b74de396ceacb8e2794d21f82d131fd9ee77237", size = 18679, upload-time = "2023-09-19T17:11:18.725Z" },
+dependencies = [
+    { name = "pyyaml" },
 ]
-
-[[package]]
-name = "fastjsonschema"
-version = "2.21.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/8b/50/4b769ce1ac4071a1ef6d86b1a3fb56cdc3a37615e8c5519e1af96cdac366/fastjsonschema-2.21.1.tar.gz", hash = "sha256:794d4f0a58f848961ba16af7b9c85a3e88cd360df008c59aac6fc5ae9323b5d4", size = 373939, upload-time = "2024-12-02T10:55:15.133Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/25/71/80cc718ff6d7abfbabacb1f57aaa42e9c1552bfdd01e64ddd704e4a03638/donfig-0.8.1.post1.tar.gz", hash = "sha256:3bef3413a4c1c601b585e8d297256d0c1470ea012afa6e8461dc28bfb7c23f52", size = 19506, upload-time = "2024-05-23T14:14:31.513Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/90/2b/0817a2b257fe88725c25589d89aec060581aabf668707a8d03b2e9e0cb2a/fastjsonschema-2.21.1-py3-none-any.whl", hash = "sha256:c9e5b7e908310918cf494a434eeb31384dd84a98b57a30bcb1f535015b554667", size = 23924, upload-time = "2024-12-02T10:55:07.599Z" },
+    { url = "https://files.pythonhosted.org/packages/0c/d5/c5db1ea3394c6e1732fb3286b3bd878b59507a8f77d32a2cebda7d7b7cd4/donfig-0.8.1.post1-py3-none-any.whl", hash = "sha256:2a3175ce74a06109ff9307d90a230f81215cbac9a751f4d1c6194644b8204f9d", size = 21592, upload-time = "2024-05-23T14:13:55.283Z" },
 ]
 
 [[package]]
@@ -867,15 +587,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d0/9c/df0ef2c51845a13043e5088f7bb988ca6cd5bb82d5d4203d6a158aa58cf2/fonttools-4.59.0-py3-none-any.whl", hash = "sha256:241313683afd3baacb32a6bd124d0bce7404bc5280e12e291bae1b9bba28711d", size = 1128050, upload-time = "2025-07-16T12:04:52.687Z" },
 ]
 
-[[package]]
-name = "fqdn"
-version = "1.5.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/30/3e/a80a8c077fd798951169626cde3e239adeba7dab75deb3555716415bd9b0/fqdn-1.5.1.tar.gz", hash = "sha256:105ed3677e767fb5ca086a0c1f4bb66ebc3c100be518f0e0d755d9eae164d89f", size = 6015, upload-time = "2021-03-11T07:16:29.08Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/cf/58/8acf1b3e91c58313ce5cb67df61001fc9dcd21be4fadb76c1a2d540e09ed/fqdn-1.5.1-py3-none-any.whl", hash = "sha256:3a179af3761e4df6eb2e026ff9e1a3033d3587bf980a0b1b2e1e5d08d7358014", size = 9121, upload-time = "2021-03-11T07:16:28.351Z" },
-]
-
 [[package]]
 name = "frozenlist"
 version = "1.7.0"
@@ -967,15 +678,6 @@ http = [
     { name = "aiohttp" },
 ]
 
-[[package]]
-name = "future"
-version = "1.0.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/a7/b2/4140c69c6a66432916b26158687e821ba631a4c9273c474343badf84d3ba/future-1.0.0.tar.gz", hash = "sha256:bd2968309307861edae1458a4f8a4f3598c03be43b97521076aebf5d94c07b05", size = 1228490, upload-time = "2024-02-21T11:52:38.461Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/da/71/ae30dadffc90b9006d77af76b393cb9dfbfc9629f339fc1574a1c52e6806/future-1.0.0-py3-none-any.whl", hash = "sha256:929292d34f5872e70396626ef385ec22355a1fae8ad29e1a734c3e43f9fbc216", size = 491326, upload-time = "2024-02-21T11:52:35.956Z" },
-]
-
 [[package]]
 name = "gitdb"
 version = "4.0.12"
@@ -1000,6 +702,31 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/01/61/d4b89fec821f72385526e1b9d9a3a0385dda4a72b206d28049e2c7cd39b8/gitpython-3.1.45-py3-none-any.whl", hash = "sha256:8908cb2e02fb3b93b7eb0f2827125cb699869470432cc885f019b8fd0fccff77", size = 208168, upload-time = "2025-07-24T03:45:52.517Z" },
 ]
 
+[[package]]
+name = "google-crc32c"
+version = "1.8.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/03/41/4b9c02f99e4c5fb477122cd5437403b552873f014616ac1d19ac8221a58d/google_crc32c-1.8.0.tar.gz", hash = "sha256:a428e25fb7691024de47fecfbff7ff957214da51eddded0da0ae0e0f03a2cf79", size = 14192, upload-time = "2025-12-16T00:35:25.142Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5d/ef/21ccfaab3d5078d41efe8612e0ed0bfc9ce22475de074162a91a25f7980d/google_crc32c-1.8.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:014a7e68d623e9a4222d663931febc3033c5c7c9730785727de2a81f87d5bab8", size = 31298, upload-time = "2025-12-16T00:20:32.241Z" },
+    { url = "https://files.pythonhosted.org/packages/c5/b8/f8413d3f4b676136e965e764ceedec904fe38ae8de0cdc52a12d8eb1096e/google_crc32c-1.8.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:86cfc00fe45a0ac7359e5214a1704e51a99e757d0272554874f419f79838c5f7", size = 30872, upload-time = "2025-12-16T00:33:58.785Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/fd/33aa4ec62b290477181c55bb1c9302c9698c58c0ce9a6ab4874abc8b0d60/google_crc32c-1.8.0-cp311-cp311-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:19b40d637a54cb71e0829179f6cb41835f0fbd9e8eb60552152a8b52c36cbe15", size = 33243, upload-time = "2025-12-16T00:40:21.46Z" },
+    { url = "https://files.pythonhosted.org/packages/71/03/4820b3bd99c9653d1a5210cb32f9ba4da9681619b4d35b6a052432df4773/google_crc32c-1.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:17446feb05abddc187e5441a45971b8394ea4c1b6efd88ab0af393fd9e0a156a", size = 33608, upload-time = "2025-12-16T00:40:22.204Z" },
+    { url = "https://files.pythonhosted.org/packages/7c/43/acf61476a11437bf9733fb2f70599b1ced11ec7ed9ea760fdd9a77d0c619/google_crc32c-1.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:71734788a88f551fbd6a97be9668a0020698e07b2bf5b3aa26a36c10cdfb27b2", size = 34439, upload-time = "2025-12-16T00:35:20.458Z" },
+    { url = "https://files.pythonhosted.org/packages/e9/5f/7307325b1198b59324c0fa9807cafb551afb65e831699f2ce211ad5c8240/google_crc32c-1.8.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:4b8286b659c1335172e39563ab0a768b8015e88e08329fa5321f774275fc3113", size = 31300, upload-time = "2025-12-16T00:21:56.723Z" },
+    { url = "https://files.pythonhosted.org/packages/21/8e/58c0d5d86e2220e6a37befe7e6a94dd2f6006044b1a33edf1ff6d9f7e319/google_crc32c-1.8.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:2a3dc3318507de089c5384cc74d54318401410f82aa65b2d9cdde9d297aca7cb", size = 30867, upload-time = "2025-12-16T00:38:31.302Z" },
+    { url = "https://files.pythonhosted.org/packages/ce/a9/a780cc66f86335a6019f557a8aaca8fbb970728f0efd2430d15ff1beae0e/google_crc32c-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:14f87e04d613dfa218d6135e81b78272c3b904e2a7053b841481b38a7d901411", size = 33364, upload-time = "2025-12-16T00:40:22.96Z" },
+    { url = "https://files.pythonhosted.org/packages/21/3f/3457ea803db0198c9aaca2dd373750972ce28a26f00544b6b85088811939/google_crc32c-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb5c869c2923d56cb0c8e6bcdd73c009c36ae39b652dbe46a05eb4ef0ad01454", size = 33740, upload-time = "2025-12-16T00:40:23.96Z" },
+    { url = "https://files.pythonhosted.org/packages/df/c0/87c2073e0c72515bb8733d4eef7b21548e8d189f094b5dad20b0ecaf64f6/google_crc32c-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:3cc0c8912038065eafa603b238abf252e204accab2a704c63b9e14837a854962", size = 34437, upload-time = "2025-12-16T00:35:21.395Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/db/000f15b41724589b0e7bc24bc7a8967898d8d3bc8caf64c513d91ef1f6c0/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:3ebb04528e83b2634857f43f9bb8ef5b2bbe7f10f140daeb01b58f972d04736b", size = 31297, upload-time = "2025-12-16T00:23:20.709Z" },
+    { url = "https://files.pythonhosted.org/packages/d7/0d/8ebed0c39c53a7e838e2a486da8abb0e52de135f1b376ae2f0b160eb4c1a/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:450dc98429d3e33ed2926fc99ee81001928d63460f8538f21a5d6060912a8e27", size = 30867, upload-time = "2025-12-16T00:43:14.628Z" },
+    { url = "https://files.pythonhosted.org/packages/ce/42/b468aec74a0354b34c8cbf748db20d6e350a68a2b0912e128cabee49806c/google_crc32c-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3b9776774b24ba76831609ffbabce8cdf6fa2bd5e9df37b594221c7e333a81fa", size = 33344, upload-time = "2025-12-16T00:40:24.742Z" },
+    { url = "https://files.pythonhosted.org/packages/1c/e8/b33784d6fc77fb5062a8a7854e43e1e618b87d5ddf610a88025e4de6226e/google_crc32c-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:89c17d53d75562edfff86679244830599ee0a48efc216200691de8b02ab6b2b8", size = 33694, upload-time = "2025-12-16T00:40:25.505Z" },
+    { url = "https://files.pythonhosted.org/packages/92/b1/d3cbd4d988afb3d8e4db94ca953df429ed6db7282ed0e700d25e6c7bfc8d/google_crc32c-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:57a50a9035b75643996fbf224d6661e386c7162d1dfdab9bc4ca790947d1007f", size = 34435, upload-time = "2025-12-16T00:35:22.107Z" },
+    { url = "https://files.pythonhosted.org/packages/52/c5/c171e4d8c44fec1422d801a6d2e5d7ddabd733eeda505c79730ee9607f07/google_crc32c-1.8.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:87fa445064e7db928226b2e6f0d5304ab4cd0339e664a4e9a25029f384d9bb93", size = 28615, upload-time = "2025-12-16T00:40:29.298Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/97/7d75fe37a7a6ed171a2cf17117177e7aab7e6e0d115858741b41e9dd4254/google_crc32c-1.8.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f639065ea2042d5c034bf258a9f085eaa7af0cd250667c0635a3118e8f92c69c", size = 28800, upload-time = "2025-12-16T00:40:30.322Z" },
+]
+
 [[package]]
 name = "h11"
 version = "0.16.0"
@@ -1055,8 +782,8 @@ name = "httpcore"
 version = "1.0.9"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "certifi" },
-    { name = "h11" },
+    { name = "certifi", marker = "python_full_version >= '3.13'" },
+    { name = "h11", marker = "python_full_version >= '3.13'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" }
 wheels = [
@@ -1068,10 +795,10 @@ name = "httpx"
 version = "0.28.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "anyio" },
-    { name = "certifi" },
-    { name = "httpcore" },
-    { name = "idna" },
+    { name = "anyio", marker = "python_full_version >= '3.13'" },
+    { name = "certifi", marker = "python_full_version >= '3.13'" },
+    { name = "httpcore", marker = "python_full_version >= '3.13'" },
+    { name = "idna", marker = "python_full_version >= '3.13'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" }
 wheels = [
@@ -1090,27 +817,14 @@ dependencies = [
     { name = "pyyaml" },
     { name = "requests" },
     { name = "tqdm" },
-    { name = "typing-extensions" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/91/b4/e6b465eca5386b52cf23cb6df8644ad318a6b0e12b4b96a7e0be09cbfbcc/huggingface_hub-0.34.3.tar.gz", hash = "sha256:d58130fd5aa7408480681475491c0abd7e835442082fbc3ef4d45b6c39f83853", size = 456800, upload-time = "2025-07-29T08:38:53.885Z" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/59/a8/4677014e771ed1591a87b63a2392ce6923baf807193deef302dcfde17542/huggingface_hub-0.34.3-py3-none-any.whl", hash = "sha256:5444550099e2d86e68b2898b09e85878fbd788fc2957b506c6a79ce060e39492", size = 558847, upload-time = "2025-07-29T08:38:51.904Z" },
 ]
 
-[[package]]
-name = "hydra-core"
-version = "1.3.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "antlr4-python3-runtime" },
-    { name = "omegaconf" },
-    { name = "packaging" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" },
-]
-
 [[package]]
 name = "id"
 version = "1.5.0"
@@ -1190,92 +904,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/2c/e1/e6716421ea10d38022b952c159d5161ca1193197fb744506875fbb87ea7b/iniconfig-2.1.0-py3-none-any.whl", hash = "sha256:9deba5723312380e77435581c6bf4935c94cbfab9b1ed33ef8d238ea168eb760", size = 6050, upload-time = "2025-03-19T20:10:01.071Z" },
 ]
 
-[[package]]
-name = "ipykernel"
-version = "6.30.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "appnope", marker = "sys_platform == 'darwin'" },
-    { name = "comm" },
-    { name = "debugpy" },
-    { name = "ipython" },
-    { name = "jupyter-client" },
-    { name = "jupyter-core" },
-    { name = "matplotlib-inline" },
-    { name = "nest-asyncio" },
-    { name = "packaging" },
-    { name = "psutil" },
-    { name = "pyzmq" },
-    { name = "tornado" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/38/27/9e6e30ed92f2ac53d29f70b09da8b2dc456e256148e289678fa0e825f46a/ipykernel-6.30.0.tar.gz", hash = "sha256:b7b808ddb2d261aae2df3a26ff3ff810046e6de3dfbc6f7de8c98ea0a6cb632c", size = 165125, upload-time = "2025-07-21T10:36:09.259Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1f/3d/00813c3d9b46e3dcd88bd4530e0a3c63c0509e5d8c9eff34723ea243ab04/ipykernel-6.30.0-py3-none-any.whl", hash = "sha256:fd2936e55c4a1c2ee8b1e5fa6a372b8eecc0ab1338750dee76f48fa5cca1301e", size = 117264, upload-time = "2025-07-21T10:36:06.854Z" },
-]
-
-[[package]]
-name = "ipython"
-version = "9.4.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
-    { name = "decorator" },
-    { name = "ipython-pygments-lexers" },
-    { name = "jedi" },
-    { name = "matplotlib-inline" },
-    { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
-    { name = "prompt-toolkit" },
-    { name = "pygments" },
-    { name = "stack-data" },
-    { name = "traitlets" },
-    { name = "typing-extensions", marker = "python_full_version < '3.12'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/54/80/406f9e3bde1c1fd9bf5a0be9d090f8ae623e401b7670d8f6fdf2ab679891/ipython-9.4.0.tar.gz", hash = "sha256:c033c6d4e7914c3d9768aabe76bbe87ba1dc66a92a05db6bfa1125d81f2ee270", size = 4385338, upload-time = "2025-07-01T11:11:30.606Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/63/f8/0031ee2b906a15a33d6bfc12dd09c3dfa966b3cb5b284ecfb7549e6ac3c4/ipython-9.4.0-py3-none-any.whl", hash = "sha256:25850f025a446d9b359e8d296ba175a36aedd32e83ca9b5060430fe16801f066", size = 611021, upload-time = "2025-07-01T11:11:27.85Z" },
-]
-
-[[package]]
-name = "ipython-pygments-lexers"
-version = "1.1.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "pygments" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl", hash = "sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c", size = 8074, upload-time = "2025-01-17T11:24:33.271Z" },
-]
-
-[[package]]
-name = "ipywidgets"
-version = "8.1.8"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "comm" },
-    { name = "ipython" },
-    { name = "jupyterlab-widgets" },
-    { name = "traitlets" },
-    { name = "widgetsnbextension" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/4c/ae/c5ce1edc1afe042eadb445e95b0671b03cee61895264357956e61c0d2ac0/ipywidgets-8.1.8.tar.gz", hash = "sha256:61f969306b95f85fba6b6986b7fe45d73124d1d9e3023a8068710d47a22ea668", size = 116739, upload-time = "2025-11-01T21:18:12.393Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/56/6d/0d9848617b9f753b87f214f1c682592f7ca42de085f564352f10f0843026/ipywidgets-8.1.8-py3-none-any.whl", hash = "sha256:ecaca67aed704a338f88f67b1181b58f821ab5dc89c1f0f5ef99db43c1c2921e", size = 139808, upload-time = "2025-11-01T21:18:10.956Z" },
-]
-
-[[package]]
-name = "isoduration"
-version = "20.11.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "arrow" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/7c/1a/3c8edc664e06e6bd06cce40c6b22da5f1429aa4224d0c590f3be21c91ead/isoduration-20.11.0.tar.gz", hash = "sha256:ac2f9015137935279eac671f94f89eb00584f940f5dc49462a0c4ee692ba1bd9", size = 11649, upload-time = "2020-11-01T11:00:00.312Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7b/55/e5326141505c5d5e34c5e0935d2908a74e4561eca44108fbfb9c13d2911a/isoduration-20.11.0-py3-none-any.whl", hash = "sha256:b2904c2a4228c3d44f409c8ae8e2370eb21a26f7ac2ec5446df141dde3452042", size = 11321, upload-time = "2020-11-01T10:59:58.02Z" },
-]
-
 [[package]]
 name = "jaraco-classes"
 version = "3.4.0"
@@ -1312,18 +940,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f3/fd/179a20f832824514df39a90bb0e5372b314fea99f217f5ab942b10a8a4e8/jaraco_functools-4.2.1-py3-none-any.whl", hash = "sha256:590486285803805f4b1f99c60ca9e94ed348d4added84b74c7a12885561e524e", size = 10349, upload-time = "2025-06-21T19:22:02.039Z" },
 ]
 
-[[package]]
-name = "jedi"
-version = "0.19.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "parso" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/72/3a/79a912fbd4d8dd6fbb02bf69afd3bb72cf0c729bb3063c6f4498603db17a/jedi-0.19.2.tar.gz", hash = "sha256:4770dc3de41bde3966b02eb84fbcf557fb33cce26ad23da12c742fb50ecb11f0", size = 1231287, upload-time = "2024-11-11T01:41:42.873Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c0/5a/9cac0c82afec3d09ccd97c8b6502d48f165f9124db81b4bcb90b4af974ee/jedi-0.19.2-py2.py3-none-any.whl", hash = "sha256:a8ef22bde8490f57fe5c7681a3c83cb58874daf72b4784de3cce5b6ef6edb5b9", size = 1572278, upload-time = "2024-11-11T01:41:40.175Z" },
-]
-
 [[package]]
 name = "jeepney"
 version = "0.9.0"
@@ -1345,15 +961,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/31/80/3a54838c3fb461f6fec263ebf3a3a41771bd05190238de3486aae8540c36/jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d", size = 133271, upload-time = "2024-05-05T23:41:59.928Z" },
 ]
 
-[[package]]
-name = "jmespath"
-version = "1.0.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/00/2a/e867e8531cf3e36b41201936b7fa7ba7b5702dbef42922193f05c8976cd6/jmespath-1.0.1.tar.gz", hash = "sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe", size = 25843, upload-time = "2022-06-17T18:00:12.224Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/31/b4/b9b800c45527aadd64d5b442f9b932b00648617eb5d63d2c7a6587b7cafc/jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980", size = 20256, upload-time = "2022-06-17T18:00:10.251Z" },
-]
-
 [[package]]
 name = "joblib"
 version = "1.5.1"
@@ -1363,264 +970,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7d/4f/1195bbac8e0c2acc5f740661631d8d750dc38d4a32b23ee5df3cde6f4e0d/joblib-1.5.1-py3-none-any.whl", hash = "sha256:4719a31f054c7d766948dcd83e9613686b27114f190f717cec7eaa2084f8a74a", size = 307746, upload-time = "2025-05-23T12:04:35.124Z" },
 ]
 
-[[package]]
-name = "json5"
-version = "0.12.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/12/ae/929aee9619e9eba9015207a9d2c1c54db18311da7eb4dcf6d41ad6f0eb67/json5-0.12.1.tar.gz", hash = "sha256:b2743e77b3242f8d03c143dd975a6ec7c52e2f2afe76ed934e53503dd4ad4990", size = 52191, upload-time = "2025-08-12T19:47:42.583Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/85/e2/05328bd2621be49a6fed9e3030b1e51a2d04537d3f816d211b9cc53c5262/json5-0.12.1-py3-none-any.whl", hash = "sha256:d9c9b3bc34a5f54d43c35e11ef7cb87d8bdd098c6ace87117a7b7e83e705c1d5", size = 36119, upload-time = "2025-08-12T19:47:41.131Z" },
-]
-
-[[package]]
-name = "jsonpointer"
-version = "3.0.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/6a/0a/eebeb1fa92507ea94016a2a790b93c2ae41a7e18778f85471dc54475ed25/jsonpointer-3.0.0.tar.gz", hash = "sha256:2b2d729f2091522d61c3b31f82e11870f60b68f43fbc705cb76bf4b832af59ef", size = 9114, upload-time = "2024-06-10T19:24:42.462Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/71/92/5e77f98553e9e75130c78900d000368476aed74276eb8ae8796f65f00918/jsonpointer-3.0.0-py2.py3-none-any.whl", hash = "sha256:13e088adc14fca8b6aa8177c044e12701e6ad4b28ff10e65f2267a90109c9942", size = 7595, upload-time = "2024-06-10T19:24:40.698Z" },
-]
-
-[[package]]
-name = "jsonschema"
-version = "4.25.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "attrs" },
-    { name = "jsonschema-specifications" },
-    { name = "referencing" },
-    { name = "rpds-py" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d5/00/a297a868e9d0784450faa7365c2172a7d6110c763e30ba861867c32ae6a9/jsonschema-4.25.0.tar.gz", hash = "sha256:e63acf5c11762c0e6672ffb61482bdf57f0876684d8d249c0fe2d730d48bc55f", size = 356830, upload-time = "2025-07-18T15:39:45.11Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/fe/54/c86cd8e011fe98803d7e382fd67c0df5ceab8d2b7ad8c5a81524f791551c/jsonschema-4.25.0-py3-none-any.whl", hash = "sha256:24c2e8da302de79c8b9382fee3e76b355e44d2a4364bb207159ce10b517bd716", size = 89184, upload-time = "2025-07-18T15:39:42.956Z" },
-]
-
-[package.optional-dependencies]
-format-nongpl = [
-    { name = "fqdn" },
-    { name = "idna" },
-    { name = "isoduration" },
-    { name = "jsonpointer" },
-    { name = "rfc3339-validator" },
-    { name = "rfc3986-validator" },
-    { name = "rfc3987-syntax" },
-    { name = "uri-template" },
-    { name = "webcolors" },
-]
-
-[[package]]
-name = "jsonschema-specifications"
-version = "2025.4.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "referencing" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/bf/ce/46fbd9c8119cfc3581ee5643ea49464d168028cfb5caff5fc0596d0cf914/jsonschema_specifications-2025.4.1.tar.gz", hash = "sha256:630159c9f4dbea161a6a2205c3011cc4f18ff381b189fff48bb39b9bf26ae608", size = 15513, upload-time = "2025-04-23T12:34:07.418Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/01/0e/b27cdbaccf30b890c40ed1da9fd4a3593a5cf94dae54fb34f8a4b74fcd3f/jsonschema_specifications-2025.4.1-py3-none-any.whl", hash = "sha256:4653bffbd6584f7de83a67e0d620ef16900b390ddc7939d56684d6c81e33f1af", size = 18437, upload-time = "2025-04-23T12:34:05.422Z" },
-]
-
-[[package]]
-name = "jupyter"
-version = "1.1.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "ipykernel" },
-    { name = "ipywidgets" },
-    { name = "jupyter-console" },
-    { name = "jupyterlab" },
-    { name = "nbconvert" },
-    { name = "notebook" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/58/f3/af28ea964ab8bc1e472dba2e82627d36d470c51f5cd38c37502eeffaa25e/jupyter-1.1.1.tar.gz", hash = "sha256:d55467bceabdea49d7e3624af7e33d59c37fff53ed3a350e1ac957bed731de7a", size = 5714959, upload-time = "2024-08-30T07:15:48.299Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/38/64/285f20a31679bf547b75602702f7800e74dbabae36ef324f716c02804753/jupyter-1.1.1-py2.py3-none-any.whl", hash = "sha256:7a59533c22af65439b24bbe60373a4e95af8f16ac65a6c00820ad378e3f7cc83", size = 2657, upload-time = "2024-08-30T07:15:47.045Z" },
-]
-
-[[package]]
-name = "jupyter-client"
-version = "8.6.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jupyter-core" },
-    { name = "python-dateutil" },
-    { name = "pyzmq" },
-    { name = "tornado" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/71/22/bf9f12fdaeae18019a468b68952a60fe6dbab5d67cd2a103cac7659b41ca/jupyter_client-8.6.3.tar.gz", hash = "sha256:35b3a0947c4a6e9d589eb97d7d4cd5e90f910ee73101611f01283732bd6d9419", size = 342019, upload-time = "2024-09-17T10:44:17.613Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/11/85/b0394e0b6fcccd2c1eeefc230978a6f8cb0c5df1e4cd3e7625735a0d7d1e/jupyter_client-8.6.3-py3-none-any.whl", hash = "sha256:e8a19cc986cc45905ac3362915f410f3af85424b4c0905e94fa5f2cb08e8f23f", size = 106105, upload-time = "2024-09-17T10:44:15.218Z" },
-]
-
-[[package]]
-name = "jupyter-console"
-version = "6.6.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "ipykernel" },
-    { name = "ipython" },
-    { name = "jupyter-client" },
-    { name = "jupyter-core" },
-    { name = "prompt-toolkit" },
-    { name = "pygments" },
-    { name = "pyzmq" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/bd/2d/e2fd31e2fc41c14e2bcb6c976ab732597e907523f6b2420305f9fc7fdbdb/jupyter_console-6.6.3.tar.gz", hash = "sha256:566a4bf31c87adbfadf22cdf846e3069b59a71ed5da71d6ba4d8aaad14a53539", size = 34363, upload-time = "2023-03-06T14:13:31.02Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ca/77/71d78d58f15c22db16328a476426f7ac4a60d3a5a7ba3b9627ee2f7903d4/jupyter_console-6.6.3-py3-none-any.whl", hash = "sha256:309d33409fcc92ffdad25f0bcdf9a4a9daa61b6f341177570fdac03de5352485", size = 24510, upload-time = "2023-03-06T14:13:28.229Z" },
-]
-
-[[package]]
-name = "jupyter-core"
-version = "5.8.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "platformdirs" },
-    { name = "pywin32", marker = "platform_python_implementation != 'PyPy' and sys_platform == 'win32'" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/99/1b/72906d554acfeb588332eaaa6f61577705e9ec752ddb486f302dafa292d9/jupyter_core-5.8.1.tar.gz", hash = "sha256:0a5f9706f70e64786b75acba995988915ebd4601c8a52e534a40b51c95f59941", size = 88923, upload-time = "2025-05-27T07:38:16.655Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/2f/57/6bffd4b20b88da3800c5d691e0337761576ee688eb01299eae865689d2df/jupyter_core-5.8.1-py3-none-any.whl", hash = "sha256:c28d268fc90fb53f1338ded2eb410704c5449a358406e8a948b75706e24863d0", size = 28880, upload-time = "2025-05-27T07:38:15.137Z" },
-]
-
-[[package]]
-name = "jupyter-events"
-version = "0.12.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jsonschema", extra = ["format-nongpl"] },
-    { name = "packaging" },
-    { name = "python-json-logger" },
-    { name = "pyyaml" },
-    { name = "referencing" },
-    { name = "rfc3339-validator" },
-    { name = "rfc3986-validator" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/9d/c3/306d090461e4cf3cd91eceaff84bede12a8e52cd821c2d20c9a4fd728385/jupyter_events-0.12.0.tar.gz", hash = "sha256:fc3fce98865f6784c9cd0a56a20644fc6098f21c8c33834a8d9fe383c17e554b", size = 62196, upload-time = "2025-02-03T17:23:41.485Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e2/48/577993f1f99c552f18a0428731a755e06171f9902fa118c379eb7c04ea22/jupyter_events-0.12.0-py3-none-any.whl", hash = "sha256:6464b2fa5ad10451c3d35fabc75eab39556ae1e2853ad0c0cc31b656731a97fb", size = 19430, upload-time = "2025-02-03T17:23:38.643Z" },
-]
-
-[[package]]
-name = "jupyter-lsp"
-version = "2.3.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jupyter-server" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/eb/5a/9066c9f8e94ee517133cd98dba393459a16cd48bba71a82f16a65415206c/jupyter_lsp-2.3.0.tar.gz", hash = "sha256:458aa59339dc868fb784d73364f17dbce8836e906cd75fd471a325cba02e0245", size = 54823, upload-time = "2025-08-27T17:47:34.671Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1a/60/1f6cee0c46263de1173894f0fafcb3475ded276c472c14d25e0280c18d6d/jupyter_lsp-2.3.0-py3-none-any.whl", hash = "sha256:e914a3cb2addf48b1c7710914771aaf1819d46b2e5a79b0f917b5478ec93f34f", size = 76687, upload-time = "2025-08-27T17:47:33.15Z" },
-]
-
-[[package]]
-name = "jupyter-server"
-version = "2.17.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "anyio" },
-    { name = "argon2-cffi" },
-    { name = "jinja2" },
-    { name = "jupyter-client" },
-    { name = "jupyter-core" },
-    { name = "jupyter-events" },
-    { name = "jupyter-server-terminals" },
-    { name = "nbconvert" },
-    { name = "nbformat" },
-    { name = "overrides", marker = "python_full_version < '3.12'" },
-    { name = "packaging" },
-    { name = "prometheus-client" },
-    { name = "pywinpty", marker = "os_name == 'nt' and sys_platform != 'linux'" },
-    { name = "pyzmq" },
-    { name = "send2trash" },
-    { name = "terminado" },
-    { name = "tornado" },
-    { name = "traitlets" },
-    { name = "websocket-client" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/5b/ac/e040ec363d7b6b1f11304cc9f209dac4517ece5d5e01821366b924a64a50/jupyter_server-2.17.0.tar.gz", hash = "sha256:c38ea898566964c888b4772ae1ed58eca84592e88251d2cfc4d171f81f7e99d5", size = 731949, upload-time = "2025-08-21T14:42:54.042Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/92/80/a24767e6ca280f5a49525d987bf3e4d7552bf67c8be07e8ccf20271f8568/jupyter_server-2.17.0-py3-none-any.whl", hash = "sha256:e8cb9c7db4251f51ed307e329b81b72ccf2056ff82d50524debde1ee1870e13f", size = 388221, upload-time = "2025-08-21T14:42:52.034Z" },
-]
-
-[[package]]
-name = "jupyter-server-terminals"
-version = "0.5.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "pywinpty", marker = "os_name == 'nt' and sys_platform != 'linux'" },
-    { name = "terminado" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/fc/d5/562469734f476159e99a55426d697cbf8e7eb5efe89fb0e0b4f83a3d3459/jupyter_server_terminals-0.5.3.tar.gz", hash = "sha256:5ae0295167220e9ace0edcfdb212afd2b01ee8d179fe6f23c899590e9b8a5269", size = 31430, upload-time = "2024-03-12T14:37:03.049Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/07/2d/2b32cdbe8d2a602f697a649798554e4f072115438e92249624e532e8aca6/jupyter_server_terminals-0.5.3-py3-none-any.whl", hash = "sha256:41ee0d7dc0ebf2809c668e0fc726dfaf258fcd3e769568996ca731b6194ae9aa", size = 13656, upload-time = "2024-03-12T14:37:00.708Z" },
-]
-
-[[package]]
-name = "jupyterlab"
-version = "4.4.5"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "async-lru" },
-    { name = "httpx" },
-    { name = "ipykernel" },
-    { name = "jinja2" },
-    { name = "jupyter-core" },
-    { name = "jupyter-lsp" },
-    { name = "jupyter-server" },
-    { name = "jupyterlab-server" },
-    { name = "notebook-shim" },
-    { name = "packaging" },
-    { name = "setuptools" },
-    { name = "tornado" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/20/89/695805a6564bafe08ef2505f3c473ae7140b8ba431d381436f11bdc2c266/jupyterlab-4.4.5.tar.gz", hash = "sha256:0bd6c18e6a3c3d91388af6540afa3d0bb0b2e76287a7b88ddf20ab41b336e595", size = 23037079, upload-time = "2025-07-20T09:21:30.151Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/47/74/e144ce85b34414e44b14c1f6bf2e3bfe17964c8e5670ebdc7629f2e53672/jupyterlab-4.4.5-py3-none-any.whl", hash = "sha256:e76244cceb2d1fb4a99341f3edc866f2a13a9e14c50368d730d75d8017be0863", size = 12267763, upload-time = "2025-07-20T09:21:26.37Z" },
-]
-
-[[package]]
-name = "jupyterlab-pygments"
-version = "0.3.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/90/51/9187be60d989df97f5f0aba133fa54e7300f17616e065d1ada7d7646b6d6/jupyterlab_pygments-0.3.0.tar.gz", hash = "sha256:721aca4d9029252b11cfa9d185e5b5af4d54772bb8072f9b7036f4170054d35d", size = 512900, upload-time = "2023-11-23T09:26:37.44Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b1/dd/ead9d8ea85bf202d90cc513b533f9c363121c7792674f78e0d8a854b63b4/jupyterlab_pygments-0.3.0-py3-none-any.whl", hash = "sha256:841a89020971da1d8693f1a99997aefc5dc424bb1b251fd6322462a1b8842780", size = 15884, upload-time = "2023-11-23T09:26:34.325Z" },
-]
-
-[[package]]
-name = "jupyterlab-server"
-version = "2.28.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "babel" },
-    { name = "jinja2" },
-    { name = "json5" },
-    { name = "jsonschema" },
-    { name = "jupyter-server" },
-    { name = "packaging" },
-    { name = "requests" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d6/2c/90153f189e421e93c4bb4f9e3f59802a1f01abd2ac5cf40b152d7f735232/jupyterlab_server-2.28.0.tar.gz", hash = "sha256:35baa81898b15f93573e2deca50d11ac0ae407ebb688299d3a5213265033712c", size = 76996, upload-time = "2025-10-22T13:59:18.37Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e0/07/a000fe835f76b7e1143242ab1122e6362ef1c03f23f83a045c38859c2ae0/jupyterlab_server-2.28.0-py3-none-any.whl", hash = "sha256:e4355b148fdcf34d312bbbc80f22467d6d20460e8b8736bf235577dd18506968", size = 59830, upload-time = "2025-10-22T13:59:16.767Z" },
-]
-
-[[package]]
-name = "jupyterlab-widgets"
-version = "3.0.16"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/26/2d/ef58fed122b268c69c0aa099da20bc67657cdfb2e222688d5731bd5b971d/jupyterlab_widgets-3.0.16.tar.gz", hash = "sha256:423da05071d55cf27a9e602216d35a3a65a3e41cdf9c5d3b643b814ce38c19e0", size = 897423, upload-time = "2025-11-01T21:11:29.724Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ab/b5/36c712098e6191d1b4e349304ef73a8d06aed77e56ceaac8c0a306c7bda1/jupyterlab_widgets-3.0.16-py3-none-any.whl", hash = "sha256:45fa36d9c6422cf2559198e4db481aa243c7a32d9926b500781c830c80f7ecf8", size = 914926, upload-time = "2025-11-01T21:11:28.008Z" },
-]
-
 [[package]]
 name = "keyring"
 version = "25.6.0"
@@ -1705,15 +1054,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/4c/fa/be89a49c640930180657482a74970cdcf6f7072c8d2471e1babe17a222dc/kiwisolver-1.4.8-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:be4816dc51c8a471749d664161b434912eee82f2ea66bd7628bd14583a833e85", size = 2349213, upload-time = "2024-12-24T18:30:40.019Z" },
 ]
 
-[[package]]
-name = "lark"
-version = "1.3.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/da/34/28fff3ab31ccff1fd4f6c7c7b0ceb2b6968d8ea4950663eadcb5720591a0/lark-1.3.1.tar.gz", hash = "sha256:b426a7a6d6d53189d318f2b6236ab5d6429eaf09259f1ca33eb716eed10d2905", size = 382732, upload-time = "2025-10-27T18:25:56.653Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/82/3d/14ce75ef66813643812f3093ab17e46d3a206942ce7376d31ec2d36229e7/lark-1.3.1-py3-none-any.whl", hash = "sha256:c629b661023a014c37da873b4ff58a817398d12635d3bbb2c5a03be7fe5d1e12", size = 113151, upload-time = "2025-10-27T18:25:54.882Z" },
-]
-
 [[package]]
 name = "legacy-api-wrap"
 version = "1.4.1"
@@ -1766,92 +1106,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d0/81/e66fc86539293282fd9cb7c9417438e897f369e79ffb62e1ae5e5154d4dd/llvmlite-0.44.0-cp313-cp313-win_amd64.whl", hash = "sha256:2fb7c4f2fb86cbae6dca3db9ab203eeea0e22d73b99bc2341cdf9de93612e930", size = 30331193, upload-time = "2025-01-20T11:14:38.578Z" },
 ]
 
-[[package]]
-name = "louvain"
-version = "0.8.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "igraph" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/4b/14/4bf2d8b73639ffedc8a949e6ad40cd2207e0dc4ed9720afeb3f69656c194/louvain-0.8.2.tar.gz", hash = "sha256:ce040bb0c5e349aead5a1e5e99738dcd9f2d10c22526d0633281b6ae230ee8d4", size = 4219094, upload-time = "2024-03-18T10:33:16.182Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/4c/a0/b08b20e4a3366ef39555bcf390cc1489a1d12d76f53413c5cc4cc0e760a6/louvain-0.8.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ad0f20e5e9ed22db3c0cfec25d5c5f004a4a2ca00f556779f5cdc421d816a250", size = 207124, upload-time = "2024-03-18T10:31:47.447Z" },
-    { url = "https://files.pythonhosted.org/packages/fc/26/d99a77f5c17ef366bbd11b3509dfad47bcd968c4c1778bfccd4279e4ea25/louvain-0.8.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3514682845c7b9d8f14e275f3826ad8aeaff079267eadaec1353a52629f1212e", size = 185202, upload-time = "2024-03-18T10:31:48.903Z" },
-    { url = "https://files.pythonhosted.org/packages/ae/e4/9854d680fc4ae73bdc545544dc596d3a5f89f2a50abe597daded7140c4e2/louvain-0.8.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4dc11889d69b596f7f5755a4a197ef96da480b6e09959d3340717836ae3cc255", size = 1051707, upload-time = "2024-03-18T10:31:50.565Z" },
-    { url = "https://files.pythonhosted.org/packages/99/15/dfa3d33935e4c90d5a45676e8d3bdf0cbd0c173e7a7f45dee23e7cc636bb/louvain-0.8.2-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b7ae4a8b11f94acd90509972a2fc5d393ad2a6f327e5506be3ed5e97b6ec94ce", size = 1057039, upload-time = "2024-03-18T10:31:52.077Z" },
-    { url = "https://files.pythonhosted.org/packages/3f/3e/b20691ad452187d88c247202cc144dd55babe0dba20f67fd9f715bb0e6fe/louvain-0.8.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:df66bf58b03ca257299e5652816bae0203a62447f31eec1008f719e0d5604b88", size = 1078332, upload-time = "2024-03-18T10:31:54.328Z" },
-    { url = "https://files.pythonhosted.org/packages/b7/12/1e8a2748dbce63eae1b05ed938834ab1568a7d60da615cf247b8918d4a83/louvain-0.8.2-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:f25d97fe0c4367ea9a9ca3edd731d742b6a0f420e922ea1b4ce4448b8caa0a3b", size = 1605109, upload-time = "2024-03-18T10:31:56.483Z" },
-    { url = "https://files.pythonhosted.org/packages/c3/10/485c648c2dba47aae745f428f77728ec4e7188f77e16c8834e3aaa7c6342/louvain-0.8.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:a4968137c22004b7ffb465f87f798d2cff1efb456c0f4ad2cea98d2eae057ad1", size = 1591018, upload-time = "2024-03-18T10:31:58.877Z" },
-    { url = "https://files.pythonhosted.org/packages/6e/8e/befb2f08063f1eef786b7c911f10588572bba11687b54056b79553730c93/louvain-0.8.2-cp311-cp311-win32.whl", hash = "sha256:5fee00958c0b45bfd43c20b9141370607c91e32cda28e04db4512187005616c9", size = 82079, upload-time = "2024-03-18T10:32:01.004Z" },
-    { url = "https://files.pythonhosted.org/packages/46/a0/ab037cddff15b67fe747bcbb04ab6e2668c78b749e031514af0cee1c6071/louvain-0.8.2-cp311-cp311-win_amd64.whl", hash = "sha256:797fa410e2d0f208202eda27ecd686ebecefd2ce3631a754475408eb7ba01421", size = 89850, upload-time = "2024-03-18T10:32:03.043Z" },
-]
-
-[[package]]
-name = "lxml"
-version = "6.0.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/aa/88/262177de60548e5a2bfc46ad28232c9e9cbde697bd94132aeb80364675cb/lxml-6.0.2.tar.gz", hash = "sha256:cd79f3367bd74b317dda655dc8fcfa304d9eb6e4fb06b7168c5cf27f96e0cd62", size = 4073426, upload-time = "2025-09-22T04:04:59.287Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/77/d5/becbe1e2569b474a23f0c672ead8a29ac50b2dc1d5b9de184831bda8d14c/lxml-6.0.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:13e35cbc684aadf05d8711a5d1b5857c92e5e580efa9a0d2be197199c8def607", size = 8634365, upload-time = "2025-09-22T04:00:45.672Z" },
-    { url = "https://files.pythonhosted.org/packages/28/66/1ced58f12e804644426b85d0bb8a4478ca77bc1761455da310505f1a3526/lxml-6.0.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3b1675e096e17c6fe9c0e8c81434f5736c0739ff9ac6123c87c2d452f48fc938", size = 4650793, upload-time = "2025-09-22T04:00:47.783Z" },
-    { url = "https://files.pythonhosted.org/packages/11/84/549098ffea39dfd167e3f174b4ce983d0eed61f9d8d25b7bf2a57c3247fc/lxml-6.0.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8ac6e5811ae2870953390452e3476694196f98d447573234592d30488147404d", size = 4944362, upload-time = "2025-09-22T04:00:49.845Z" },
-    { url = "https://files.pythonhosted.org/packages/ac/bd/f207f16abf9749d2037453d56b643a7471d8fde855a231a12d1e095c4f01/lxml-6.0.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5aa0fc67ae19d7a64c3fe725dc9a1bb11f80e01f78289d05c6f62545affec438", size = 5083152, upload-time = "2025-09-22T04:00:51.709Z" },
-    { url = "https://files.pythonhosted.org/packages/15/ae/bd813e87d8941d52ad5b65071b1affb48da01c4ed3c9c99e40abb266fbff/lxml-6.0.2-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:de496365750cc472b4e7902a485d3f152ecf57bd3ba03ddd5578ed8ceb4c5964", size = 5023539, upload-time = "2025-09-22T04:00:53.593Z" },
-    { url = "https://files.pythonhosted.org/packages/02/cd/9bfef16bd1d874fbe0cb51afb00329540f30a3283beb9f0780adbb7eec03/lxml-6.0.2-cp311-cp311-manylinux_2_26_i686.manylinux_2_28_i686.whl", hash = "sha256:200069a593c5e40b8f6fc0d84d86d970ba43138c3e68619ffa234bc9bb806a4d", size = 5344853, upload-time = "2025-09-22T04:00:55.524Z" },
-    { url = "https://files.pythonhosted.org/packages/b8/89/ea8f91594bc5dbb879734d35a6f2b0ad50605d7fb419de2b63d4211765cc/lxml-6.0.2-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7d2de809c2ee3b888b59f995625385f74629707c9355e0ff856445cdcae682b7", size = 5225133, upload-time = "2025-09-22T04:00:57.269Z" },
-    { url = "https://files.pythonhosted.org/packages/b9/37/9c735274f5dbec726b2db99b98a43950395ba3d4a1043083dba2ad814170/lxml-6.0.2-cp311-cp311-manylinux_2_31_armv7l.whl", hash = "sha256:b2c3da8d93cf5db60e8858c17684c47d01fee6405e554fb55018dd85fc23b178", size = 4677944, upload-time = "2025-09-22T04:00:59.052Z" },
-    { url = "https://files.pythonhosted.org/packages/20/28/7dfe1ba3475d8bfca3878365075abe002e05d40dfaaeb7ec01b4c587d533/lxml-6.0.2-cp311-cp311-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:442de7530296ef5e188373a1ea5789a46ce90c4847e597856570439621d9c553", size = 5284535, upload-time = "2025-09-22T04:01:01.335Z" },
-    { url = "https://files.pythonhosted.org/packages/e7/cf/5f14bc0de763498fc29510e3532bf2b4b3a1c1d5d0dff2e900c16ba021ef/lxml-6.0.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:2593c77efde7bfea7f6389f1ab249b15ed4aa5bc5cb5131faa3b843c429fbedb", size = 5067343, upload-time = "2025-09-22T04:01:03.13Z" },
-    { url = "https://files.pythonhosted.org/packages/1c/b0/bb8275ab5472f32b28cfbbcc6db7c9d092482d3439ca279d8d6fa02f7025/lxml-6.0.2-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:3e3cb08855967a20f553ff32d147e14329b3ae70ced6edc2f282b94afbc74b2a", size = 4725419, upload-time = "2025-09-22T04:01:05.013Z" },
-    { url = "https://files.pythonhosted.org/packages/25/4c/7c222753bc72edca3b99dbadba1b064209bc8ed4ad448af990e60dcce462/lxml-6.0.2-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:2ed6c667fcbb8c19c6791bbf40b7268ef8ddf5a96940ba9404b9f9a304832f6c", size = 5275008, upload-time = "2025-09-22T04:01:07.327Z" },
-    { url = "https://files.pythonhosted.org/packages/6c/8c/478a0dc6b6ed661451379447cdbec77c05741a75736d97e5b2b729687828/lxml-6.0.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b8f18914faec94132e5b91e69d76a5c1d7b0c73e2489ea8929c4aaa10b76bbf7", size = 5248906, upload-time = "2025-09-22T04:01:09.452Z" },
-    { url = "https://files.pythonhosted.org/packages/2d/d9/5be3a6ab2784cdf9accb0703b65e1b64fcdd9311c9f007630c7db0cfcce1/lxml-6.0.2-cp311-cp311-win32.whl", hash = "sha256:6605c604e6daa9e0d7f0a2137bdc47a2e93b59c60a65466353e37f8272f47c46", size = 3610357, upload-time = "2025-09-22T04:01:11.102Z" },
-    { url = "https://files.pythonhosted.org/packages/e2/7d/ca6fb13349b473d5732fb0ee3eec8f6c80fc0688e76b7d79c1008481bf1f/lxml-6.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:e5867f2651016a3afd8dd2c8238baa66f1e2802f44bc17e236f547ace6647078", size = 4036583, upload-time = "2025-09-22T04:01:12.766Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/a2/51363b5ecd3eab46563645f3a2c3836a2fc67d01a1b87c5017040f39f567/lxml-6.0.2-cp311-cp311-win_arm64.whl", hash = "sha256:4197fb2534ee05fd3e7afaab5d8bfd6c2e186f65ea7f9cd6a82809c887bd1285", size = 3680591, upload-time = "2025-09-22T04:01:14.874Z" },
-    { url = "https://files.pythonhosted.org/packages/f3/c8/8ff2bc6b920c84355146cd1ab7d181bc543b89241cfb1ebee824a7c81457/lxml-6.0.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a59f5448ba2ceccd06995c95ea59a7674a10de0810f2ce90c9006f3cbc044456", size = 8661887, upload-time = "2025-09-22T04:01:17.265Z" },
-    { url = "https://files.pythonhosted.org/packages/37/6f/9aae1008083bb501ef63284220ce81638332f9ccbfa53765b2b7502203cf/lxml-6.0.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:e8113639f3296706fbac34a30813929e29247718e88173ad849f57ca59754924", size = 4667818, upload-time = "2025-09-22T04:01:19.688Z" },
-    { url = "https://files.pythonhosted.org/packages/f1/ca/31fb37f99f37f1536c133476674c10b577e409c0a624384147653e38baf2/lxml-6.0.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a8bef9b9825fa8bc816a6e641bb67219489229ebc648be422af695f6e7a4fa7f", size = 4950807, upload-time = "2025-09-22T04:01:21.487Z" },
-    { url = "https://files.pythonhosted.org/packages/da/87/f6cb9442e4bada8aab5ae7e1046264f62fdbeaa6e3f6211b93f4c0dd97f1/lxml-6.0.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:65ea18d710fd14e0186c2f973dc60bb52039a275f82d3c44a0e42b43440ea534", size = 5109179, upload-time = "2025-09-22T04:01:23.32Z" },
-    { url = "https://files.pythonhosted.org/packages/c8/20/a7760713e65888db79bbae4f6146a6ae5c04e4a204a3c48896c408cd6ed2/lxml-6.0.2-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c371aa98126a0d4c739ca93ceffa0fd7a5d732e3ac66a46e74339acd4d334564", size = 5023044, upload-time = "2025-09-22T04:01:25.118Z" },
-    { url = "https://files.pythonhosted.org/packages/a2/b0/7e64e0460fcb36471899f75831509098f3fd7cd02a3833ac517433cb4f8f/lxml-6.0.2-cp312-cp312-manylinux_2_26_i686.manylinux_2_28_i686.whl", hash = "sha256:700efd30c0fa1a3581d80a748157397559396090a51d306ea59a70020223d16f", size = 5359685, upload-time = "2025-09-22T04:01:27.398Z" },
-    { url = "https://files.pythonhosted.org/packages/b9/e1/e5df362e9ca4e2f48ed6411bd4b3a0ae737cc842e96877f5bf9428055ab4/lxml-6.0.2-cp312-cp312-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c33e66d44fe60e72397b487ee92e01da0d09ba2d66df8eae42d77b6d06e5eba0", size = 5654127, upload-time = "2025-09-22T04:01:29.629Z" },
-    { url = "https://files.pythonhosted.org/packages/c6/d1/232b3309a02d60f11e71857778bfcd4acbdb86c07db8260caf7d008b08f8/lxml-6.0.2-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90a345bbeaf9d0587a3aaffb7006aa39ccb6ff0e96a57286c0cb2fd1520ea192", size = 5253958, upload-time = "2025-09-22T04:01:31.535Z" },
-    { url = "https://files.pythonhosted.org/packages/35/35/d955a070994725c4f7d80583a96cab9c107c57a125b20bb5f708fe941011/lxml-6.0.2-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:064fdadaf7a21af3ed1dcaa106b854077fbeada827c18f72aec9346847cd65d0", size = 4711541, upload-time = "2025-09-22T04:01:33.801Z" },
-    { url = "https://files.pythonhosted.org/packages/1e/be/667d17363b38a78c4bd63cfd4b4632029fd68d2c2dc81f25ce9eb5224dd5/lxml-6.0.2-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fbc74f42c3525ac4ffa4b89cbdd00057b6196bcefe8bce794abd42d33a018092", size = 5267426, upload-time = "2025-09-22T04:01:35.639Z" },
-    { url = "https://files.pythonhosted.org/packages/ea/47/62c70aa4a1c26569bc958c9ca86af2bb4e1f614e8c04fb2989833874f7ae/lxml-6.0.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6ddff43f702905a4e32bc24f3f2e2edfe0f8fde3277d481bffb709a4cced7a1f", size = 5064917, upload-time = "2025-09-22T04:01:37.448Z" },
-    { url = "https://files.pythonhosted.org/packages/bd/55/6ceddaca353ebd0f1908ef712c597f8570cc9c58130dbb89903198e441fd/lxml-6.0.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:6da5185951d72e6f5352166e3da7b0dc27aa70bd1090b0eb3f7f7212b53f1bb8", size = 4788795, upload-time = "2025-09-22T04:01:39.165Z" },
-    { url = "https://files.pythonhosted.org/packages/cf/e8/fd63e15da5e3fd4c2146f8bbb3c14e94ab850589beab88e547b2dbce22e1/lxml-6.0.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:57a86e1ebb4020a38d295c04fc79603c7899e0df71588043eb218722dabc087f", size = 5676759, upload-time = "2025-09-22T04:01:41.506Z" },
-    { url = "https://files.pythonhosted.org/packages/76/47/b3ec58dc5c374697f5ba37412cd2728f427d056315d124dd4b61da381877/lxml-6.0.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:2047d8234fe735ab77802ce5f2297e410ff40f5238aec569ad7c8e163d7b19a6", size = 5255666, upload-time = "2025-09-22T04:01:43.363Z" },
-    { url = "https://files.pythonhosted.org/packages/19/93/03ba725df4c3d72afd9596eef4a37a837ce8e4806010569bedfcd2cb68fd/lxml-6.0.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6f91fd2b2ea15a6800c8e24418c0775a1694eefc011392da73bc6cef2623b322", size = 5277989, upload-time = "2025-09-22T04:01:45.215Z" },
-    { url = "https://files.pythonhosted.org/packages/c6/80/c06de80bfce881d0ad738576f243911fccf992687ae09fd80b734712b39c/lxml-6.0.2-cp312-cp312-win32.whl", hash = "sha256:3ae2ce7d6fedfb3414a2b6c5e20b249c4c607f72cb8d2bb7cc9c6ec7c6f4e849", size = 3611456, upload-time = "2025-09-22T04:01:48.243Z" },
-    { url = "https://files.pythonhosted.org/packages/f7/d7/0cdfb6c3e30893463fb3d1e52bc5f5f99684a03c29a0b6b605cfae879cd5/lxml-6.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:72c87e5ee4e58a8354fb9c7c84cbf95a1c8236c127a5d1b7683f04bed8361e1f", size = 4011793, upload-time = "2025-09-22T04:01:50.042Z" },
-    { url = "https://files.pythonhosted.org/packages/ea/7b/93c73c67db235931527301ed3785f849c78991e2e34f3fd9a6663ffda4c5/lxml-6.0.2-cp312-cp312-win_arm64.whl", hash = "sha256:61cb10eeb95570153e0c0e554f58df92ecf5109f75eacad4a95baa709e26c3d6", size = 3672836, upload-time = "2025-09-22T04:01:52.145Z" },
-    { url = "https://files.pythonhosted.org/packages/53/fd/4e8f0540608977aea078bf6d79f128e0e2c2bba8af1acf775c30baa70460/lxml-6.0.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:9b33d21594afab46f37ae58dfadd06636f154923c4e8a4d754b0127554eb2e77", size = 8648494, upload-time = "2025-09-22T04:01:54.242Z" },
-    { url = "https://files.pythonhosted.org/packages/5d/f4/2a94a3d3dfd6c6b433501b8d470a1960a20ecce93245cf2db1706adf6c19/lxml-6.0.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:6c8963287d7a4c5c9a432ff487c52e9c5618667179c18a204bdedb27310f022f", size = 4661146, upload-time = "2025-09-22T04:01:56.282Z" },
-    { url = "https://files.pythonhosted.org/packages/25/2e/4efa677fa6b322013035d38016f6ae859d06cac67437ca7dc708a6af7028/lxml-6.0.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1941354d92699fb5ffe6ed7b32f9649e43c2feb4b97205f75866f7d21aa91452", size = 4946932, upload-time = "2025-09-22T04:01:58.989Z" },
-    { url = "https://files.pythonhosted.org/packages/ce/0f/526e78a6d38d109fdbaa5049c62e1d32fdd70c75fb61c4eadf3045d3d124/lxml-6.0.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bb2f6ca0ae2d983ded09357b84af659c954722bbf04dea98030064996d156048", size = 5100060, upload-time = "2025-09-22T04:02:00.812Z" },
-    { url = "https://files.pythonhosted.org/packages/81/76/99de58d81fa702cc0ea7edae4f4640416c2062813a00ff24bd70ac1d9c9b/lxml-6.0.2-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eb2a12d704f180a902d7fa778c6d71f36ceb7b0d317f34cdc76a5d05aa1dd1df", size = 5019000, upload-time = "2025-09-22T04:02:02.671Z" },
-    { url = "https://files.pythonhosted.org/packages/b5/35/9e57d25482bc9a9882cb0037fdb9cc18f4b79d85df94fa9d2a89562f1d25/lxml-6.0.2-cp313-cp313-manylinux_2_26_i686.manylinux_2_28_i686.whl", hash = "sha256:6ec0e3f745021bfed19c456647f0298d60a24c9ff86d9d051f52b509663feeb1", size = 5348496, upload-time = "2025-09-22T04:02:04.904Z" },
-    { url = "https://files.pythonhosted.org/packages/a6/8e/cb99bd0b83ccc3e8f0f528e9aa1f7a9965dfec08c617070c5db8d63a87ce/lxml-6.0.2-cp313-cp313-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:846ae9a12d54e368933b9759052d6206a9e8b250291109c48e350c1f1f49d916", size = 5643779, upload-time = "2025-09-22T04:02:06.689Z" },
-    { url = "https://files.pythonhosted.org/packages/d0/34/9e591954939276bb679b73773836c6684c22e56d05980e31d52a9a8deb18/lxml-6.0.2-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ef9266d2aa545d7374938fb5c484531ef5a2ec7f2d573e62f8ce722c735685fd", size = 5244072, upload-time = "2025-09-22T04:02:08.587Z" },
-    { url = "https://files.pythonhosted.org/packages/8d/27/b29ff065f9aaca443ee377aff699714fcbffb371b4fce5ac4ca759e436d5/lxml-6.0.2-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:4077b7c79f31755df33b795dc12119cb557a0106bfdab0d2c2d97bd3cf3dffa6", size = 4718675, upload-time = "2025-09-22T04:02:10.783Z" },
-    { url = "https://files.pythonhosted.org/packages/2b/9f/f756f9c2cd27caa1a6ef8c32ae47aadea697f5c2c6d07b0dae133c244fbe/lxml-6.0.2-cp313-cp313-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a7c5d5e5f1081955358533be077166ee97ed2571d6a66bdba6ec2f609a715d1a", size = 5255171, upload-time = "2025-09-22T04:02:12.631Z" },
-    { url = "https://files.pythonhosted.org/packages/61/46/bb85ea42d2cb1bd8395484fd72f38e3389611aa496ac7772da9205bbda0e/lxml-6.0.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:8f8d0cbd0674ee89863a523e6994ac25fd5be9c8486acfc3e5ccea679bad2679", size = 5057175, upload-time = "2025-09-22T04:02:14.718Z" },
-    { url = "https://files.pythonhosted.org/packages/95/0c/443fc476dcc8e41577f0af70458c50fe299a97bb6b7505bb1ae09aa7f9ac/lxml-6.0.2-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:2cbcbf6d6e924c28f04a43f3b6f6e272312a090f269eff68a2982e13e5d57659", size = 4785688, upload-time = "2025-09-22T04:02:16.957Z" },
-    { url = "https://files.pythonhosted.org/packages/48/78/6ef0b359d45bb9697bc5a626e1992fa5d27aa3f8004b137b2314793b50a0/lxml-6.0.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:dfb874cfa53340009af6bdd7e54ebc0d21012a60a4e65d927c2e477112e63484", size = 5660655, upload-time = "2025-09-22T04:02:18.815Z" },
-    { url = "https://files.pythonhosted.org/packages/ff/ea/e1d33808f386bc1339d08c0dcada6e4712d4ed8e93fcad5f057070b7988a/lxml-6.0.2-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:fb8dae0b6b8b7f9e96c26fdd8121522ce5de9bb5538010870bd538683d30e9a2", size = 5247695, upload-time = "2025-09-22T04:02:20.593Z" },
-    { url = "https://files.pythonhosted.org/packages/4f/47/eba75dfd8183673725255247a603b4ad606f4ae657b60c6c145b381697da/lxml-6.0.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:358d9adae670b63e95bc59747c72f4dc97c9ec58881d4627fe0120da0f90d314", size = 5269841, upload-time = "2025-09-22T04:02:22.489Z" },
-    { url = "https://files.pythonhosted.org/packages/76/04/5c5e2b8577bc936e219becb2e98cdb1aca14a4921a12995b9d0c523502ae/lxml-6.0.2-cp313-cp313-win32.whl", hash = "sha256:e8cd2415f372e7e5a789d743d133ae474290a90b9023197fd78f32e2dc6873e2", size = 3610700, upload-time = "2025-09-22T04:02:24.465Z" },
-    { url = "https://files.pythonhosted.org/packages/fe/0a/4643ccc6bb8b143e9f9640aa54e38255f9d3b45feb2cbe7ae2ca47e8782e/lxml-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:b30d46379644fbfc3ab81f8f82ae4de55179414651f110a1514f0b1f8f6cb2d7", size = 4010347, upload-time = "2025-09-22T04:02:26.286Z" },
-    { url = "https://files.pythonhosted.org/packages/31/ef/dcf1d29c3f530577f61e5fe2f1bd72929acf779953668a8a47a479ae6f26/lxml-6.0.2-cp313-cp313-win_arm64.whl", hash = "sha256:13dcecc9946dca97b11b7c40d29fba63b55ab4170d3c0cf8c0c164343b9bfdcf", size = 3671248, upload-time = "2025-09-22T04:02:27.918Z" },
-    { url = "https://files.pythonhosted.org/packages/0b/11/29d08bc103a62c0eba8016e7ed5aeebbf1e4312e83b0b1648dd203b0e87d/lxml-6.0.2-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:1c06035eafa8404b5cf475bb37a9f6088b0aca288d4ccc9d69389750d5543700", size = 3949829, upload-time = "2025-09-22T04:04:45.608Z" },
-    { url = "https://files.pythonhosted.org/packages/12/b3/52ab9a3b31e5ab8238da241baa19eec44d2ab426532441ee607165aebb52/lxml-6.0.2-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c7d13103045de1bdd6fe5d61802565f1a3537d70cd3abf596aa0af62761921ee", size = 4226277, upload-time = "2025-09-22T04:04:47.754Z" },
-    { url = "https://files.pythonhosted.org/packages/a0/33/1eaf780c1baad88224611df13b1c2a9dfa460b526cacfe769103ff50d845/lxml-6.0.2-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0a3c150a95fbe5ac91de323aa756219ef9cf7fde5a3f00e2281e30f33fa5fa4f", size = 4330433, upload-time = "2025-09-22T04:04:49.907Z" },
-    { url = "https://files.pythonhosted.org/packages/7a/c1/27428a2ff348e994ab4f8777d3a0ad510b6b92d37718e5887d2da99952a2/lxml-6.0.2-pp311-pypy311_pp73-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:60fa43be34f78bebb27812ed90f1925ec99560b0fa1decdb7d12b84d857d31e9", size = 4272119, upload-time = "2025-09-22T04:04:51.801Z" },
-    { url = "https://files.pythonhosted.org/packages/f0/d0/3020fa12bcec4ab62f97aab026d57c2f0cfd480a558758d9ca233bb6a79d/lxml-6.0.2-pp311-pypy311_pp73-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:21c73b476d3cfe836be731225ec3421fa2f048d84f6df6a8e70433dff1376d5a", size = 4417314, upload-time = "2025-09-22T04:04:55.024Z" },
-    { url = "https://files.pythonhosted.org/packages/6c/77/d7f491cbc05303ac6801651aabeb262d43f319288c1ea96c66b1d2692ff3/lxml-6.0.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:27220da5be049e936c3aca06f174e8827ca6445a4353a1995584311487fc4e3e", size = 3518768, upload-time = "2025-09-22T04:04:57.097Z" },
-]
-
 [[package]]
 name = "markdown-it-py"
 version = "3.0.0"
@@ -1935,18 +1189,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1b/92/9a45c91089c3cf690b5badd4be81e392ff086ccca8a1d4e3a08463d8a966/matplotlib-3.10.3-cp313-cp313t-win_amd64.whl", hash = "sha256:4f23ffe95c5667ef8a2b56eea9b53db7f43910fa4a2d5472ae0f72b64deab4d5", size = 8139044, upload-time = "2025-05-08T19:10:44.551Z" },
 ]
 
-[[package]]
-name = "matplotlib-inline"
-version = "0.1.7"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/99/5b/a36a337438a14116b16480db471ad061c36c3694df7c2084a0da7ba538b7/matplotlib_inline-0.1.7.tar.gz", hash = "sha256:8423b23ec666be3d16e16b60bdd8ac4e86e840ebd1dd11a30b9f117f2fa0ab90", size = 8159, upload-time = "2024-04-15T13:44:44.803Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/8f/8e/9ad090d3553c280a8060fbf6e24dc1c0c29704ee7d1c372f0c174aa59285/matplotlib_inline-0.1.7-py3-none-any.whl", hash = "sha256:df192d39a4ff8f21b1895d72e6a13f5fcc5099f00fa84384e0ea28c2cc0653ca", size = 9899, upload-time = "2024-04-15T13:44:43.265Z" },
-]
-
 [[package]]
 name = "mdurl"
 version = "0.1.2"
@@ -1956,77 +1198,71 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" },
 ]
 
-[[package]]
-name = "mistune"
-version = "3.1.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d7/02/a7fb8b21d4d55ac93cdcde9d3638da5dd0ebdd3a4fed76c7725e10b81cbe/mistune-3.1.4.tar.gz", hash = "sha256:b5a7f801d389f724ec702840c11d8fc48f2b33519102fc7ee739e8177b672164", size = 94588, upload-time = "2025-08-29T07:20:43.594Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7a/f0/8282d9641415e9e33df173516226b404d367a0fc55e1a60424a152913abc/mistune-3.1.4-py3-none-any.whl", hash = "sha256:93691da911e5d9d2e23bc54472892aff676df27a75274962ff9edc210364266d", size = 53481, upload-time = "2025-08-29T07:20:42.218Z" },
-]
-
 [[package]]
 name = "mmcontext"
-version = "0.0.1"
+version = "2.0.0"
 source = { editable = "." }
 dependencies = [
     { name = "accelerate" },
-    { name = "adata-hf-datasets" },
     { name = "anndata" },
     { name = "datasets", version = "3.6.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" },
     { name = "datasets", version = "4.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
     { name = "huggingface-hub" },
-    { name = "hydra-core" },
-    { name = "ipykernel" },
-    { name = "louvain" },
-    { name = "psutil" },
+    { name = "numpy" },
     { name = "python-dotenv" },
     { name = "scanpy" },
-    { name = "scib" },
     { name = "sentence-transformers" },
     { name = "torch" },
     { name = "transformers" },
-    { name = "wandb" },
-    { name = "zarr" },
+    { name = "zarr", version = "3.1.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "zarr", version = "3.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
 
 [package.optional-dependencies]
 dev = [
+    { name = "build" },
     { name = "pre-commit" },
+    { name = "ruff" },
     { name = "twine" },
 ]
+eval = [
+    { name = "scib" },
+    { name = "umap-learn" },
+]
 test = [
     { name = "coverage" },
     { name = "pytest" },
     { name = "python-json-logger" },
 ]
+train = [
+    { name = "wandb" },
+]
 
 [package.metadata]
 requires-dist = [
     { name = "accelerate", specifier = ">=1.9" },
-    { name = "adata-hf-datasets", git = "https://github.com/mengerj/adata_hf_datasets.git" },
-    { name = "anndata", specifier = ">=0.11.3" },
+    { name = "anndata", specifier = ">=0.12" },
+    { name = "build", marker = "extra == 'dev'" },
     { name = "coverage", marker = "extra == 'test'" },
     { name = "datasets", specifier = ">=2.19.1,!=4" },
     { name = "huggingface-hub", specifier = ">=0.34.3" },
-    { name = "hydra-core", specifier = ">=1.3.2" },
-    { name = "ipykernel", specifier = ">=6.30" },
-    { name = "louvain", specifier = ">=0.8.2" },
+    { name = "numpy", specifier = ">=1.24" },
     { name = "pre-commit", marker = "extra == 'dev'" },
-    { name = "psutil", specifier = ">=7" },
     { name = "pytest", marker = "extra == 'test'" },
     { name = "python-dotenv", specifier = ">=1.1.1" },
     { name = "python-json-logger", marker = "extra == 'test'" },
+    { name = "ruff", marker = "extra == 'dev'" },
     { name = "scanpy", specifier = ">=1.11.3" },
-    { name = "scib", specifier = ">=1.1.7" },
-    { name = "sentence-transformers", specifier = ">=5" },
+    { name = "scib", marker = "extra == 'eval'", specifier = ">=1.1.7" },
+    { name = "sentence-transformers", specifier = ">=5.4" },
     { name = "torch", specifier = ">=2.5" },
     { name = "transformers", specifier = ">=4.57.1" },
     { name = "twine", marker = "extra == 'dev'" },
-    { name = "wandb", specifier = ">=0.21" },
-    { name = "zarr", specifier = "<3" },
+    { name = "umap-learn", marker = "extra == 'eval'", specifier = ">=0.5" },
+    { name = "wandb", marker = "extra == 'train'", specifier = ">=0.27" },
+    { name = "zarr", specifier = ">=3.1.6" },
 ]
-provides-extras = ["dev", "test"]
+provides-extras = ["eval", "train", "dev", "test"]
 
 [[package]]
 name = "more-itertools"
@@ -2143,18 +1379,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/da/d9/f7f9379981e39b8c2511c9e0326d212accacb82f12fbfdc1aa2ce2a7b2b6/multiprocess-0.70.16-py39-none-any.whl", hash = "sha256:a0bafd3ae1b732eac64be2e72038231c1ba97724b60b09400d68f229fcc2fbf3", size = 133351, upload-time = "2024-01-28T18:52:31.981Z" },
 ]
 
-[[package]]
-name = "mygene"
-version = "3.2.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "biothings-client" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/0a/ec/a256003f84196aa3fdd65a7c6f5adfc0688398fb66442eba75b39c9b7627/mygene-3.2.2.tar.gz", hash = "sha256:e729cabbc28cf5afb221bca1ab637883b375cb1a3e2f067587ec79f71affdaea", size = 5399, upload-time = "2021-04-05T21:24:30.934Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a7/b7/132b1673c0ec00881d49d56c09624942fa0ebd2fc21d73d80647efa082e9/mygene-3.2.2-py2.py3-none-any.whl", hash = "sha256:18d85d1b28ecee2be31d844607fb0c5f7d7c58573278432df819ee2a5e88fe46", size = 5357, upload-time = "2021-04-05T21:24:29.07Z" },
-]
-
 [[package]]
 name = "natsort"
 version = "8.4.0"
@@ -2164,70 +1388,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ef/82/7a9d0550484a62c6da82858ee9419f3dd1ccc9aa1c26a1e43da3ecd20b0d/natsort-8.4.0-py3-none-any.whl", hash = "sha256:4732914fb471f56b5cce04d7bae6f164a592c7712e1c85f9ef585e197299521c", size = 38268, upload-time = "2023-06-20T04:17:17.522Z" },
 ]
 
-[[package]]
-name = "nbclient"
-version = "0.10.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jupyter-client" },
-    { name = "jupyter-core" },
-    { name = "nbformat" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/87/66/7ffd18d58eae90d5721f9f39212327695b749e23ad44b3881744eaf4d9e8/nbclient-0.10.2.tar.gz", hash = "sha256:90b7fc6b810630db87a6d0c2250b1f0ab4cf4d3c27a299b0cde78a4ed3fd9193", size = 62424, upload-time = "2024-12-19T10:32:27.164Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/34/6d/e7fa07f03a4a7b221d94b4d586edb754a9b0dc3c9e2c93353e9fa4e0d117/nbclient-0.10.2-py3-none-any.whl", hash = "sha256:4ffee11e788b4a27fabeb7955547e4318a5298f34342a4bfd01f2e1faaeadc3d", size = 25434, upload-time = "2024-12-19T10:32:24.139Z" },
-]
-
-[[package]]
-name = "nbconvert"
-version = "7.16.6"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "beautifulsoup4" },
-    { name = "bleach", extra = ["css"] },
-    { name = "defusedxml" },
-    { name = "jinja2" },
-    { name = "jupyter-core" },
-    { name = "jupyterlab-pygments" },
-    { name = "markupsafe" },
-    { name = "mistune" },
-    { name = "nbclient" },
-    { name = "nbformat" },
-    { name = "packaging" },
-    { name = "pandocfilters" },
-    { name = "pygments" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a3/59/f28e15fc47ffb73af68a8d9b47367a8630d76e97ae85ad18271b9db96fdf/nbconvert-7.16.6.tar.gz", hash = "sha256:576a7e37c6480da7b8465eefa66c17844243816ce1ccc372633c6b71c3c0f582", size = 857715, upload-time = "2025-01-28T09:29:14.724Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/cc/9a/cd673b2f773a12c992f41309ef81b99da1690426bd2f96957a7ade0d3ed7/nbconvert-7.16.6-py3-none-any.whl", hash = "sha256:1375a7b67e0c2883678c48e506dc320febb57685e5ee67faa51b18a90f3a712b", size = 258525, upload-time = "2025-01-28T09:29:12.551Z" },
-]
-
-[[package]]
-name = "nbformat"
-version = "5.10.4"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "fastjsonschema" },
-    { name = "jsonschema" },
-    { name = "jupyter-core" },
-    { name = "traitlets" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/6d/fd/91545e604bc3dad7dca9ed03284086039b294c6b3d75c0d2fa45f9e9caf3/nbformat-5.10.4.tar.gz", hash = "sha256:322168b14f937a5d11362988ecac2a4952d3d8e3a2cbeb2319584631226d5b3a", size = 142749, upload-time = "2024-04-04T11:20:37.371Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a9/82/0340caa499416c78e5d8f5f05947ae4bc3cba53c9f038ab6e9ed964e22f1/nbformat-5.10.4-py3-none-any.whl", hash = "sha256:3b48d6c8fbca4b299bf3982ea7db1af21580e4fec269ad087b9e81588891200b", size = 78454, upload-time = "2024-04-04T11:20:34.895Z" },
-]
-
-[[package]]
-name = "nest-asyncio"
-version = "1.6.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/83/f8/51569ac65d696c8ecbee95938f89d4abf00f47d58d48f6fbabfe8f0baefe/nest_asyncio-1.6.0.tar.gz", hash = "sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe", size = 7418, upload-time = "2024-01-21T14:25:19.227Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a0/c4/c2971a3ba4c6103a3d10c4b0f24f461ddc027f0f09763220cf35ca1401b3/nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c", size = 5195, upload-time = "2024-01-21T14:25:17.223Z" },
-]
-
 [[package]]
 name = "networkx"
 version = "3.3"
@@ -2279,34 +1439,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d2/1d/1b658dbd2b9fa9c4c9f32accbfc0205d532c8c6194dc0f2a4c0428e7128a/nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9", size = 22314, upload-time = "2024-06-04T18:44:08.352Z" },
 ]
 
-[[package]]
-name = "notebook"
-version = "7.4.5"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jupyter-server" },
-    { name = "jupyterlab" },
-    { name = "jupyterlab-server" },
-    { name = "notebook-shim" },
-    { name = "tornado" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/9f/21/9669982f9569e7478763837e0d35b9fd9f43de0eb5ab5d6ca620b8019cfc/notebook-7.4.5.tar.gz", hash = "sha256:7c2c4ea245913c3ad8ab3e5d36b34a842c06e524556f5c2e1f5d7d08c986615e", size = 13888993, upload-time = "2025-08-05T07:40:56.529Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/fe/c7/207fd1138bd82435d13b6d8640a240be4d855b8ddb41f6bf31aca5be64df/notebook-7.4.5-py3-none-any.whl", hash = "sha256:351635461aca9dad08cf8946a4216f963e2760cc1bf7b1aaaecb23afc33ec046", size = 14295193, upload-time = "2025-08-05T07:40:52.586Z" },
-]
-
-[[package]]
-name = "notebook-shim"
-version = "0.2.4"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jupyter-server" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/54/d2/92fa3243712b9a3e8bafaf60aac366da1cada3639ca767ff4b5b3654ec28/notebook_shim-0.2.4.tar.gz", hash = "sha256:b4b2cfa1b65d98307ca24361f5b30fe785b53c3fd07b7a47e89acb5e6ac638cb", size = 13167, upload-time = "2024-02-14T23:35:18.353Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f9/33/bd5b9137445ea4b680023eb0469b2bb969d61303dedb2aac6560ff3d14a1/notebook_shim-0.2.4-py3-none-any.whl", hash = "sha256:411a5be4e9dc882a074ccbcae671eda64cceb068767e9a3419096986560e1cef", size = 13307, upload-time = "2024-02-14T23:35:16.286Z" },
-]
-
 [[package]]
 name = "numba"
 version = "0.61.2"
@@ -2336,26 +1468,30 @@ wheels = [
 
 [[package]]
 name = "numcodecs"
-version = "0.15.1"
+version = "0.16.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "deprecated" },
     { name = "numpy" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/63/fc/bb532969eb8236984ba65e4f0079a7da885b8ac0ce1f0835decbb3938a62/numcodecs-0.15.1.tar.gz", hash = "sha256:eeed77e4d6636641a2cc605fbc6078c7a8f2cc40f3dfa2b3f61e52e6091b04ff", size = 6267275, upload-time = "2025-02-10T10:23:33.254Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/44/bd/8a391e7c356366224734efd24da929cc4796fff468bfb179fe1af6548535/numcodecs-0.16.5.tar.gz", hash = "sha256:0d0fb60852f84c0bd9543cc4d2ab9eefd37fc8efcc410acd4777e62a1d300318", size = 6276387, upload-time = "2025-11-21T02:49:48.986Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/e4/fc/410f1cacaef0931f5daf06813b1b8a2442f7418ee284ec73fe5e830dca48/numcodecs-0.15.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:698f1d59511488b8fe215fadc1e679a4c70d894de2cca6d8bf2ab770eed34dfd", size = 1649501, upload-time = "2025-02-10T10:23:01.828Z" },
-    { url = "https://files.pythonhosted.org/packages/85/29/dff62fae04323035912c419a82dc9624fad7d08541dbfcd9ab78a3a40074/numcodecs-0.15.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bef8c8e64fab76677324a07672b10c31861775d03fc63ed5012ca384144e4bb9", size = 1187306, upload-time = "2025-02-10T10:23:04.569Z" },
-    { url = "https://files.pythonhosted.org/packages/a6/a8/908a226632ffabf19caf8c99f1b2898f2f22aac02795a6fe9d018fd6d9dd/numcodecs-0.15.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cdfaef9f5f2ed8f65858db801f1953f1007c9613ee490a1c56233cd78b505ed5", size = 8891971, upload-time = "2025-02-10T10:23:07.689Z" },
-    { url = "https://files.pythonhosted.org/packages/2b/e8/058aac43e1300d588e99b2d0d5b771c8a43fa92ce9c9517da596869fc146/numcodecs-0.15.1-cp311-cp311-win_amd64.whl", hash = "sha256:e2547fa3a7ffc9399cfd2936aecb620a3db285f2630c86c8a678e477741a4b3c", size = 840035, upload-time = "2025-02-10T10:23:10.761Z" },
-    { url = "https://files.pythonhosted.org/packages/e7/7e/f12fc32d3beedc6a8f1ec69ea0ba72e93cb99c0350feed2cff5d04679bc3/numcodecs-0.15.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b0a9d9cd29a0088220682dda4a9898321f7813ff7802be2bbb545f6e3d2f10ff", size = 1691889, upload-time = "2025-02-10T10:23:12.934Z" },
-    { url = "https://files.pythonhosted.org/packages/81/38/88e40d40288b73c3b3a390ed5614a34b0661d00255bdd4cfb91c32101364/numcodecs-0.15.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a34f0fe5e5f3b837bbedbeb98794a6d4a12eeeef8d4697b523905837900b5e1c", size = 1189149, upload-time = "2025-02-10T10:23:15.803Z" },
-    { url = "https://files.pythonhosted.org/packages/28/7d/7527d9180bc76011d6163c848c9cf02cd28a623c2c66cf543e1e86de7c5e/numcodecs-0.15.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c3a09e22140f2c691f7df26303ff8fa2dadcf26d7d0828398c0bc09b69e5efa3", size = 8879163, upload-time = "2025-02-10T10:23:18.582Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/bc/b6c3cde91c754860a3467a8c058dcf0b1a5ca14d82b1c5397c700cf8b1eb/numcodecs-0.15.1-cp312-cp312-win_amd64.whl", hash = "sha256:daed6066ffcf40082da847d318b5ab6123d69ceb433ba603cb87c323a541a8bc", size = 836785, upload-time = "2025-02-10T10:23:22.314Z" },
-    { url = "https://files.pythonhosted.org/packages/78/57/acbc54b3419e5be65015e47177c76c0a73e037fd3ae2cde5808169194d4d/numcodecs-0.15.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e3d82b70500cf61e8d115faa0d0a76be6ecdc24a16477ee3279d711699ad85f3", size = 1688220, upload-time = "2025-02-10T10:23:23.79Z" },
-    { url = "https://files.pythonhosted.org/packages/b6/56/9863fa6dc679f40a31bea5e9713ee5507a31dcd3ee82ea4b1a9268ce52e8/numcodecs-0.15.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1d471a1829ce52d3f365053a2bd1379e32e369517557c4027ddf5ac0d99c591e", size = 1180294, upload-time = "2025-02-10T10:23:25.533Z" },
-    { url = "https://files.pythonhosted.org/packages/fa/91/d96999b41e3146b6c0ce6bddc5ad85803cb4d743c95394562c2a4bb8cded/numcodecs-0.15.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1dfdea4a67108205edfce99c1cb6cd621343bc7abb7e16a041c966776920e7de", size = 8834323, upload-time = "2025-02-10T10:23:27.46Z" },
-    { url = "https://files.pythonhosted.org/packages/c3/32/233e5ede6568bdb044e6f99aaa9fa39827ff3109c6487fc137315f733586/numcodecs-0.15.1-cp313-cp313-win_amd64.whl", hash = "sha256:a4f7bdb26f1b34423cb56d48e75821223be38040907c9b5954eeb7463e7eb03c", size = 831955, upload-time = "2025-02-10T10:23:30.601Z" },
+    { url = "https://files.pythonhosted.org/packages/af/85/1ac101a40ead81eaa1c7dc49a8827a30e2e436211b43ebdc63c590eb1347/numcodecs-0.16.5-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:78382dcea50622f2ef1e6e7a71dbe7f861d8fe376b27b7c297c26907304fef1e", size = 1621795, upload-time = "2025-11-21T02:49:17.418Z" },
+    { url = "https://files.pythonhosted.org/packages/0e/cc/0d97ef55dda48cb0f93d7b92d761208e7a99bd2eea6b0e859426e6a99a21/numcodecs-0.16.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e2d04a19cb57a3c519b4127ac377cca6471aee1990d7c18f5b1e3a4fe1306689", size = 1153030, upload-time = "2025-11-21T02:49:19.089Z" },
+    { url = "https://files.pythonhosted.org/packages/5e/41/e120ee1b390730ac5987cde2afd82e2b8442cec315ab40b94b0373e93e73/numcodecs-0.16.5-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c043af648eb280cd61785c99c22ff5c3c3460f906eb51a8511327c4f5111b283", size = 8510503, upload-time = "2025-11-21T02:49:20.324Z" },
+    { url = "https://files.pythonhosted.org/packages/54/4b/195ac84cc8f6077b4f0f421e8daee21b7f1bd88cb7716414234379fe68ec/numcodecs-0.16.5-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c398919ef2eb0e56b8e97456f622640bfd3deed06de3acc976989cbcb22628a3", size = 9123428, upload-time = "2025-11-21T02:49:22.328Z" },
+    { url = "https://files.pythonhosted.org/packages/0f/5b/af02c417954f46e5c7bd5163ac251f535877d909fce54861c99ae197f6f6/numcodecs-0.16.5-cp311-cp311-win_amd64.whl", hash = "sha256:3820860ed302d4d84a1c66e70981ff959d5eb712555be4e7d8ced49888594773", size = 801542, upload-time = "2025-11-21T02:49:24.265Z" },
+    { url = "https://files.pythonhosted.org/packages/75/cc/55420f3641a67f78392dc0bc5d02cb9eb0a9dcebf2848d1ac77253ca61fa/numcodecs-0.16.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:24e675dc8d1550cd976a99479b87d872cb142632c75cc402fea04c08c4898523", size = 1656287, upload-time = "2025-11-21T02:49:25.755Z" },
+    { url = "https://files.pythonhosted.org/packages/f5/6c/86644987505dcb90ba6d627d6989c27bafb0699f9fd00187e06d05ea8594/numcodecs-0.16.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:94ddfa4341d1a3ab99989d13b01b5134abb687d3dab2ead54b450aefe4ad5bd6", size = 1148899, upload-time = "2025-11-21T02:49:26.87Z" },
+    { url = "https://files.pythonhosted.org/packages/97/1e/98aaddf272552d9fef1f0296a9939d1487914a239e98678f6b20f8b0a5c8/numcodecs-0.16.5-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b554ab9ecf69de7ca2b6b5e8bc696bd9747559cb4dd5127bd08d7a28bec59c3a", size = 8534814, upload-time = "2025-11-21T02:49:28.547Z" },
+    { url = "https://files.pythonhosted.org/packages/fb/53/78c98ef5c8b2b784453487f3e4d6c017b20747c58b470393e230c78d18e8/numcodecs-0.16.5-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ad1a379a45bd3491deab8ae6548313946744f868c21d5340116977ea3be5b1d6", size = 9173471, upload-time = "2025-11-21T02:49:30.444Z" },
+    { url = "https://files.pythonhosted.org/packages/1c/20/2fdec87fc7f8cec950d2b0bea603c12dc9f05b4966dc5924ba5a36a61bf6/numcodecs-0.16.5-cp312-cp312-win_amd64.whl", hash = "sha256:845a9857886ffe4a3172ba1c537ae5bcc01e65068c31cf1fce1a844bd1da050f", size = 801412, upload-time = "2025-11-21T02:49:32.123Z" },
+    { url = "https://files.pythonhosted.org/packages/38/38/071ced5a5fd1c85ba0e14ba721b66b053823e5176298c2f707e50bed11d9/numcodecs-0.16.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:25be3a516ab677dad890760d357cfe081a371d9c0a2e9a204562318ac5969de3", size = 1654359, upload-time = "2025-11-21T02:49:33.673Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/c0/5f84ba7525577c1b9909fc2d06ef11314825fc4ad4378f61d0e4c9883b4a/numcodecs-0.16.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0107e839ef75b854e969cb577e140b1aadb9847893937636582d23a2a4c6ce50", size = 1144237, upload-time = "2025-11-21T02:49:35.294Z" },
+    { url = "https://files.pythonhosted.org/packages/0b/00/787ea5f237b8ea7bc67140c99155f9c00b5baf11c49afc5f3bfefa298f95/numcodecs-0.16.5-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:015a7c859ecc2a06e2a548f64008c0ec3aaecabc26456c2c62f4278d8fc20597", size = 8483064, upload-time = "2025-11-21T02:49:36.454Z" },
+    { url = "https://files.pythonhosted.org/packages/c4/e6/d359fdd37498e74d26a167f7a51e54542e642ea47181eb4e643a69a066c3/numcodecs-0.16.5-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:84230b4b9dad2392f2a84242bd6e3e659ac137b5a1ce3571d6965fca673e0903", size = 9126063, upload-time = "2025-11-21T02:49:38.018Z" },
+    { url = "https://files.pythonhosted.org/packages/27/72/6663cc0382ddbb866136c255c837bcb96cc7ce5e83562efec55e1b995941/numcodecs-0.16.5-cp313-cp313-win_amd64.whl", hash = "sha256:5088145502ad1ebf677ec47d00eb6f0fd600658217db3e0c070c321c85d6cf3d", size = 799275, upload-time = "2025-11-21T02:49:39.558Z" },
 ]
 
 [[package]]
@@ -2506,15 +1642,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" },
 ]
 
-[[package]]
-name = "nvidia-ml-py"
-version = "12.575.51"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d2/4d/6f017814ed5ac28e08e1b8a62e3a258957da27582c89b7f8f8b15ac3d2e7/nvidia_ml_py-12.575.51.tar.gz", hash = "sha256:6490e93fea99eb4e966327ae18c6eec6256194c921f23459c8767aee28c54581", size = 46597, upload-time = "2025-05-06T20:46:37.962Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/db/24/552ebea28f0570b9e65e62b50287a273804c9f997cc1c2dcd4e2d64b9e7d/nvidia_ml_py-12.575.51-py3-none-any.whl", hash = "sha256:eb8641800d98ce40a22f479873f34b482e214a7e80349c63be51c3919845446e", size = 47547, upload-time = "2025-05-06T20:46:36.457Z" },
-]
-
 [[package]]
 name = "nvidia-nccl-cu12"
 version = "2.27.3"
@@ -2539,35 +1666,13 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" },
 ]
 
-[[package]]
-name = "omegaconf"
-version = "2.3.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "antlr4-python3-runtime" },
-    { name = "pyyaml" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" },
-]
-
-[[package]]
-name = "overrides"
-version = "7.7.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/36/86/b585f53236dec60aba864e050778b25045f857e17f6e5ea0ae95fe80edd2/overrides-7.7.0.tar.gz", hash = "sha256:55158fa3d93b98cc75299b1e67078ad9003ca27945c76162c1c0766d6f91820a", size = 22812, upload-time = "2024-01-27T21:01:33.423Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/2c/ab/fc8290c6a4c722e5514d80f62b2dc4c4df1a68a41d1364e625c35990fcf3/overrides-7.7.0-py3-none-any.whl", hash = "sha256:c7ed9d062f78b8e4c1a7b70bd8796b35ead4d9f510227ef9c5dc7626c60d7e49", size = 17832, upload-time = "2024-01-27T21:01:31.393Z" },
-]
-
 [[package]]
 name = "packaging"
-version = "24.1"
+version = "26.2"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/51/65/50db4dda066951078f0a96cf12f4b9ada6e4b811516bf0262c0f4f7064d4/packaging-24.1.tar.gz", hash = "sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002", size = 148788, upload-time = "2024-06-09T23:19:24.956Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/08/aa/cc0199a5f0ad350994d660967a8efb233fe0416e4639146c089643407ce6/packaging-24.1-py3-none-any.whl", hash = "sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124", size = 53985, upload-time = "2024-06-09T23:19:21.909Z" },
+    { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" },
 ]
 
 [[package]]
@@ -2612,45 +1717,15 @@ wheels = [
 ]
 
 [[package]]
-name = "pandocfilters"
-version = "1.5.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/70/6f/3dd4940bbe001c06a65f88e36bad298bc7a0de5036115639926b0c5c0458/pandocfilters-1.5.1.tar.gz", hash = "sha256:002b4a555ee4ebc03f8b66307e287fa492e4a77b4ea14d3f934328297bb4939e", size = 8454, upload-time = "2024-01-18T20:08:13.726Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ef/af/4fbc8cab944db5d21b7e2a5b8e9211a03a79852b1157e2c102fcc61ac440/pandocfilters-1.5.1-py2.py3-none-any.whl", hash = "sha256:93be382804a9cdb0a7267585f157e5d1731bbe5545a85b268d6f5fe6232de2bc", size = 8663, upload-time = "2024-01-18T20:08:11.28Z" },
-]
-
-[[package]]
-name = "parso"
-version = "0.8.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/66/94/68e2e17afaa9169cf6412ab0f28623903be73d1b32e208d9e8e541bb086d/parso-0.8.4.tar.gz", hash = "sha256:eb3a7b58240fb99099a345571deecc0f9540ea5f4dd2fe14c2a99d6b281ab92d", size = 400609, upload-time = "2024-04-05T09:43:55.897Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c6/ac/dac4a63f978e4dcb3c6d3a78c4d8e0192a113d288502a1216950c41b1027/parso-0.8.4-py2.py3-none-any.whl", hash = "sha256:a418670a20291dacd2dddc80c377c5c3791378ee1e8d12bffc35420643d43f18", size = 103650, upload-time = "2024-04-05T09:43:53.299Z" },
-]
-
-[[package]]
-name = "patsy"
-version = "1.0.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d1/81/74f6a65b848ffd16c18f920620ce999fe45fe27f01ab3911260ce4ed85e4/patsy-1.0.1.tar.gz", hash = "sha256:e786a9391eec818c054e359b737bbce692f051aee4c661f4141cc88fb459c0c4", size = 396010, upload-time = "2024-11-12T14:10:54.642Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/87/2b/b50d3d08ea0fc419c183a84210571eba005328efa62b6b98bc28e9ead32a/patsy-1.0.1-py2.py3-none-any.whl", hash = "sha256:751fb38f9e97e62312e921a1954b81e1bb2bcda4f5eeabaf94db251ee791509c", size = 232923, upload-time = "2024-11-12T14:10:52.85Z" },
-]
-
-[[package]]
-name = "pexpect"
-version = "4.9.0"
+name = "patsy"
+version = "1.0.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "ptyprocess" },
+    { name = "numpy" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/d1/81/74f6a65b848ffd16c18f920620ce999fe45fe27f01ab3911260ce4ed85e4/patsy-1.0.1.tar.gz", hash = "sha256:e786a9391eec818c054e359b737bbce692f051aee4c661f4141cc88fb459c0c4", size = 396010, upload-time = "2024-11-12T14:10:54.642Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/9e/c3/059298687310d527a58bb01f3b1965787ee3b40dce76752eda8b44e9a2c5/pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523", size = 63772, upload-time = "2023-11-25T06:56:14.81Z" },
+    { url = "https://files.pythonhosted.org/packages/87/2b/b50d3d08ea0fc419c183a84210571eba005328efa62b6b98bc28e9ead32a/patsy-1.0.1-py2.py3-none-any.whl", hash = "sha256:751fb38f9e97e62312e921a1954b81e1bb2bcda4f5eeabaf94db251ee791509c", size = 232923, upload-time = "2024-11-12T14:10:52.85Z" },
 ]
 
 [[package]]
@@ -2736,27 +1811,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/88/74/a88bf1b1efeae488a0c0b7bdf71429c313722d1fc0f377537fbe554e6180/pre_commit-4.2.0-py2.py3-none-any.whl", hash = "sha256:a009ca7205f1eb497d10b845e52c838a98b6cdd2102a6c8e4540e94ee75c58bd", size = 220707, upload-time = "2025-03-18T21:35:19.343Z" },
 ]
 
-[[package]]
-name = "prometheus-client"
-version = "0.23.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/23/53/3edb5d68ecf6b38fcbcc1ad28391117d2a322d9a1a3eff04bfdb184d8c3b/prometheus_client-0.23.1.tar.gz", hash = "sha256:6ae8f9081eaaaf153a2e959d2e6c4f4fb57b12ef76c8c7980202f1e57b48b2ce", size = 80481, upload-time = "2025-09-18T20:47:25.043Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b8/db/14bafcb4af2139e046d03fd00dea7873e48eafe18b7d2797e73d6681f210/prometheus_client-0.23.1-py3-none-any.whl", hash = "sha256:dd1913e6e76b59cfe44e7a4b83e01afc9873c1bdfd2ed8739f1e76aeca115f99", size = 61145, upload-time = "2025-09-18T20:47:23.875Z" },
-]
-
-[[package]]
-name = "prompt-toolkit"
-version = "3.0.51"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "wcwidth" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/bb/6e/9d084c929dfe9e3bfe0c6a47e31f78a25c54627d64a66e884a8bf5474f1c/prompt_toolkit-3.0.51.tar.gz", hash = "sha256:931a162e3b27fc90c86f1b48bb1fb2c528c2761475e57c9c06de13311c7b54ed", size = 428940, upload-time = "2025-04-15T09:18:47.731Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ce/4f/5249960887b1fbe561d9ff265496d170b55a735b76724f10ef19f9e40716/prompt_toolkit-3.0.51-py3-none-any.whl", hash = "sha256:52742911fde84e2d423e2f9a4cf1de7d7ac4e51958f648d9540e0fb8db077b07", size = 387810, upload-time = "2025-04-15T09:18:44.753Z" },
-]
-
 [[package]]
 name = "propcache"
 version = "0.3.2"
@@ -2859,24 +1913,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/50/1b/6921afe68c74868b4c9fa424dad3be35b095e16687989ebbb50ce4fceb7c/psutil-7.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:4cf3d4eb1aa9b348dec30105c55cd9b7d4629285735a102beb4441e38db90553", size = 244885, upload-time = "2025-02-13T21:54:37.486Z" },
 ]
 
-[[package]]
-name = "ptyprocess"
-version = "0.7.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/20/e5/16ff212c1e452235a90aeb09066144d0c5a6a8c0834397e03f5224495c4e/ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220", size = 70762, upload-time = "2020-12-28T15:15:30.155Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" },
-]
-
-[[package]]
-name = "pure-eval"
-version = "0.2.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/cd/05/0a34433a064256a578f1783a10da6df098ceaa4a57bbeaa96a6c0352786b/pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42", size = 19752, upload-time = "2024-07-21T12:58:21.801Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" },
-]
-
 [[package]]
 name = "pyarrow"
 version = "17.0.0"
@@ -2948,21 +1984,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/60/ae/b5a5811e11f25788ccfdaa8f26b6791c9807119dffcf80514505527c384c/pyarrow-22.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:bba208d9c7decf9961998edf5c65e3ea4355d5818dd6cd0f6809bec1afb951cc", size = 28262504, upload-time = "2025-10-24T10:08:00.932Z" },
 ]
 
-[[package]]
-name = "pybiomart"
-version = "0.2.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "future" },
-    { name = "pandas" },
-    { name = "requests" },
-    { name = "requests-cache" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/91/22/3c7f20669238c228c2768c733377078e57b60bc58b93a7b517f56b795df0/pybiomart-0.2.0.tar.gz", hash = "sha256:e9eac20db921820670c646d99725b0ee279e407379e5e8c3ec7245a07425d8fe", size = 195840, upload-time = "2017-05-10T21:17:59.211Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/4a/6d/23cf5ae8b0a549a60164c49836c8e8eb393747b5fd00a7b5d9445886e0d8/pybiomart-0.2.0-py3-none-any.whl", hash = "sha256:c88bdd97887e2144a8be44036901aa032a824c82b42c33e24515b6c786b80933", size = 10586, upload-time = "2017-05-10T21:17:57.826Z" },
-]
-
 [[package]]
 name = "pycparser"
 version = "2.22"
@@ -2979,7 +2000,8 @@ source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "annotated-types" },
     { name = "pydantic-core" },
-    { name = "typing-extensions" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
     { name = "typing-inspection" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/00/dd/4325abf92c39ba8623b5af936ddb36ffcfe0beae70405d456ab1fb2f5b8c/pydantic-2.11.7.tar.gz", hash = "sha256:d989c3c6cb79469287b1569f7447a17848c998458d49ebe294e975b9baf0f0db", size = 788350, upload-time = "2025-06-14T08:33:17.137Z" }
@@ -2992,7 +2014,8 @@ name = "pydantic-core"
 version = "2.33.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "typing-extensions" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/ad/88/5f2260bdfae97aabf98f1778d43f69574390ad787afb646292a638c923d4/pydantic_core-2.33.2.tar.gz", hash = "sha256:7cb8bc3605c29176e1b105350d2e6474142d7c1bd1d9327c4a9bdb46bf827acc", size = 435195, upload-time = "2025-04-23T18:33:52.104Z" }
 wheels = [
@@ -3089,18 +2112,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d2/53/d23a97e0a2c690d40b165d1062e2c4ccc796be458a1ce59f6ba030434663/pynndescent-0.5.13-py3-none-any.whl", hash = "sha256:69aabb8f394bc631b6ac475a1c7f3994c54adf3f51cd63b2730fefba5771b949", size = 56850, upload-time = "2024-06-17T15:48:31.184Z" },
 ]
 
-[[package]]
-name = "pynvml"
-version = "12.0.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "nvidia-ml-py" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/26/6f/6b5880ed0239e85b9a39aed103b65b2ef81425beef9f45e5c035bf008330/pynvml-12.0.0.tar.gz", hash = "sha256:299ce2451a6a17e6822d6faee750103e25b415f06f59abb8db65d30f794166f5", size = 33636, upload-time = "2024-12-02T15:04:36.631Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ed/df/f7cf07a65a96dd11d71f346f9c2863accdd4784da83af7181b067d556cbc/pynvml-12.0.0-py3-none-any.whl", hash = "sha256:fdff84b62a27dbe98e08e1a647eb77342bef1aebe0878bcd15e99a83fcbecb9e", size = 26560, upload-time = "2024-12-02T15:04:35.047Z" },
-]
-
 [[package]]
 name = "pyparsing"
 version = "3.2.3"
@@ -3111,19 +2122,12 @@ wheels = [
 ]
 
 [[package]]
-name = "pysradb"
-version = "2.5.1"
+name = "pyproject-hooks"
+version = "1.2.0"
 source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "lxml" },
-    { name = "pandas" },
-    { name = "requests" },
-    { name = "tqdm" },
-    { name = "xmltodict" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/74/d4/b7ff41d870f44ec288bbd4e52bbc43d52610c4f436a2e624d123b4924acf/pysradb-2.5.1.tar.gz", hash = "sha256:030d3f70c263d0dae28048c2ed158b0dbf49023bb7d26f1fa9939c13227644de", size = 175574, upload-time = "2025-10-20T10:15:04.611Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e7/82/28175b2414effca1cdac8dc99f76d660e7a4fb0ceefa4b4ab8f5f6742925/pyproject_hooks-1.2.0.tar.gz", hash = "sha256:1e859bd5c40fae9448642dd871adf459e5e2084186e8d2c2a79a824c970da1f8", size = 19228, upload-time = "2024-09-29T09:24:13.293Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/81/b8/a631ba05927d98513446644c247f8d66f461115ae6d8c2d7e2884f809b76/pysradb-2.5.1-py3-none-any.whl", hash = "sha256:b37892427bd890a70d97e7bf908cf69f8d9356e892e51e27f6ec8b61e8bf64fc", size = 181484, upload-time = "2025-10-20T10:14:59.772Z" },
+    { url = "https://files.pythonhosted.org/packages/bd/24/12818598c362d7f300f18e74db45963dbcb85150324092410c8b49405e42/pyproject_hooks-1.2.0-py3-none-any.whl", hash = "sha256:9e5c6bfa8dcc30091c74b0cf803c81fdd29d94f01992a7707bc97babb1141913", size = 10216, upload-time = "2024-09-29T09:24:11.978Z" },
 ]
 
 [[package]]
@@ -3181,22 +2185,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/81/c4/34e93fe5f5429d7570ec1fa436f1986fb1f00c3e0f43a589fe2bbcd22c3f/pytz-2025.2-py2.py3-none-any.whl", hash = "sha256:5ddf76296dd8c44c26eb8f4b6f35488f3ccbf6fbbd7adee0b7262d43f0ec2f00", size = 509225, upload-time = "2025-03-25T02:24:58.468Z" },
 ]
 
-[[package]]
-name = "pywin32"
-version = "311"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7c/af/449a6a91e5d6db51420875c54f6aff7c97a86a3b13a0b4f1a5c13b988de3/pywin32-311-cp311-cp311-win32.whl", hash = "sha256:184eb5e436dea364dcd3d2316d577d625c0351bf237c4e9a5fabbcfa5a58b151", size = 8697031, upload-time = "2025-07-14T20:13:13.266Z" },
-    { url = "https://files.pythonhosted.org/packages/51/8f/9bb81dd5bb77d22243d33c8397f09377056d5c687aa6d4042bea7fbf8364/pywin32-311-cp311-cp311-win_amd64.whl", hash = "sha256:3ce80b34b22b17ccbd937a6e78e7225d80c52f5ab9940fe0506a1a16f3dab503", size = 9508308, upload-time = "2025-07-14T20:13:15.147Z" },
-    { url = "https://files.pythonhosted.org/packages/44/7b/9c2ab54f74a138c491aba1b1cd0795ba61f144c711daea84a88b63dc0f6c/pywin32-311-cp311-cp311-win_arm64.whl", hash = "sha256:a733f1388e1a842abb67ffa8e7aad0e70ac519e09b0f6a784e65a136ec7cefd2", size = 8703930, upload-time = "2025-07-14T20:13:16.945Z" },
-    { url = "https://files.pythonhosted.org/packages/e7/ab/01ea1943d4eba0f850c3c61e78e8dd59757ff815ff3ccd0a84de5f541f42/pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31", size = 8706543, upload-time = "2025-07-14T20:13:20.765Z" },
-    { url = "https://files.pythonhosted.org/packages/d1/a8/a0e8d07d4d051ec7502cd58b291ec98dcc0c3fff027caad0470b72cfcc2f/pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067", size = 9495040, upload-time = "2025-07-14T20:13:22.543Z" },
-    { url = "https://files.pythonhosted.org/packages/ba/3a/2ae996277b4b50f17d61f0603efd8253cb2d79cc7ae159468007b586396d/pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852", size = 8710102, upload-time = "2025-07-14T20:13:24.682Z" },
-    { url = "https://files.pythonhosted.org/packages/a5/be/3fd5de0979fcb3994bfee0d65ed8ca9506a8a1260651b86174f6a86f52b3/pywin32-311-cp313-cp313-win32.whl", hash = "sha256:f95ba5a847cba10dd8c4d8fefa9f2a6cf283b8b88ed6178fa8a6c1ab16054d0d", size = 8705700, upload-time = "2025-07-14T20:13:26.471Z" },
-    { url = "https://files.pythonhosted.org/packages/e3/28/e0a1909523c6890208295a29e05c2adb2126364e289826c0a8bc7297bd5c/pywin32-311-cp313-cp313-win_amd64.whl", hash = "sha256:718a38f7e5b058e76aee1c56ddd06908116d35147e133427e59a3983f703a20d", size = 9494700, upload-time = "2025-07-14T20:13:28.243Z" },
-    { url = "https://files.pythonhosted.org/packages/04/bf/90339ac0f55726dce7d794e6d79a18a91265bdf3aa70b6b9ca52f35e022a/pywin32-311-cp313-cp313-win_arm64.whl", hash = "sha256:7b4075d959648406202d92a2310cb990fea19b535c7f4a78d3f5e10b926eeb8a", size = 8709318, upload-time = "2025-07-14T20:13:30.348Z" },
-]
-
 [[package]]
 name = "pywin32-ctypes"
 version = "0.2.3"
@@ -3206,18 +2194,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl", hash = "sha256:8a1513379d709975552d202d942d9837758905c8d01eb82b8bcc30918929e7b8", size = 30756, upload-time = "2024-08-14T10:15:33.187Z" },
 ]
 
-[[package]]
-name = "pywinpty"
-version = "3.0.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/f3/bb/a7cc2967c5c4eceb6cc49cfe39447d4bfc56e6c865e7c2249b6eb978935f/pywinpty-3.0.2.tar.gz", hash = "sha256:1505cc4cb248af42cb6285a65c9c2086ee9e7e574078ee60933d5d7fa86fb004", size = 30669, upload-time = "2025-10-03T21:16:29.205Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a6/a1/409c1651c9f874d598c10f51ff586c416625601df4bca315d08baec4c3e3/pywinpty-3.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:327790d70e4c841ebd9d0f295a780177149aeb405bca44c7115a3de5c2054b23", size = 2050304, upload-time = "2025-10-03T21:19:29.466Z" },
-    { url = "https://files.pythonhosted.org/packages/02/4e/1098484e042c9485f56f16eb2b69b43b874bd526044ee401512234cf9e04/pywinpty-3.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:99fdd9b455f0ad6419aba6731a7a0d2f88ced83c3c94a80ff9533d95fa8d8a9e", size = 2050391, upload-time = "2025-10-03T21:19:01.642Z" },
-    { url = "https://files.pythonhosted.org/packages/fc/19/b757fe28008236a4a713e813283721b8a40aa60cd7d3f83549f2e25a3155/pywinpty-3.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:18f78b81e4cfee6aabe7ea8688441d30247b73e52cd9657138015c5f4ee13a51", size = 2050057, upload-time = "2025-10-03T21:19:26.732Z" },
-    { url = "https://files.pythonhosted.org/packages/cb/44/cbae12ecf6f4fa4129c36871fd09c6bef4f98d5f625ecefb5e2449765508/pywinpty-3.0.2-cp313-cp313t-win_amd64.whl", hash = "sha256:663383ecfab7fc382cc97ea5c4f7f0bb32c2f889259855df6ea34e5df42d305b", size = 2049874, upload-time = "2025-10-03T21:18:53.923Z" },
-]
-
 [[package]]
 name = "pyyaml"
 version = "6.0.2"
@@ -3253,51 +2229,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446, upload-time = "2024-08-06T20:33:04.33Z" },
 ]
 
-[[package]]
-name = "pyzmq"
-version = "27.0.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "cffi", marker = "implementation_name == 'pypy'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/f1/06/50a4e9648b3e8b992bef8eb632e457307553a89d294103213cfd47b3da69/pyzmq-27.0.0.tar.gz", hash = "sha256:b1f08eeb9ce1510e6939b6e5dcd46a17765e2333daae78ecf4606808442e52cf", size = 280478, upload-time = "2025-06-13T14:09:07.087Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/44/df/84c630654106d9bd9339cdb564aa941ed41b023a0264251d6743766bb50e/pyzmq-27.0.0-cp311-cp311-macosx_10_15_universal2.whl", hash = "sha256:21457825249b2a53834fa969c69713f8b5a79583689387a5e7aed880963ac564", size = 1332718, upload-time = "2025-06-13T14:07:16.555Z" },
-    { url = "https://files.pythonhosted.org/packages/c1/8e/f6a5461a07654d9840d256476434ae0ff08340bba562a455f231969772cb/pyzmq-27.0.0-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:1958947983fef513e6e98eff9cb487b60bf14f588dc0e6bf35fa13751d2c8251", size = 908248, upload-time = "2025-06-13T14:07:18.033Z" },
-    { url = "https://files.pythonhosted.org/packages/7c/93/82863e8d695a9a3ae424b63662733ae204a295a2627d52af2f62c2cd8af9/pyzmq-27.0.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0dc628b5493f9a8cd9844b8bee9732ef587ab00002157c9329e4fc0ef4d3afa", size = 668647, upload-time = "2025-06-13T14:07:19.378Z" },
-    { url = "https://files.pythonhosted.org/packages/f3/85/15278769b348121eacdbfcbd8c4d40f1102f32fa6af5be1ffc032ed684be/pyzmq-27.0.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f7bbe9e1ed2c8d3da736a15694d87c12493e54cc9dc9790796f0321794bbc91f", size = 856600, upload-time = "2025-06-13T14:07:20.906Z" },
-    { url = "https://files.pythonhosted.org/packages/d4/af/1c469b3d479bd095edb28e27f12eee10b8f00b356acbefa6aeb14dd295d1/pyzmq-27.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:dc1091f59143b471d19eb64f54bae4f54bcf2a466ffb66fe45d94d8d734eb495", size = 1657748, upload-time = "2025-06-13T14:07:22.549Z" },
-    { url = "https://files.pythonhosted.org/packages/8c/f4/17f965d0ee6380b1d6326da842a50e4b8b9699745161207945f3745e8cb5/pyzmq-27.0.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:7011ade88c8e535cf140f8d1a59428676fbbce7c6e54fefce58bf117aefb6667", size = 2034311, upload-time = "2025-06-13T14:07:23.966Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/6e/7c391d81fa3149fd759de45d298003de6cfab343fb03e92c099821c448db/pyzmq-27.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:2c386339d7e3f064213aede5d03d054b237937fbca6dd2197ac8cf3b25a6b14e", size = 1893630, upload-time = "2025-06-13T14:07:25.899Z" },
-    { url = "https://files.pythonhosted.org/packages/0e/e0/eaffe7a86f60e556399e224229e7769b717f72fec0706b70ab2c03aa04cb/pyzmq-27.0.0-cp311-cp311-win32.whl", hash = "sha256:0546a720c1f407b2172cb04b6b094a78773491497e3644863cf5c96c42df8cff", size = 567706, upload-time = "2025-06-13T14:07:27.595Z" },
-    { url = "https://files.pythonhosted.org/packages/c9/05/89354a8cffdcce6e547d48adaaf7be17007fc75572123ff4ca90a4ca04fc/pyzmq-27.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:15f39d50bd6c9091c67315ceb878a4f531957b121d2a05ebd077eb35ddc5efed", size = 630322, upload-time = "2025-06-13T14:07:28.938Z" },
-    { url = "https://files.pythonhosted.org/packages/fa/07/4ab976d5e1e63976719389cc4f3bfd248a7f5f2bb2ebe727542363c61b5f/pyzmq-27.0.0-cp311-cp311-win_arm64.whl", hash = "sha256:c5817641eebb391a2268c27fecd4162448e03538387093cdbd8bf3510c316b38", size = 558435, upload-time = "2025-06-13T14:07:30.256Z" },
-    { url = "https://files.pythonhosted.org/packages/93/a7/9ad68f55b8834ede477842214feba6a4c786d936c022a67625497aacf61d/pyzmq-27.0.0-cp312-abi3-macosx_10_15_universal2.whl", hash = "sha256:cbabc59dcfaac66655c040dfcb8118f133fb5dde185e5fc152628354c1598e52", size = 1305438, upload-time = "2025-06-13T14:07:31.676Z" },
-    { url = "https://files.pythonhosted.org/packages/ba/ee/26aa0f98665a22bc90ebe12dced1de5f3eaca05363b717f6fb229b3421b3/pyzmq-27.0.0-cp312-abi3-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:cb0ac5179cba4b2f94f1aa208fbb77b62c4c9bf24dd446278b8b602cf85fcda3", size = 895095, upload-time = "2025-06-13T14:07:33.104Z" },
-    { url = "https://files.pythonhosted.org/packages/cf/85/c57e7ab216ecd8aa4cc7e3b83b06cc4e9cf45c87b0afc095f10cd5ce87c1/pyzmq-27.0.0-cp312-abi3-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:53a48f0228eab6cbf69fde3aa3c03cbe04e50e623ef92ae395fce47ef8a76152", size = 651826, upload-time = "2025-06-13T14:07:34.831Z" },
-    { url = "https://files.pythonhosted.org/packages/69/9a/9ea7e230feda9400fb0ae0d61d7d6ddda635e718d941c44eeab22a179d34/pyzmq-27.0.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:111db5f395e09f7e775f759d598f43cb815fc58e0147623c4816486e1a39dc22", size = 839750, upload-time = "2025-06-13T14:07:36.553Z" },
-    { url = "https://files.pythonhosted.org/packages/08/66/4cebfbe71f3dfbd417011daca267539f62ed0fbc68105357b68bbb1a25b7/pyzmq-27.0.0-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c8878011653dcdc27cc2c57e04ff96f0471e797f5c19ac3d7813a245bcb24371", size = 1641357, upload-time = "2025-06-13T14:07:38.21Z" },
-    { url = "https://files.pythonhosted.org/packages/ac/f6/b0f62578c08d2471c791287149cb8c2aaea414ae98c6e995c7dbe008adfb/pyzmq-27.0.0-cp312-abi3-musllinux_1_2_i686.whl", hash = "sha256:c0ed2c1f335ba55b5fdc964622254917d6b782311c50e138863eda409fbb3b6d", size = 2020281, upload-time = "2025-06-13T14:07:39.599Z" },
-    { url = "https://files.pythonhosted.org/packages/37/b9/4f670b15c7498495da9159edc374ec09c88a86d9cd5a47d892f69df23450/pyzmq-27.0.0-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e918d70862d4cfd4b1c187310015646a14e1f5917922ab45b29f28f345eeb6be", size = 1877110, upload-time = "2025-06-13T14:07:41.027Z" },
-    { url = "https://files.pythonhosted.org/packages/66/31/9dee25c226295b740609f0d46db2fe972b23b6f5cf786360980524a3ba92/pyzmq-27.0.0-cp312-abi3-win32.whl", hash = "sha256:88b4e43cab04c3c0f0d55df3b1eef62df2b629a1a369b5289a58f6fa8b07c4f4", size = 559297, upload-time = "2025-06-13T14:07:42.533Z" },
-    { url = "https://files.pythonhosted.org/packages/9b/12/52da5509800f7ff2d287b2f2b4e636e7ea0f001181cba6964ff6c1537778/pyzmq-27.0.0-cp312-abi3-win_amd64.whl", hash = "sha256:dce4199bf5f648a902ce37e7b3afa286f305cd2ef7a8b6ec907470ccb6c8b371", size = 619203, upload-time = "2025-06-13T14:07:43.843Z" },
-    { url = "https://files.pythonhosted.org/packages/93/6d/7f2e53b19d1edb1eb4f09ec7c3a1f945ca0aac272099eab757d15699202b/pyzmq-27.0.0-cp312-abi3-win_arm64.whl", hash = "sha256:56e46bbb85d52c1072b3f809cc1ce77251d560bc036d3a312b96db1afe76db2e", size = 551927, upload-time = "2025-06-13T14:07:45.51Z" },
-    { url = "https://files.pythonhosted.org/packages/19/62/876b27c4ff777db4ceba1c69ea90d3c825bb4f8d5e7cd987ce5802e33c55/pyzmq-27.0.0-cp313-cp313t-macosx_10_15_universal2.whl", hash = "sha256:c36ad534c0c29b4afa088dc53543c525b23c0797e01b69fef59b1a9c0e38b688", size = 1340826, upload-time = "2025-06-13T14:07:46.881Z" },
-    { url = "https://files.pythonhosted.org/packages/43/69/58ef8f4f59d3bcd505260c73bee87b008850f45edca40ddaba54273c35f4/pyzmq-27.0.0-cp313-cp313t-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:67855c14173aec36395d7777aaba3cc527b393821f30143fd20b98e1ff31fd38", size = 897283, upload-time = "2025-06-13T14:07:49.562Z" },
-    { url = "https://files.pythonhosted.org/packages/43/15/93a0d0396700a60475ad3c5d42c5f1c308d3570bc94626b86c71ef9953e0/pyzmq-27.0.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8617c7d43cd8ccdb62aebe984bfed77ca8f036e6c3e46dd3dddda64b10f0ab7a", size = 660567, upload-time = "2025-06-13T14:07:51.364Z" },
-    { url = "https://files.pythonhosted.org/packages/0e/b3/fe055513e498ca32f64509abae19b9c9eb4d7c829e02bd8997dd51b029eb/pyzmq-27.0.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:67bfbcbd0a04c575e8103a6061d03e393d9f80ffdb9beb3189261e9e9bc5d5e9", size = 847681, upload-time = "2025-06-13T14:07:52.77Z" },
-    { url = "https://files.pythonhosted.org/packages/b6/4f/ff15300b00b5b602191f3df06bbc8dd4164e805fdd65bb77ffbb9c5facdc/pyzmq-27.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:5cd11d46d7b7e5958121b3eaf4cd8638eff3a720ec527692132f05a57f14341d", size = 1650148, upload-time = "2025-06-13T14:07:54.178Z" },
-    { url = "https://files.pythonhosted.org/packages/c4/6f/84bdfff2a224a6f26a24249a342e5906993c50b0761e311e81b39aef52a7/pyzmq-27.0.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:b801c2e40c5aa6072c2f4876de8dccd100af6d9918d4d0d7aa54a1d982fd4f44", size = 2023768, upload-time = "2025-06-13T14:07:55.714Z" },
-    { url = "https://files.pythonhosted.org/packages/64/39/dc2db178c26a42228c5ac94a9cc595030458aa64c8d796a7727947afbf55/pyzmq-27.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:20d5cb29e8c5f76a127c75b6e7a77e846bc4b655c373baa098c26a61b7ecd0ef", size = 1885199, upload-time = "2025-06-13T14:07:57.166Z" },
-    { url = "https://files.pythonhosted.org/packages/c7/21/dae7b06a1f8cdee5d8e7a63d99c5d129c401acc40410bef2cbf42025e26f/pyzmq-27.0.0-cp313-cp313t-win32.whl", hash = "sha256:a20528da85c7ac7a19b7384e8c3f8fa707841fd85afc4ed56eda59d93e3d98ad", size = 575439, upload-time = "2025-06-13T14:07:58.959Z" },
-    { url = "https://files.pythonhosted.org/packages/eb/bc/1709dc55f0970cf4cb8259e435e6773f9946f41a045c2cb90e870b7072da/pyzmq-27.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:d8229f2efece6a660ee211d74d91dbc2a76b95544d46c74c615e491900dc107f", size = 639933, upload-time = "2025-06-13T14:08:00.777Z" },
-    { url = "https://files.pythonhosted.org/packages/98/a6/92394373b8dbc1edc9d53c951e8d3989d518185174ee54492ec27711779d/pyzmq-27.0.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:cd1dc59763effd1576f8368047c9c31468fce0af89d76b5067641137506792ae", size = 835948, upload-time = "2025-06-13T14:08:43.516Z" },
-    { url = "https://files.pythonhosted.org/packages/56/f3/4dc38d75d9995bfc18773df3e41f2a2ca9b740b06f1a15dbf404077e7588/pyzmq-27.0.0-pp311-pypy311_pp73-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:60e8cc82d968174650c1860d7b716366caab9973787a1c060cf8043130f7d0f7", size = 799874, upload-time = "2025-06-13T14:08:45.017Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/ba/64af397e0f421453dc68e31d5e0784d554bf39013a2de0872056e96e58af/pyzmq-27.0.0-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:14fe7aaac86e4e93ea779a821967360c781d7ac5115b3f1a171ced77065a0174", size = 567400, upload-time = "2025-06-13T14:08:46.855Z" },
-    { url = "https://files.pythonhosted.org/packages/63/87/ec956cbe98809270b59a22891d5758edae147a258e658bf3024a8254c855/pyzmq-27.0.0-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6ad0562d4e6abb785be3e4dd68599c41be821b521da38c402bc9ab2a8e7ebc7e", size = 747031, upload-time = "2025-06-13T14:08:48.419Z" },
-    { url = "https://files.pythonhosted.org/packages/be/8a/4a3764a68abc02e2fbb0668d225b6fda5cd39586dd099cee8b2ed6ab0452/pyzmq-27.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:9df43a2459cd3a3563404c1456b2c4c69564daa7dbaf15724c09821a3329ce46", size = 544726, upload-time = "2025-06-13T14:08:49.903Z" },
-]
-
 [[package]]
 name = "readme-renderer"
 version = "44.0"
@@ -3312,20 +2243,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl", hash = "sha256:2fbca89b81a08526aadf1357a8c2ae889ec05fb03f5da67f9769c9a592166151", size = 13310, upload-time = "2024-07-08T15:00:56.577Z" },
 ]
 
-[[package]]
-name = "referencing"
-version = "0.36.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "attrs" },
-    { name = "rpds-py" },
-    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/2f/db/98b5c277be99dd18bfd91dd04e1b759cad18d1a338188c936e92f921c7e2/referencing-0.36.2.tar.gz", hash = "sha256:df2e89862cd09deabbdba16944cc3f10feb6b3e6f18e902f7cc25609a34775aa", size = 74744, upload-time = "2025-01-25T08:48:16.138Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c1/b1/3baf80dc6d2b7bc27a95a67752d0208e410351e3feb4eb78de5f77454d8d/referencing-0.36.2-py3-none-any.whl", hash = "sha256:e8699adbbf8b5c7de96d8ffa0eb5c158b3beafce084968e2ea8bb08c6794dcd0", size = 26775, upload-time = "2025-01-25T08:48:14.241Z" },
-]
-
 [[package]]
 name = "regex"
 version = "2024.11.6"
@@ -3394,23 +2311,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" },
 ]
 
-[[package]]
-name = "requests-cache"
-version = "1.2.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "attrs" },
-    { name = "cattrs" },
-    { name = "platformdirs" },
-    { name = "requests" },
-    { name = "url-normalize" },
-    { name = "urllib3" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/1a/be/7b2a95a9e7a7c3e774e43d067c51244e61dea8b120ae2deff7089a93fb2b/requests_cache-1.2.1.tar.gz", hash = "sha256:68abc986fdc5b8d0911318fbb5f7c80eebcd4d01bfacc6685ecf8876052511d1", size = 3018209, upload-time = "2024-06-18T17:18:03.774Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/4e/2e/8f4051119f460cfc786aa91f212165bb6e643283b533db572d7b33952bd2/requests_cache-1.2.1-py3-none-any.whl", hash = "sha256:1285151cddf5331067baa82598afe2d47c7495a1334bfe7a7d329b43e9fd3603", size = 61425, upload-time = "2024-06-18T17:17:45Z" },
-]
-
 [[package]]
 name = "requests-toolbelt"
 version = "1.0.0"
@@ -3423,18 +2323,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/3f/51/d4db610ef29373b879047326cbf6fa98b6c1969d6f6dc423279de2b1be2c/requests_toolbelt-1.0.0-py2.py3-none-any.whl", hash = "sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06", size = 54481, upload-time = "2023-05-01T04:11:28.427Z" },
 ]
 
-[[package]]
-name = "rfc3339-validator"
-version = "0.1.4"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "six" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/28/ea/a9387748e2d111c3c2b275ba970b735e04e15cdb1eb30693b6b5708c4dbd/rfc3339_validator-0.1.4.tar.gz", hash = "sha256:138a2abdf93304ad60530167e51d2dfb9549521a836871b88d7f4695d0022f6b", size = 5513, upload-time = "2021-05-12T16:37:54.178Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7b/44/4e421b96b67b2daff264473f7465db72fbdf36a07e05494f50300cc7b0c6/rfc3339_validator-0.1.4-py2.py3-none-any.whl", hash = "sha256:24f6ec1eda14ef823da9e36ec7113124b39c04d50a4d3d3a3c2859577e7791fa", size = 3490, upload-time = "2021-05-12T16:37:52.536Z" },
-]
-
 [[package]]
 name = "rfc3986"
 version = "2.0.0"
@@ -3444,27 +2332,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl", hash = "sha256:50b1502b60e289cb37883f3dfd34532b8873c7de9f49bb546641ce9cbd256ebd", size = 31326, upload-time = "2022-01-10T00:52:29.594Z" },
 ]
 
-[[package]]
-name = "rfc3986-validator"
-version = "0.1.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/da/88/f270de456dd7d11dcc808abfa291ecdd3f45ff44e3b549ffa01b126464d0/rfc3986_validator-0.1.1.tar.gz", hash = "sha256:3d44bde7921b3b9ec3ae4e3adca370438eccebc676456449b145d533b240d055", size = 6760, upload-time = "2019-10-28T16:00:19.144Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/9e/51/17023c0f8f1869d8806b979a2bffa3f861f26a3f1a66b094288323fba52f/rfc3986_validator-0.1.1-py2.py3-none-any.whl", hash = "sha256:2f235c432ef459970b4306369336b9d5dbdda31b510ca1e327636e01f528bfa9", size = 4242, upload-time = "2019-10-28T16:00:13.976Z" },
-]
-
-[[package]]
-name = "rfc3987-syntax"
-version = "1.1.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "lark" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/2c/06/37c1a5557acf449e8e406a830a05bf885ac47d33270aec454ef78675008d/rfc3987_syntax-1.1.0.tar.gz", hash = "sha256:717a62cbf33cffdd16dfa3a497d81ce48a660ea691b1ddd7be710c22f00b4a0d", size = 14239, upload-time = "2025-07-18T01:05:05.015Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7e/71/44ce230e1b7fadd372515a97e32a83011f906ddded8d03e3c6aafbdedbb7/rfc3987_syntax-1.1.0-py3-none-any.whl", hash = "sha256:6c3d97604e4c5ce9f714898e05401a0445a641cfa276432b0a648c80856f6a3f", size = 8046, upload-time = "2025-07-18T01:05:03.843Z" },
-]
-
 [[package]]
 name = "rich"
 version = "14.1.0"
@@ -3479,89 +2346,28 @@ wheels = [
 ]
 
 [[package]]
-name = "rpds-py"
-version = "0.26.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/a5/aa/4456d84bbb54adc6a916fb10c9b374f78ac840337644e4a5eda229c81275/rpds_py-0.26.0.tar.gz", hash = "sha256:20dae58a859b0906f0685642e591056f1e787f3a8b39c8e8749a45dc7d26bdb0", size = 27385, upload-time = "2025-07-01T15:57:13.958Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/09/4c/4ee8f7e512030ff79fda1df3243c88d70fc874634e2dbe5df13ba4210078/rpds_py-0.26.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:9e8cb77286025bdb21be2941d64ac6ca016130bfdcd228739e8ab137eb4406ed", size = 372610, upload-time = "2025-07-01T15:53:58.844Z" },
-    { url = "https://files.pythonhosted.org/packages/fa/9d/3dc16be00f14fc1f03c71b1d67c8df98263ab2710a2fbd65a6193214a527/rpds_py-0.26.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5e09330b21d98adc8ccb2dbb9fc6cb434e8908d4c119aeaa772cb1caab5440a0", size = 358032, upload-time = "2025-07-01T15:53:59.985Z" },
-    { url = "https://files.pythonhosted.org/packages/e7/5a/7f1bf8f045da2866324a08ae80af63e64e7bfaf83bd31f865a7b91a58601/rpds_py-0.26.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2c9c1b92b774b2e68d11193dc39620d62fd8ab33f0a3c77ecdabe19c179cdbc1", size = 381525, upload-time = "2025-07-01T15:54:01.162Z" },
-    { url = "https://files.pythonhosted.org/packages/45/8a/04479398c755a066ace10e3d158866beb600867cacae194c50ffa783abd0/rpds_py-0.26.0-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:824e6d3503ab990d7090768e4dfd9e840837bae057f212ff9f4f05ec6d1975e7", size = 397089, upload-time = "2025-07-01T15:54:02.319Z" },
-    { url = "https://files.pythonhosted.org/packages/72/88/9203f47268db488a1b6d469d69c12201ede776bb728b9d9f29dbfd7df406/rpds_py-0.26.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8ad7fd2258228bf288f2331f0a6148ad0186b2e3643055ed0db30990e59817a6", size = 514255, upload-time = "2025-07-01T15:54:03.38Z" },
-    { url = "https://files.pythonhosted.org/packages/f5/b4/01ce5d1e853ddf81fbbd4311ab1eff0b3cf162d559288d10fd127e2588b5/rpds_py-0.26.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0dc23bbb3e06ec1ea72d515fb572c1fea59695aefbffb106501138762e1e915e", size = 402283, upload-time = "2025-07-01T15:54:04.923Z" },
-    { url = "https://files.pythonhosted.org/packages/34/a2/004c99936997bfc644d590a9defd9e9c93f8286568f9c16cdaf3e14429a7/rpds_py-0.26.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d80bf832ac7b1920ee29a426cdca335f96a2b5caa839811803e999b41ba9030d", size = 383881, upload-time = "2025-07-01T15:54:06.482Z" },
-    { url = "https://files.pythonhosted.org/packages/05/1b/ef5fba4a8f81ce04c427bfd96223f92f05e6cd72291ce9d7523db3b03a6c/rpds_py-0.26.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:0919f38f5542c0a87e7b4afcafab6fd2c15386632d249e9a087498571250abe3", size = 415822, upload-time = "2025-07-01T15:54:07.605Z" },
-    { url = "https://files.pythonhosted.org/packages/16/80/5c54195aec456b292f7bd8aa61741c8232964063fd8a75fdde9c1e982328/rpds_py-0.26.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d422b945683e409000c888e384546dbab9009bb92f7c0b456e217988cf316107", size = 558347, upload-time = "2025-07-01T15:54:08.591Z" },
-    { url = "https://files.pythonhosted.org/packages/f2/1c/1845c1b1fd6d827187c43afe1841d91678d7241cbdb5420a4c6de180a538/rpds_py-0.26.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:77a7711fa562ba2da1aa757e11024ad6d93bad6ad7ede5afb9af144623e5f76a", size = 587956, upload-time = "2025-07-01T15:54:09.963Z" },
-    { url = "https://files.pythonhosted.org/packages/2e/ff/9e979329dd131aa73a438c077252ddabd7df6d1a7ad7b9aacf6261f10faa/rpds_py-0.26.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:238e8c8610cb7c29460e37184f6799547f7e09e6a9bdbdab4e8edb90986a2318", size = 554363, upload-time = "2025-07-01T15:54:11.073Z" },
-    { url = "https://files.pythonhosted.org/packages/00/8b/d78cfe034b71ffbe72873a136e71acc7a831a03e37771cfe59f33f6de8a2/rpds_py-0.26.0-cp311-cp311-win32.whl", hash = "sha256:893b022bfbdf26d7bedb083efeea624e8550ca6eb98bf7fea30211ce95b9201a", size = 220123, upload-time = "2025-07-01T15:54:12.382Z" },
-    { url = "https://files.pythonhosted.org/packages/94/c1/3c8c94c7dd3905dbfde768381ce98778500a80db9924731d87ddcdb117e9/rpds_py-0.26.0-cp311-cp311-win_amd64.whl", hash = "sha256:87a5531de9f71aceb8af041d72fc4cab4943648d91875ed56d2e629bef6d4c03", size = 231732, upload-time = "2025-07-01T15:54:13.434Z" },
-    { url = "https://files.pythonhosted.org/packages/67/93/e936fbed1b734eabf36ccb5d93c6a2e9246fbb13c1da011624b7286fae3e/rpds_py-0.26.0-cp311-cp311-win_arm64.whl", hash = "sha256:de2713f48c1ad57f89ac25b3cb7daed2156d8e822cf0eca9b96a6f990718cc41", size = 221917, upload-time = "2025-07-01T15:54:14.559Z" },
-    { url = "https://files.pythonhosted.org/packages/ea/86/90eb87c6f87085868bd077c7a9938006eb1ce19ed4d06944a90d3560fce2/rpds_py-0.26.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:894514d47e012e794f1350f076c427d2347ebf82f9b958d554d12819849a369d", size = 363933, upload-time = "2025-07-01T15:54:15.734Z" },
-    { url = "https://files.pythonhosted.org/packages/63/78/4469f24d34636242c924626082b9586f064ada0b5dbb1e9d096ee7a8e0c6/rpds_py-0.26.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc921b96fa95a097add244da36a1d9e4f3039160d1d30f1b35837bf108c21136", size = 350447, upload-time = "2025-07-01T15:54:16.922Z" },
-    { url = "https://files.pythonhosted.org/packages/ad/91/c448ed45efdfdade82348d5e7995e15612754826ea640afc20915119734f/rpds_py-0.26.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3e1157659470aa42a75448b6e943c895be8c70531c43cb78b9ba990778955582", size = 384711, upload-time = "2025-07-01T15:54:18.101Z" },
-    { url = "https://files.pythonhosted.org/packages/ec/43/e5c86fef4be7f49828bdd4ecc8931f0287b1152c0bb0163049b3218740e7/rpds_py-0.26.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:521ccf56f45bb3a791182dc6b88ae5f8fa079dd705ee42138c76deb1238e554e", size = 400865, upload-time = "2025-07-01T15:54:19.295Z" },
-    { url = "https://files.pythonhosted.org/packages/55/34/e00f726a4d44f22d5c5fe2e5ddd3ac3d7fd3f74a175607781fbdd06fe375/rpds_py-0.26.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9def736773fd56b305c0eef698be5192c77bfa30d55a0e5885f80126c4831a15", size = 517763, upload-time = "2025-07-01T15:54:20.858Z" },
-    { url = "https://files.pythonhosted.org/packages/52/1c/52dc20c31b147af724b16104500fba13e60123ea0334beba7b40e33354b4/rpds_py-0.26.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cdad4ea3b4513b475e027be79e5a0ceac8ee1c113a1a11e5edc3c30c29f964d8", size = 406651, upload-time = "2025-07-01T15:54:22.508Z" },
-    { url = "https://files.pythonhosted.org/packages/2e/77/87d7bfabfc4e821caa35481a2ff6ae0b73e6a391bb6b343db2c91c2b9844/rpds_py-0.26.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:82b165b07f416bdccf5c84546a484cc8f15137ca38325403864bfdf2b5b72f6a", size = 386079, upload-time = "2025-07-01T15:54:23.987Z" },
-    { url = "https://files.pythonhosted.org/packages/e3/d4/7f2200c2d3ee145b65b3cddc4310d51f7da6a26634f3ac87125fd789152a/rpds_py-0.26.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d04cab0a54b9dba4d278fe955a1390da3cf71f57feb78ddc7cb67cbe0bd30323", size = 421379, upload-time = "2025-07-01T15:54:25.073Z" },
-    { url = "https://files.pythonhosted.org/packages/ae/13/9fdd428b9c820869924ab62236b8688b122baa22d23efdd1c566938a39ba/rpds_py-0.26.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:79061ba1a11b6a12743a2b0f72a46aa2758613d454aa6ba4f5a265cc48850158", size = 562033, upload-time = "2025-07-01T15:54:26.225Z" },
-    { url = "https://files.pythonhosted.org/packages/f3/e1/b69686c3bcbe775abac3a4c1c30a164a2076d28df7926041f6c0eb5e8d28/rpds_py-0.26.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:f405c93675d8d4c5ac87364bb38d06c988e11028a64b52a47158a355079661f3", size = 591639, upload-time = "2025-07-01T15:54:27.424Z" },
-    { url = "https://files.pythonhosted.org/packages/5c/c9/1e3d8c8863c84a90197ac577bbc3d796a92502124c27092413426f670990/rpds_py-0.26.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:dafd4c44b74aa4bed4b250f1aed165b8ef5de743bcca3b88fc9619b6087093d2", size = 557105, upload-time = "2025-07-01T15:54:29.93Z" },
-    { url = "https://files.pythonhosted.org/packages/9f/c5/90c569649057622959f6dcc40f7b516539608a414dfd54b8d77e3b201ac0/rpds_py-0.26.0-cp312-cp312-win32.whl", hash = "sha256:3da5852aad63fa0c6f836f3359647870e21ea96cf433eb393ffa45263a170d44", size = 223272, upload-time = "2025-07-01T15:54:31.128Z" },
-    { url = "https://files.pythonhosted.org/packages/7d/16/19f5d9f2a556cfed454eebe4d354c38d51c20f3db69e7b4ce6cff904905d/rpds_py-0.26.0-cp312-cp312-win_amd64.whl", hash = "sha256:cf47cfdabc2194a669dcf7a8dbba62e37a04c5041d2125fae0233b720da6f05c", size = 234995, upload-time = "2025-07-01T15:54:32.195Z" },
-    { url = "https://files.pythonhosted.org/packages/83/f0/7935e40b529c0e752dfaa7880224771b51175fce08b41ab4a92eb2fbdc7f/rpds_py-0.26.0-cp312-cp312-win_arm64.whl", hash = "sha256:20ab1ae4fa534f73647aad289003f1104092890849e0266271351922ed5574f8", size = 223198, upload-time = "2025-07-01T15:54:33.271Z" },
-    { url = "https://files.pythonhosted.org/packages/6a/67/bb62d0109493b12b1c6ab00de7a5566aa84c0e44217c2d94bee1bd370da9/rpds_py-0.26.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:696764a5be111b036256c0b18cd29783fab22154690fc698062fc1b0084b511d", size = 363917, upload-time = "2025-07-01T15:54:34.755Z" },
-    { url = "https://files.pythonhosted.org/packages/4b/f3/34e6ae1925a5706c0f002a8d2d7f172373b855768149796af87bd65dcdb9/rpds_py-0.26.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1e6c15d2080a63aaed876e228efe4f814bc7889c63b1e112ad46fdc8b368b9e1", size = 350073, upload-time = "2025-07-01T15:54:36.292Z" },
-    { url = "https://files.pythonhosted.org/packages/75/83/1953a9d4f4e4de7fd0533733e041c28135f3c21485faaef56a8aadbd96b5/rpds_py-0.26.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:390e3170babf42462739a93321e657444f0862c6d722a291accc46f9d21ed04e", size = 384214, upload-time = "2025-07-01T15:54:37.469Z" },
-    { url = "https://files.pythonhosted.org/packages/48/0e/983ed1b792b3322ea1d065e67f4b230f3b96025f5ce3878cc40af09b7533/rpds_py-0.26.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7da84c2c74c0f5bc97d853d9e17bb83e2dcafcff0dc48286916001cc114379a1", size = 400113, upload-time = "2025-07-01T15:54:38.954Z" },
-    { url = "https://files.pythonhosted.org/packages/69/7f/36c0925fff6f660a80be259c5b4f5e53a16851f946eb080351d057698528/rpds_py-0.26.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4c5fe114a6dd480a510b6d3661d09d67d1622c4bf20660a474507aaee7eeeee9", size = 515189, upload-time = "2025-07-01T15:54:40.57Z" },
-    { url = "https://files.pythonhosted.org/packages/13/45/cbf07fc03ba7a9b54662c9badb58294ecfb24f828b9732970bd1a431ed5c/rpds_py-0.26.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3100b3090269f3a7ea727b06a6080d4eb7439dca4c0e91a07c5d133bb1727ea7", size = 406998, upload-time = "2025-07-01T15:54:43.025Z" },
-    { url = "https://files.pythonhosted.org/packages/6c/b0/8fa5e36e58657997873fd6a1cf621285ca822ca75b4b3434ead047daa307/rpds_py-0.26.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2c03c9b0c64afd0320ae57de4c982801271c0c211aa2d37f3003ff5feb75bb04", size = 385903, upload-time = "2025-07-01T15:54:44.752Z" },
-    { url = "https://files.pythonhosted.org/packages/4b/f7/b25437772f9f57d7a9fbd73ed86d0dcd76b4c7c6998348c070d90f23e315/rpds_py-0.26.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5963b72ccd199ade6ee493723d18a3f21ba7d5b957017607f815788cef50eaf1", size = 419785, upload-time = "2025-07-01T15:54:46.043Z" },
-    { url = "https://files.pythonhosted.org/packages/a7/6b/63ffa55743dfcb4baf2e9e77a0b11f7f97ed96a54558fcb5717a4b2cd732/rpds_py-0.26.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:9da4e873860ad5bab3291438525cae80169daecbfafe5657f7f5fb4d6b3f96b9", size = 561329, upload-time = "2025-07-01T15:54:47.64Z" },
-    { url = "https://files.pythonhosted.org/packages/2f/07/1f4f5e2886c480a2346b1e6759c00278b8a69e697ae952d82ae2e6ee5db0/rpds_py-0.26.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:5afaddaa8e8c7f1f7b4c5c725c0070b6eed0228f705b90a1732a48e84350f4e9", size = 590875, upload-time = "2025-07-01T15:54:48.9Z" },
-    { url = "https://files.pythonhosted.org/packages/cc/bc/e6639f1b91c3a55f8c41b47d73e6307051b6e246254a827ede730624c0f8/rpds_py-0.26.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4916dc96489616a6f9667e7526af8fa693c0fdb4f3acb0e5d9f4400eb06a47ba", size = 556636, upload-time = "2025-07-01T15:54:50.619Z" },
-    { url = "https://files.pythonhosted.org/packages/05/4c/b3917c45566f9f9a209d38d9b54a1833f2bb1032a3e04c66f75726f28876/rpds_py-0.26.0-cp313-cp313-win32.whl", hash = "sha256:2a343f91b17097c546b93f7999976fd6c9d5900617aa848c81d794e062ab302b", size = 222663, upload-time = "2025-07-01T15:54:52.023Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/0b/0851bdd6025775aaa2365bb8de0697ee2558184c800bfef8d7aef5ccde58/rpds_py-0.26.0-cp313-cp313-win_amd64.whl", hash = "sha256:0a0b60701f2300c81b2ac88a5fb893ccfa408e1c4a555a77f908a2596eb875a5", size = 234428, upload-time = "2025-07-01T15:54:53.692Z" },
-    { url = "https://files.pythonhosted.org/packages/ed/e8/a47c64ed53149c75fb581e14a237b7b7cd18217e969c30d474d335105622/rpds_py-0.26.0-cp313-cp313-win_arm64.whl", hash = "sha256:257d011919f133a4746958257f2c75238e3ff54255acd5e3e11f3ff41fd14256", size = 222571, upload-time = "2025-07-01T15:54:54.822Z" },
-    { url = "https://files.pythonhosted.org/packages/89/bf/3d970ba2e2bcd17d2912cb42874107390f72873e38e79267224110de5e61/rpds_py-0.26.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:529c8156d7506fba5740e05da8795688f87119cce330c244519cf706a4a3d618", size = 360475, upload-time = "2025-07-01T15:54:56.228Z" },
-    { url = "https://files.pythonhosted.org/packages/82/9f/283e7e2979fc4ec2d8ecee506d5a3675fce5ed9b4b7cb387ea5d37c2f18d/rpds_py-0.26.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:f53ec51f9d24e9638a40cabb95078ade8c99251945dad8d57bf4aabe86ecee35", size = 346692, upload-time = "2025-07-01T15:54:58.561Z" },
-    { url = "https://files.pythonhosted.org/packages/e3/03/7e50423c04d78daf391da3cc4330bdb97042fc192a58b186f2d5deb7befd/rpds_py-0.26.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7ab504c4d654e4a29558eaa5bb8cea5fdc1703ea60a8099ffd9c758472cf913f", size = 379415, upload-time = "2025-07-01T15:54:59.751Z" },
-    { url = "https://files.pythonhosted.org/packages/57/00/d11ee60d4d3b16808432417951c63df803afb0e0fc672b5e8d07e9edaaae/rpds_py-0.26.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:fd0641abca296bc1a00183fe44f7fced8807ed49d501f188faa642d0e4975b83", size = 391783, upload-time = "2025-07-01T15:55:00.898Z" },
-    { url = "https://files.pythonhosted.org/packages/08/b3/1069c394d9c0d6d23c5b522e1f6546b65793a22950f6e0210adcc6f97c3e/rpds_py-0.26.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:69b312fecc1d017b5327afa81d4da1480f51c68810963a7336d92203dbb3d4f1", size = 512844, upload-time = "2025-07-01T15:55:02.201Z" },
-    { url = "https://files.pythonhosted.org/packages/08/3b/c4fbf0926800ed70b2c245ceca99c49f066456755f5d6eb8863c2c51e6d0/rpds_py-0.26.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c741107203954f6fc34d3066d213d0a0c40f7bb5aafd698fb39888af277c70d8", size = 402105, upload-time = "2025-07-01T15:55:03.698Z" },
-    { url = "https://files.pythonhosted.org/packages/1c/b0/db69b52ca07413e568dae9dc674627a22297abb144c4d6022c6d78f1e5cc/rpds_py-0.26.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc3e55a7db08dc9a6ed5fb7103019d2c1a38a349ac41901f9f66d7f95750942f", size = 383440, upload-time = "2025-07-01T15:55:05.398Z" },
-    { url = "https://files.pythonhosted.org/packages/4c/e1/c65255ad5b63903e56b3bb3ff9dcc3f4f5c3badde5d08c741ee03903e951/rpds_py-0.26.0-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:9e851920caab2dbcae311fd28f4313c6953993893eb5c1bb367ec69d9a39e7ed", size = 412759, upload-time = "2025-07-01T15:55:08.316Z" },
-    { url = "https://files.pythonhosted.org/packages/e4/22/bb731077872377a93c6e93b8a9487d0406c70208985831034ccdeed39c8e/rpds_py-0.26.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:dfbf280da5f876d0b00c81f26bedce274e72a678c28845453885a9b3c22ae632", size = 556032, upload-time = "2025-07-01T15:55:09.52Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/8b/393322ce7bac5c4530fb96fc79cc9ea2f83e968ff5f6e873f905c493e1c4/rpds_py-0.26.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:1cc81d14ddfa53d7f3906694d35d54d9d3f850ef8e4e99ee68bc0d1e5fed9a9c", size = 585416, upload-time = "2025-07-01T15:55:11.216Z" },
-    { url = "https://files.pythonhosted.org/packages/49/ae/769dc372211835bf759319a7aae70525c6eb523e3371842c65b7ef41c9c6/rpds_py-0.26.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dca83c498b4650a91efcf7b88d669b170256bf8017a5db6f3e06c2bf031f57e0", size = 554049, upload-time = "2025-07-01T15:55:13.004Z" },
-    { url = "https://files.pythonhosted.org/packages/6b/f9/4c43f9cc203d6ba44ce3146246cdc38619d92c7bd7bad4946a3491bd5b70/rpds_py-0.26.0-cp313-cp313t-win32.whl", hash = "sha256:4d11382bcaf12f80b51d790dee295c56a159633a8e81e6323b16e55d81ae37e9", size = 218428, upload-time = "2025-07-01T15:55:14.486Z" },
-    { url = "https://files.pythonhosted.org/packages/7e/8b/9286b7e822036a4a977f2f1e851c7345c20528dbd56b687bb67ed68a8ede/rpds_py-0.26.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff110acded3c22c033e637dd8896e411c7d3a11289b2edf041f86663dbc791e9", size = 231524, upload-time = "2025-07-01T15:55:15.745Z" },
-    { url = "https://files.pythonhosted.org/packages/51/f2/b5c85b758a00c513bb0389f8fc8e61eb5423050c91c958cdd21843faa3e6/rpds_py-0.26.0-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:f61a9326f80ca59214d1cceb0a09bb2ece5b2563d4e0cd37bfd5515c28510674", size = 373505, upload-time = "2025-07-01T15:56:34.716Z" },
-    { url = "https://files.pythonhosted.org/packages/23/e0/25db45e391251118e915e541995bb5f5ac5691a3b98fb233020ba53afc9b/rpds_py-0.26.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:183f857a53bcf4b1b42ef0f57ca553ab56bdd170e49d8091e96c51c3d69ca696", size = 359468, upload-time = "2025-07-01T15:56:36.219Z" },
-    { url = "https://files.pythonhosted.org/packages/0b/73/dd5ee6075bb6491be3a646b301dfd814f9486d924137a5098e61f0487e16/rpds_py-0.26.0-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:941c1cfdf4799d623cf3aa1d326a6b4fdb7a5799ee2687f3516738216d2262fb", size = 382680, upload-time = "2025-07-01T15:56:37.644Z" },
-    { url = "https://files.pythonhosted.org/packages/2f/10/84b522ff58763a5c443f5bcedc1820240e454ce4e620e88520f04589e2ea/rpds_py-0.26.0-pp311-pypy311_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:72a8d9564a717ee291f554eeb4bfeafe2309d5ec0aa6c475170bdab0f9ee8e88", size = 397035, upload-time = "2025-07-01T15:56:39.241Z" },
-    { url = "https://files.pythonhosted.org/packages/06/ea/8667604229a10a520fcbf78b30ccc278977dcc0627beb7ea2c96b3becef0/rpds_py-0.26.0-pp311-pypy311_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:511d15193cbe013619dd05414c35a7dedf2088fcee93c6bbb7c77859765bd4e8", size = 514922, upload-time = "2025-07-01T15:56:40.645Z" },
-    { url = "https://files.pythonhosted.org/packages/24/e6/9ed5b625c0661c4882fc8cdf302bf8e96c73c40de99c31e0b95ed37d508c/rpds_py-0.26.0-pp311-pypy311_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:aea1f9741b603a8d8fedb0ed5502c2bc0accbc51f43e2ad1337fe7259c2b77a5", size = 402822, upload-time = "2025-07-01T15:56:42.137Z" },
-    { url = "https://files.pythonhosted.org/packages/8a/58/212c7b6fd51946047fb45d3733da27e2fa8f7384a13457c874186af691b1/rpds_py-0.26.0-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4019a9d473c708cf2f16415688ef0b4639e07abaa569d72f74745bbeffafa2c7", size = 384336, upload-time = "2025-07-01T15:56:44.239Z" },
-    { url = "https://files.pythonhosted.org/packages/aa/f5/a40ba78748ae8ebf4934d4b88e77b98497378bc2c24ba55ebe87a4e87057/rpds_py-0.26.0-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:093d63b4b0f52d98ebae33b8c50900d3d67e0666094b1be7a12fffd7f65de74b", size = 416871, upload-time = "2025-07-01T15:56:46.284Z" },
-    { url = "https://files.pythonhosted.org/packages/d5/a6/33b1fc0c9f7dcfcfc4a4353daa6308b3ece22496ceece348b3e7a7559a09/rpds_py-0.26.0-pp311-pypy311_pp73-musllinux_1_2_aarch64.whl", hash = "sha256:2abe21d8ba64cded53a2a677e149ceb76dcf44284202d737178afe7ba540c1eb", size = 559439, upload-time = "2025-07-01T15:56:48.549Z" },
-    { url = "https://files.pythonhosted.org/packages/71/2d/ceb3f9c12f8cfa56d34995097f6cd99da1325642c60d1b6680dd9df03ed8/rpds_py-0.26.0-pp311-pypy311_pp73-musllinux_1_2_i686.whl", hash = "sha256:4feb7511c29f8442cbbc28149a92093d32e815a28aa2c50d333826ad2a20fdf0", size = 588380, upload-time = "2025-07-01T15:56:50.086Z" },
-    { url = "https://files.pythonhosted.org/packages/c8/ed/9de62c2150ca8e2e5858acf3f4f4d0d180a38feef9fdab4078bea63d8dba/rpds_py-0.26.0-pp311-pypy311_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:e99685fc95d386da368013e7fb4269dd39c30d99f812a8372d62f244f662709c", size = 555334, upload-time = "2025-07-01T15:56:51.703Z" },
-]
-
-[[package]]
-name = "s3transfer"
-version = "0.14.0"
+name = "ruff"
+version = "0.15.17"
 source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "botocore" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/62/74/8d69dcb7a9efe8baa2046891735e5dfe433ad558ae23d9e3c14c633d1d58/s3transfer-0.14.0.tar.gz", hash = "sha256:eff12264e7c8b4985074ccce27a3b38a485bb7f7422cc8046fee9be4983e4125", size = 151547, upload-time = "2025-09-09T19:23:31.089Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/8c/a9/3abdf488f1bf3d24c699415e454ed554a6350d5d89ce183be1ee0a3361ac/ruff-0.15.17.tar.gz", hash = "sha256:2ec446937fd16c8c4de2674a209cc5af64d9c6f17d21fbf1151054fa0bcf5219", size = 4743346, upload-time = "2026-06-11T17:54:47.663Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/48/f0/ae7ca09223a81a1d890b2557186ea015f6e0502e9b8cb8e1813f1d8cfa4e/s3transfer-0.14.0-py3-none-any.whl", hash = "sha256:ea3b790c7077558ed1f02a3072fb3cb992bbbd253392f4b6e9e8976941c7d456", size = 85712, upload-time = "2025-09-09T19:23:30.041Z" },
+    { url = "https://files.pythonhosted.org/packages/db/4d/e11259f5da07cb6afb2d074c31bf09da9671993f7329d4f15d2fdc458301/ruff-0.15.17-py3-none-linux_armv6l.whl", hash = "sha256:d9feddb927fc68bd295f5eebc587a7e42cfaf9b65f60ca4a2386febff575da8f", size = 10856677, upload-time = "2026-06-11T17:54:49.533Z" },
+    { url = "https://files.pythonhosted.org/packages/29/3e/772d679e1a0dc058e58875bd2c0cb713a0530877b4a76fee3c7966df0d49/ruff-0.15.17-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:25805a226d741c47d274a35ad5c10a7dde175fcddfa511d7cf3da0a21eb3eab7", size = 11223443, upload-time = "2026-06-11T17:55:00.573Z" },
+    { url = "https://files.pythonhosted.org/packages/68/58/bd41f7688b2fd5623012605130ed70e60aa7f2244baa3d5066bdd61530c8/ruff-0.15.17-py3-none-macosx_11_0_arm64.whl", hash = "sha256:f6ad73b14c2d18a3bf8ad7cb6974294d7f613a7898604826058e6ac64918ef4d", size = 10566458, upload-time = "2026-06-11T17:55:07.52Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/5b/733371013fcf1ec339e477ece6ab42bfe10bdd9bba8ee88a9516aa56bfc0/ruff-0.15.17-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ba0c1e4f95bcb3869d0d30cbd5917071ef2e28665abfec970cdab0492c713ed", size = 10914483, upload-time = "2026-06-11T17:55:05.501Z" },
+    { url = "https://files.pythonhosted.org/packages/bd/cc/6f24251cc0252f7239391ccb85833f320efad14ebe5b443943f37ced6332/ruff-0.15.17-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:81647960f10bff57d2e51cadd0c3950fe598400c852863a038720ef5b8cca91e", size = 10647497, upload-time = "2026-06-11T17:54:57.733Z" },
+    { url = "https://files.pythonhosted.org/packages/68/dd/0d10c17ce1a1624d6fc3156309c3f834fdb5dfaad026ec90c85684f3990e/ruff-0.15.17-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e01a84ddbc8c16c23055ba3924476850f1bbc1917cebbb9376665a63e74260d", size = 11416967, upload-time = "2026-06-11T17:54:51.461Z" },
+    { url = "https://files.pythonhosted.org/packages/2f/91/556bfb156f6144f355e831c23db00b2fc4120f86b3ce81cc5f7fd2df51f3/ruff-0.15.17-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:84fe9f653152f8f294f9f7e03bf3a453d8b4a27f7a59c78c8666167f2b17b96c", size = 12335770, upload-time = "2026-06-11T17:54:45.793Z" },
+    { url = "https://files.pythonhosted.org/packages/88/82/8b5999aa13355e926f06d9f42a32dcca862f623bf0363785ff89d607dffd/ruff-0.15.17-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8c0fe88a7676e7a05b73174d4d4a59cb2ac21ff8263583f87a81a6018475a978", size = 11575441, upload-time = "2026-06-11T17:54:32.661Z" },
+    { url = "https://files.pythonhosted.org/packages/11/93/f10377bb04109ca0e8cbc483ff1982c54b6d418210041776f93e8cdc7fa9/ruff-0.15.17-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ecfc3c7878fff94633ab0348524e093f9ce3243080416dd7d14f8ba400174719", size = 11557614, upload-time = "2026-06-11T17:54:34.698Z" },
+    { url = "https://files.pythonhosted.org/packages/c7/a6/eeeae7f7d5493df41649ab3db92f086b2d0a30199e4efdf8e3dd7a033f24/ruff-0.15.17-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:b8461180b22420b1bdc289909410930761629fddf2a5aaf60fae1ab26cedc4c4", size = 11544450, upload-time = "2026-06-11T17:54:39.042Z" },
+    { url = "https://files.pythonhosted.org/packages/32/88/5991ce565129a24dd4a00db1254b3b5db2e53018cbe4018ea5a89738e727/ruff-0.15.17-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6eccbe50a038b503e7140b441aa9c7fc8c1f36edf23ebef9f4165c2f28f568b7", size = 10892524, upload-time = "2026-06-11T17:55:09.432Z" },
+    { url = "https://files.pythonhosted.org/packages/f5/1d/0fdd248313425f55223968af04b0a42125466a8d88d21c1d99c6af0a51e8/ruff-0.15.17-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:382fc0521025f5a8ad447d8bdd523545d0d7646adb718eb1c2dac5065ec27c0f", size = 10659573, upload-time = "2026-06-11T17:54:36.824Z" },
+    { url = "https://files.pythonhosted.org/packages/9e/0e/072e8260deb9461062ce9311ced27a8e541229a6ffd483013dd37661e43e/ruff-0.15.17-py3-none-musllinux_1_2_i686.whl", hash = "sha256:456d41fcd1b2777ad63f09a6e7121d43f7b688bbc76a800c10f7f8fb1f912c3f", size = 11127818, upload-time = "2026-06-11T17:55:03.124Z" },
+    { url = "https://files.pythonhosted.org/packages/ab/b4/55060a34163121498014696b5f656db5b8c6963768f227dbf0d76b311073/ruff-0.15.17-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:b1a04bcc94ae6194e9db05d16ad31f298a7194bfbcb08258bbe589cee1d587b8", size = 11655901, upload-time = "2026-06-11T17:54:53.562Z" },
+    { url = "https://files.pythonhosted.org/packages/49/71/9b29d6b87cef468d697f43c6a91e3fae4a80185779d7d5a4ef27d173439f/ruff-0.15.17-py3-none-win32.whl", hash = "sha256:596065960ab1ff593f744220c9fe6580eda00a95003cffa9f4048bb5b1bf0392", size = 10925574, upload-time = "2026-06-11T17:54:55.723Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/b2/8fc77f3723228836fa5d12497eb71c808f83782e10d058d2b15cfa14640b/ruff-0.15.17-py3-none-win_amd64.whl", hash = "sha256:6769e5fa1710b179b92e0bfa5a51735b35baea9013dadb06d5f44cbcf9547084", size = 12058788, upload-time = "2026-06-11T17:54:41.042Z" },
+    { url = "https://files.pythonhosted.org/packages/2d/c7/c53e8dbff9c9dc4b7928773421ae294a5d28fcb8dcda1a089579d3a7e510/ruff-0.15.17-py3-none-win_arm64.whl", hash = "sha256:f3be1fbb34bcdfd146240d8fb92a709d4c2c8191348580a3c044ec60fa0b4456", size = 11355275, upload-time = "2026-06-11T17:54:43.635Z" },
 ]
 
 [[package]]
@@ -3610,7 +2416,8 @@ dependencies = [
     { name = "session-info2" },
     { name = "statsmodels" },
     { name = "tqdm" },
-    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.12.*'" },
     { name = "umap-learn" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/2f/be/e5d84f12a33ef7dbf856f28522c4f40d016df007f022922721a527a5ec19/scanpy-1.11.3.tar.gz", hash = "sha256:ead448f112bb07351897531c24e05627f4b3e0bee0c02f50adf378f4e08e338b", size = 14089431, upload-time = "2025-07-01T11:17:02.891Z" }
@@ -3644,6 +2451,7 @@ dependencies = [
 sdist = { url = "https://files.pythonhosted.org/packages/aa/3c/fba04785b15a74a152cd265cbdb1383bbff0bac0945f9551aab2036fceb1/scib-1.1.7.tar.gz", hash = "sha256:3bd5fed6b89adf265c317bba1a73e9418aa94574b08fab46356f5ceb98990202", size = 78796, upload-time = "2025-01-13T18:53:25.983Z" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/e5/6b/92f74b30e716d5a3e5318a3c0872bb0a0fdd253f5deee6151126275d0923/scib-1.1.7-1-py3-none-any.whl", hash = "sha256:096fb2181471182e1ef6728b675884e7fb0611ba71330068386578dcd4ce2cc7", size = 84629, upload-time = "2025-01-13T18:53:24.482Z" },
+    { url = "https://files.pythonhosted.org/packages/49/ca/fc987cae754a1741350c05d51e5325657355e30ae4b05f4c5ff55fbb9523/scib-1.1.7-py3-none-any.whl", hash = "sha256:5e153810e7ce59915ccb2ea5d61567bca059f67ae3e5209147b99f56c54e053f", size = 89981, upload-time = "2026-04-28T10:27:31.122Z" },
 ]
 
 [[package]]
@@ -3751,6 +2559,42 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/81/06/0a5e5349474e1cbc5757975b21bd4fad0e72ebf138c5592f191646154e06/scipy-1.15.3-cp313-cp313t-win_amd64.whl", hash = "sha256:76ad1fb5f8752eabf0fa02e4cc0336b4e8f021e2d5f061ed37d6d264db35e3ca", size = 40308097, upload-time = "2025-05-08T16:08:27.627Z" },
 ]
 
+[[package]]
+name = "scverse-misc"
+version = "0.0.3"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version < '3.12' and sys_platform == 'linux'",
+    "python_full_version < '3.12' and sys_platform != 'linux'",
+]
+dependencies = [
+    { name = "session-info2", marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f7/05/6123a4362e2810ef216f152e249a66799f9c37975dabf89b69abb2d68c42/scverse_misc-0.0.3.tar.gz", hash = "sha256:18c46eeeac8ccef8f435e41a8ee86173b3d7ef6ea1167fde97a17553f70d3210", size = 23128, upload-time = "2026-04-10T15:12:21.731Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/bd/3e/9dad188e00d08f5381b95dd1f8f7b1dc9eb244250d93fea3e9cb7d0e61f0/scverse_misc-0.0.3-py3-none-any.whl", hash = "sha256:295702368db8f5fe946b5296135990492b3fc891a45bd6dd27798775db4c61c6", size = 8845, upload-time = "2026-04-10T15:12:20.366Z" },
+]
+
+[[package]]
+name = "scverse-misc"
+version = "0.0.8"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version >= '3.13' and sys_platform == 'linux'",
+    "python_full_version >= '3.13' and sys_platform != 'linux'",
+    "python_full_version == '3.12.*' and sys_platform == 'linux'",
+    "python_full_version == '3.12.*' and sys_platform != 'linux'",
+]
+dependencies = [
+    { name = "session-info2", marker = "python_full_version >= '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.12.*'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a3/2e/2fc93b1c9a7dc3c9c44ba8d239a094d21c8cbf3219a57213fc14a7e9bf28/scverse_misc-0.0.8.tar.gz", hash = "sha256:7de6d46e50fc111d366d60b07f165531de97a836a4d702bbbfa2a47a0015fd9e", size = 31621, upload-time = "2026-06-08T12:01:27.683Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/82/58/3c753c43fed0a96db29b787daa1dacd84acc8f484da5dd5f8627601a4e4a/scverse_misc-0.0.8-py3-none-any.whl", hash = "sha256:dc61dbd9a1bdb1171cda9f73c52567e29159d51ff1a0f24bee2e568ca1f96a4a", size = 13955, upload-time = "2026-06-08T12:01:26.639Z" },
+]
+
 [[package]]
 name = "seaborn"
 version = "0.13.2"
@@ -3778,32 +2622,24 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/54/24/b4293291fa1dd830f353d2cb163295742fa87f179fcc8a20a306a81978b7/SecretStorage-3.3.3-py3-none-any.whl", hash = "sha256:f356e6628222568e3af06f2eba8df495efa13b3b63081dafd4f7d9a7b7bc9f99", size = 15221, upload-time = "2022-08-13T16:22:44.457Z" },
 ]
 
-[[package]]
-name = "send2trash"
-version = "1.8.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/fd/3a/aec9b02217bb79b87bbc1a21bc6abc51e3d5dcf65c30487ac96c0908c722/Send2Trash-1.8.3.tar.gz", hash = "sha256:b18e7a3966d99871aefeb00cfbcfdced55ce4871194810fc71f4aa484b953abf", size = 17394, upload-time = "2024-04-07T00:01:09.267Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/40/b0/4562db6223154aa4e22f939003cb92514c79f3d4dccca3444253fd17f902/Send2Trash-1.8.3-py3-none-any.whl", hash = "sha256:0c31227e0bd08961c7665474a3d1ef7193929fedda4233843689baa056be46c9", size = 18072, upload-time = "2024-04-07T00:01:07.438Z" },
-]
-
 [[package]]
 name = "sentence-transformers"
-version = "5.0.0"
+version = "5.5.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub" },
-    { name = "pillow" },
+    { name = "numpy" },
     { name = "scikit-learn" },
     { name = "scipy" },
     { name = "torch" },
     { name = "tqdm" },
     { name = "transformers" },
-    { name = "typing-extensions" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/99/69/2a29773b43a24ee04eb26af492d85d520b30a86cfef22a0885e77e9c4a16/sentence_transformers-5.0.0.tar.gz", hash = "sha256:e5a411845910275fd166bacb01d28b7f79537d3550628ae42309dbdd3d5670d1", size = 366847, upload-time = "2025-07-01T13:01:33.04Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/cf/d4/7ef93157485e978c016f49da05363c1e4e7237beb5343b64b5631101f0f1/sentence_transformers-5.5.1.tar.gz", hash = "sha256:02b7740dfc60bdbbcb6061625f5d97a5c1a4e2d3baac5f9391b912bb5eae2290", size = 445161, upload-time = "2026-05-20T07:37:44.465Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/6f/ff/178f08ea5ebc1f9193d9de7f601efe78c01748347875c8438f66f5cecc19/sentence_transformers-5.0.0-py3-none-any.whl", hash = "sha256:346240f9cc6b01af387393f03e103998190dfb0826a399d0c38a81a05c7a5d76", size = 470191, upload-time = "2025-07-01T13:01:31.619Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/03/ee99a6b030e7a2e056547729f8a4709dd93e13d9c6f07590f74c395c4017/sentence_transformers-5.5.1-py3-none-any.whl", hash = "sha256:4fe11d433badc5282d32f7fc08bc714216b7a5aca426f9df77a45a554756deb7", size = 588887, upload-time = "2026-05-20T07:37:43.004Z" },
 ]
 
 [[package]]
@@ -3864,29 +2700,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" },
 ]
 
-[[package]]
-name = "soupsieve"
-version = "2.7"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/3f/f4/4a80cd6ef364b2e8b65b15816a843c0980f7a5a2b4dc701fc574952aa19f/soupsieve-2.7.tar.gz", hash = "sha256:ad282f9b6926286d2ead4750552c8a6142bc4c783fd66b0293547c8fe6ae126a", size = 103418, upload-time = "2025-04-20T18:50:08.518Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e7/9c/0e6afc12c269578be5c0c1c9f4b49a8d32770a080260c333ac04cc1c832d/soupsieve-2.7-py3-none-any.whl", hash = "sha256:6e60cc5c1ffaf1cebcc12e8188320b72071e922c2e897f737cadce79ad5d30c4", size = 36677, upload-time = "2025-04-20T18:50:07.196Z" },
-]
-
-[[package]]
-name = "stack-data"
-version = "0.6.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "asttokens" },
-    { name = "executing" },
-    { name = "pure-eval" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/28/e3/55dcc2cfbc3ca9c29519eb6884dd1415ecb53b0e934862d3559ddcb7e20b/stack_data-0.6.3.tar.gz", hash = "sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9", size = 44707, upload-time = "2023-09-30T13:58:05.479Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695", size = 24521, upload-time = "2023-09-30T13:58:03.53Z" },
-]
-
 [[package]]
 name = "statsmodels"
 version = "0.14.5"
@@ -3932,20 +2745,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" },
 ]
 
-[[package]]
-name = "terminado"
-version = "0.18.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "ptyprocess", marker = "os_name != 'nt'" },
-    { name = "pywinpty", marker = "os_name == 'nt' and sys_platform != 'linux'" },
-    { name = "tornado" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/8a/11/965c6fd8e5cc254f1fe142d547387da17a8ebfd75a3455f637c663fb38a0/terminado-0.18.1.tar.gz", hash = "sha256:de09f2c4b85de4765f7714688fff57d3e75bad1f909b589fde880460c753fd2e", size = 32701, upload-time = "2024-03-12T14:34:39.026Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/6a/9e/2064975477fdc887e47ad42157e214526dcad8f317a948dee17e1659a62f/terminado-0.18.1-py3-none-any.whl", hash = "sha256:a4468e1b37bb318f8a86514f65814e1afc977cf29b3992a4500d9dd305dcceb0", size = 14154, upload-time = "2024-03-12T14:34:36.569Z" },
-]
-
 [[package]]
 name = "texttable"
 version = "1.7.0"
@@ -3964,18 +2763,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" },
 ]
 
-[[package]]
-name = "tinycss2"
-version = "1.4.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "webencodings" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/7a/fd/7a5ee21fd08ff70d3d33a5781c255cbe779659bd03278feb98b19ee550f4/tinycss2-1.4.0.tar.gz", hash = "sha256:10c0972f6fc0fbee87c3edb76549357415e94548c1ae10ebccdea16fb404a9b7", size = 87085, upload-time = "2024-10-24T14:58:29.895Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e6/34/ebdc18bae6aa14fbee1a08b63c015c72b64868ff7dae68808ab500c492e2/tinycss2-1.4.0-py3-none-any.whl", hash = "sha256:3a49cf47b7675da0b15d0c6e1df8df4ebd96e9394bb905a5775adb0d884c5289", size = 26610, upload-time = "2024-10-24T14:58:28.029Z" },
-]
-
 [[package]]
 name = "tokenizers"
 version = "0.22.1"
@@ -4027,7 +2814,8 @@ dependencies = [
     { name = "setuptools", marker = "python_full_version >= '3.12'" },
     { name = "sympy" },
     { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "typing-extensions" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/8f/c4/3e7a3887eba14e815e614db70b3b529112d1513d9dae6f4d43e373360b7f/torch-2.8.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:220a06fd7af8b653c35d359dfe1aaf32f65aa85befa342629f716acb134b9710", size = 102073391, upload-time = "2025-08-06T14:53:20.937Z" },
@@ -4048,25 +2836,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/04/6e/650bb7f28f771af0cb791b02348db8b7f5f64f40f6829ee82aa6ce99aabe/torch-2.8.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:7b677e17f5a3e69fdef7eb3b9da72622f8d322692930297e4ccb52fefc6c8211", size = 73632395, upload-time = "2025-08-06T14:55:28.645Z" },
 ]
 
-[[package]]
-name = "tornado"
-version = "6.5.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/51/89/c72771c81d25d53fe33e3dca61c233b665b2780f21820ba6fd2c6793c12b/tornado-6.5.1.tar.gz", hash = "sha256:84ceece391e8eb9b2b95578db65e920d2a61070260594819589609ba9bc6308c", size = 509934, upload-time = "2025-05-22T18:15:38.788Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/77/89/f4532dee6843c9e0ebc4e28d4be04c67f54f60813e4bf73d595fe7567452/tornado-6.5.1-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:d50065ba7fd11d3bd41bcad0825227cc9a95154bad83239357094c36708001f7", size = 441948, upload-time = "2025-05-22T18:15:20.862Z" },
-    { url = "https://files.pythonhosted.org/packages/15/9a/557406b62cffa395d18772e0cdcf03bed2fff03b374677348eef9f6a3792/tornado-6.5.1-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9e9ca370f717997cb85606d074b0e5b247282cf5e2e1611568b8821afe0342d6", size = 440112, upload-time = "2025-05-22T18:15:22.591Z" },
-    { url = "https://files.pythonhosted.org/packages/55/82/7721b7319013a3cf881f4dffa4f60ceff07b31b394e459984e7a36dc99ec/tornado-6.5.1-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b77e9dfa7ed69754a54c89d82ef746398be82f749df69c4d3abe75c4d1ff4888", size = 443672, upload-time = "2025-05-22T18:15:24.027Z" },
-    { url = "https://files.pythonhosted.org/packages/7d/42/d11c4376e7d101171b94e03cef0cbce43e823ed6567ceda571f54cf6e3ce/tornado-6.5.1-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:253b76040ee3bab8bcf7ba9feb136436a3787208717a1fb9f2c16b744fba7331", size = 443019, upload-time = "2025-05-22T18:15:25.735Z" },
-    { url = "https://files.pythonhosted.org/packages/7d/f7/0c48ba992d875521ac761e6e04b0a1750f8150ae42ea26df1852d6a98942/tornado-6.5.1-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:308473f4cc5a76227157cdf904de33ac268af770b2c5f05ca6c1161d82fdd95e", size = 443252, upload-time = "2025-05-22T18:15:27.499Z" },
-    { url = "https://files.pythonhosted.org/packages/89/46/d8d7413d11987e316df4ad42e16023cd62666a3c0dfa1518ffa30b8df06c/tornado-6.5.1-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:caec6314ce8a81cf69bd89909f4b633b9f523834dc1a352021775d45e51d9401", size = 443930, upload-time = "2025-05-22T18:15:29.299Z" },
-    { url = "https://files.pythonhosted.org/packages/78/b2/f8049221c96a06df89bed68260e8ca94beca5ea532ffc63b1175ad31f9cc/tornado-6.5.1-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:13ce6e3396c24e2808774741331638ee6c2f50b114b97a55c5b442df65fd9692", size = 443351, upload-time = "2025-05-22T18:15:31.038Z" },
-    { url = "https://files.pythonhosted.org/packages/76/ff/6a0079e65b326cc222a54720a748e04a4db246870c4da54ece4577bfa702/tornado-6.5.1-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:5cae6145f4cdf5ab24744526cc0f55a17d76f02c98f4cff9daa08ae9a217448a", size = 443328, upload-time = "2025-05-22T18:15:32.426Z" },
-    { url = "https://files.pythonhosted.org/packages/49/18/e3f902a1d21f14035b5bc6246a8c0f51e0eef562ace3a2cea403c1fb7021/tornado-6.5.1-cp39-abi3-win32.whl", hash = "sha256:e0a36e1bc684dca10b1aa75a31df8bdfed656831489bc1e6a6ebed05dc1ec365", size = 444396, upload-time = "2025-05-22T18:15:34.205Z" },
-    { url = "https://files.pythonhosted.org/packages/7b/09/6526e32bf1049ee7de3bebba81572673b19a2a8541f795d887e92af1a8bc/tornado-6.5.1-cp39-abi3-win_amd64.whl", hash = "sha256:908e7d64567cecd4c2b458075589a775063453aeb1d2a1853eedb806922f568b", size = 444840, upload-time = "2025-05-22T18:15:36.1Z" },
-    { url = "https://files.pythonhosted.org/packages/55/a7/535c44c7bea4578e48281d83c615219f3ab19e6abc67625ef637c73987be/tornado-6.5.1-cp39-abi3-win_arm64.whl", hash = "sha256:02420a0eb7bf617257b9935e2b754d1b63897525d8a289c9d65690d580b4dcf7", size = 443596, upload-time = "2025-05-22T18:15:37.433Z" },
-]
-
 [[package]]
 name = "tqdm"
 version = "4.66.5"
@@ -4079,15 +2848,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/48/5d/acf5905c36149bbaec41ccf7f2b68814647347b72075ac0b1fe3022fdc73/tqdm-4.66.5-py3-none-any.whl", hash = "sha256:90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd", size = 78351, upload-time = "2024-08-03T22:35:36.644Z" },
 ]
 
-[[package]]
-name = "traitlets"
-version = "5.14.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/eb/79/72064e6a701c2183016abbbfedaba506d81e30e232a68c9f0d6f6fcd1574/traitlets-5.14.3.tar.gz", hash = "sha256:9ed0579d3502c94b4b3732ac120375cda96f923114522847de4b3bb98b96b6b7", size = 161621, upload-time = "2024-04-19T11:11:49.746Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/00/c0/8f5d070730d7836adc9c9b6408dec68c6ced86b304a9b26a14df072a6e8c/traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f", size = 85359, upload-time = "2024-04-19T11:11:46.763Z" },
-]
-
 [[package]]
 name = "transformers"
 version = "4.57.1"
@@ -4147,17 +2907,37 @@ wheels = [
 name = "typing-extensions"
 version = "4.12.2"
 source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version < '3.12' and sys_platform == 'linux'",
+    "python_full_version < '3.12' and sys_platform != 'linux'",
+]
 sdist = { url = "https://files.pythonhosted.org/packages/df/db/f35a00659bc03fec321ba8bce9420de607a1d37f8342eee1863174c69557/typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8", size = 85321, upload-time = "2024-06-07T18:52:15.995Z" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/26/9f/ad63fc0248c5379346306f8668cda6e2e2e9c95e01216d2b8ffd9ff037d0/typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d", size = 37438, upload-time = "2024-06-07T18:52:13.582Z" },
 ]
 
+[[package]]
+name = "typing-extensions"
+version = "4.15.0"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version >= '3.13' and sys_platform == 'linux'",
+    "python_full_version >= '3.13' and sys_platform != 'linux'",
+    "python_full_version == '3.12.*' and sys_platform == 'linux'",
+    "python_full_version == '3.12.*' and sys_platform != 'linux'",
+]
+sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" },
+]
+
 [[package]]
 name = "typing-inspection"
 version = "0.4.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "typing-extensions" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/f8/b1/0c11f5058406b3af7609f121aaa6b609744687f1d158b3c3a5bf4cc94238/typing_inspection-0.4.1.tar.gz", hash = "sha256:6ae134cc0203c33377d43188d4064e9b357dba58cff3185f22924610e70a9d28", size = 75726, upload-time = "2025-05-21T18:55:23.885Z" }
 wheels = [
@@ -4190,27 +2970,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/6b/b1/c24deeda9baf1fd491aaad941ed89e0fed6c583a117fd7b79e0a33a1e6c0/umap_learn-0.5.9.post2-py3-none-any.whl", hash = "sha256:fbe51166561e0e7fab00ef3d516ac2621243b8d15cf4bef9f656d701736b16a0", size = 90146, upload-time = "2025-07-03T00:18:01.042Z" },
 ]
 
-[[package]]
-name = "uri-template"
-version = "1.3.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/31/c7/0336f2bd0bcbada6ccef7aaa25e443c118a704f828a0620c6fa0207c1b64/uri-template-1.3.0.tar.gz", hash = "sha256:0e00f8eb65e18c7de20d595a14336e9f337ead580c70934141624b6d1ffdacc7", size = 21678, upload-time = "2023-06-21T01:49:05.374Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e7/00/3fca040d7cf8a32776d3d81a00c8ee7457e00f80c649f1e4a863c8321ae9/uri_template-1.3.0-py3-none-any.whl", hash = "sha256:a44a133ea12d44a0c0f06d7d42a52d71282e77e2f937d8abd5655b8d56fc1363", size = 11140, upload-time = "2023-06-21T01:49:03.467Z" },
-]
-
-[[package]]
-name = "url-normalize"
-version = "2.2.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "idna" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/80/31/febb777441e5fcdaacb4522316bf2a527c44551430a4873b052d545e3279/url_normalize-2.2.1.tar.gz", hash = "sha256:74a540a3b6eba1d95bdc610c24f2c0141639f3ba903501e61a52a8730247ff37", size = 18846, upload-time = "2025-04-26T20:37:58.553Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/bc/d9/5ec15501b675f7bc07c5d16aa70d8d778b12375686b6efd47656efdc67cd/url_normalize-2.2.1-py3-none-any.whl", hash = "sha256:3deb687587dc91f7b25c9ae5162ffc0f057ae85d22b1e15cf5698311247f567b", size = 14728, upload-time = "2025-04-26T20:37:57.217Z" },
-]
-
 [[package]]
 name = "urllib3"
 version = "1.26.13"
@@ -4236,7 +2995,7 @@ wheels = [
 
 [[package]]
 name = "wandb"
-version = "0.21.0"
+version = "0.27.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "click" },
@@ -4248,65 +3007,20 @@ dependencies = [
     { name = "pyyaml" },
     { name = "requests" },
     { name = "sentry-sdk" },
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/73/09/c84264a219e20efd615e4d5d150cc7d359d57d51328d3fa94ee02d70ed9c/wandb-0.21.0.tar.gz", hash = "sha256:473e01ef200b59d780416062991effa7349a34e51425d4be5ff482af2dc39e02", size = 40085784, upload-time = "2025-07-02T00:24:15.516Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/38/dd/65eac086e1bc337bb5f0eed65ba1fe4a6dbc62c97f094e8e9df1ef83ffed/wandb-0.21.0-py3-none-any.whl", hash = "sha256:316e8cd4329738f7562f7369e6eabeeb28ef9d473203f7ead0d03e5dba01c90d", size = 6504284, upload-time = "2025-07-02T00:23:46.671Z" },
-    { url = "https://files.pythonhosted.org/packages/17/a7/80556ce9097f59e10807aa68f4a9b29d736a90dca60852a9e2af1641baf8/wandb-0.21.0-py3-none-macosx_10_14_x86_64.whl", hash = "sha256:701d9cbdfcc8550a330c1b54a26f1585519180e0f19247867446593d34ace46b", size = 21717388, upload-time = "2025-07-02T00:23:49.348Z" },
-    { url = "https://files.pythonhosted.org/packages/23/ae/660bc75aa37bd23409822ea5ed616177d94873172d34271693c80405c820/wandb-0.21.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:01689faa6b691df23ba2367e0a1ecf6e4d0be44474905840098eedd1fbcb8bdf", size = 21141465, upload-time = "2025-07-02T00:23:52.602Z" },
-    { url = "https://files.pythonhosted.org/packages/23/ab/9861929530be56557c74002868c85d0d8ac57050cc21863afe909ae3d46f/wandb-0.21.0-py3-none-macosx_11_0_x86_64.whl", hash = "sha256:55d3f42ddb7971d1699752dff2b85bcb5906ad098d18ab62846c82e9ce5a238d", size = 21793511, upload-time = "2025-07-02T00:23:55.447Z" },
-    { url = "https://files.pythonhosted.org/packages/de/52/e5cad2eff6fbed1ac06f4a5b718457fa2fd437f84f5c8f0d31995a2ef046/wandb-0.21.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:893508f0c7da48917448daa5cd622c27ce7ce15119adaa861185034c2bd7b14c", size = 20704643, upload-time = "2025-07-02T00:23:58.255Z" },
-    { url = "https://files.pythonhosted.org/packages/83/8f/6bed9358cc33767c877b221d4f565e1ddf00caf4bbbe54d2e3bbc932c6a7/wandb-0.21.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a4e8245a8912247ddf7654f7b5330f583a6c56ab88fee65589158490d583c57d", size = 22243012, upload-time = "2025-07-02T00:24:01.423Z" },
-    { url = "https://files.pythonhosted.org/packages/be/61/9048015412ea5ca916844af55add4fed7c21fe1ad70bb137951e70b550c5/wandb-0.21.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2e4c4f951e0d02755e315679bfdcb5bc38c1b02e2e5abc5432b91a91bb0cf246", size = 20716440, upload-time = "2025-07-02T00:24:04.198Z" },
-    { url = "https://files.pythonhosted.org/packages/02/d9/fcd2273d8ec3f79323e40a031aba5d32d6fa9065702010eb428b5ffbab62/wandb-0.21.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:873749966eeac0069e0e742e6210641b6227d454fb1dae2cf5c437c6ed42d3ca", size = 22320652, upload-time = "2025-07-02T00:24:07.175Z" },
-    { url = "https://files.pythonhosted.org/packages/80/68/b8308db6b9c3c96dcd03be17c019aee105e1d7dc1e74d70756cdfb9241c6/wandb-0.21.0-py3-none-win32.whl", hash = "sha256:9d3cccfba658fa011d6cab9045fa4f070a444885e8902ae863802549106a5dab", size = 21484296, upload-time = "2025-07-02T00:24:10.147Z" },
-    { url = "https://files.pythonhosted.org/packages/cf/96/71cc033e8abd00e54465e68764709ed945e2da2d66d764f72f4660262b22/wandb-0.21.0-py3-none-win_amd64.whl", hash = "sha256:28a0b2dad09d7c7344ac62b0276be18a2492a5578e4d7c84937a3e1991edaac7", size = 21484301, upload-time = "2025-07-02T00:24:12.658Z" },
-]
-
-[[package]]
-name = "wcwidth"
-version = "0.2.13"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/6c/63/53559446a878410fc5a5974feb13d31d78d752eb18aeba59c7fef1af7598/wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5", size = 101301, upload-time = "2024-01-06T02:10:57.829Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/fd/84/fd2ba7aafacbad3c4201d395674fc6348826569da3c0937e75505ead3528/wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859", size = 34166, upload-time = "2024-01-06T02:10:55.763Z" },
-]
-
-[[package]]
-name = "webcolors"
-version = "25.10.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/1d/7a/eb316761ec35664ea5174709a68bbd3389de60d4a1ebab8808bfc264ed67/webcolors-25.10.0.tar.gz", hash = "sha256:62abae86504f66d0f6364c2a8520de4a0c47b80c03fc3a5f1815fedbef7c19bf", size = 53491, upload-time = "2025-10-31T07:51:03.977Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e2/cc/e097523dd85c9cf5d354f78310927f1656c422bd7b2613b2db3e3f9a0f2c/webcolors-25.10.0-py3-none-any.whl", hash = "sha256:032c727334856fc0b968f63daa252a1ac93d33db2f5267756623c210e57a4f1d", size = 14905, upload-time = "2025-10-31T07:51:01.778Z" },
-]
-
-[[package]]
-name = "webencodings"
-version = "0.5.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/0b/02/ae6ceac1baeda530866a85075641cec12989bd8d31af6d5ab4a3e8c92f47/webencodings-0.5.1.tar.gz", hash = "sha256:b36a1c245f2d304965eb4e0a82848379241dc04b865afcc4aab16748587e1923", size = 9721, upload-time = "2017-04-05T20:21:34.189Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f4/24/2a3e3df732393fed8b3ebf2ec078f05546de641fe1b667ee316ec1dcf3b7/webencodings-0.5.1-py2.py3-none-any.whl", hash = "sha256:a0af1213f3c2226497a97e2b3aa01a7e4bee4f403f95be16fc9acd2947514a78", size = 11774, upload-time = "2017-04-05T20:21:32.581Z" },
-]
-
-[[package]]
-name = "websocket-client"
-version = "1.9.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/2c/41/aa4bf9664e4cda14c3b39865b12251e8e7d239f4cd0e3cc1b6c2ccde25c1/websocket_client-1.9.0.tar.gz", hash = "sha256:9e813624b6eb619999a97dc7958469217c3176312b3a16a4bd1bc7e08a46ec98", size = 70576, upload-time = "2025-10-07T21:16:36.495Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/34/db/b10e48aa8fff7407e67470363eac595018441cf32d5e1001567a7aeba5d2/websocket_client-1.9.0-py3-none-any.whl", hash = "sha256:af248a825037ef591efbf6ed20cc5faa03d3b47b9e5a2230a529eeee1c1fc3ef", size = 82616, upload-time = "2025-10-07T21:16:34.951Z" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
-
-[[package]]
-name = "widgetsnbextension"
-version = "4.0.15"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/bd/f4/c67440c7fb409a71b7404b7aefcd7569a9c0d6bd071299bf4198ae7a5d95/widgetsnbextension-4.0.15.tar.gz", hash = "sha256:de8610639996f1567952d763a5a41af8af37f2575a41f9852a38f947eb82a3b9", size = 1097402, upload-time = "2025-11-01T21:15:55.178Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/8e/31/fe53d06b75ef0a7f2f0ee5931a89f7aedc27d233840b1839616860fed256/wandb-0.27.0.tar.gz", hash = "sha256:579e75300173059f9334e1f513a79ef15f6d9ea5c74e20d695633648cdd02031", size = 41090732, upload-time = "2026-05-14T03:44:08.894Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/3f/0e/fa3b193432cfc60c93b42f3be03365f5f909d2b3ea410295cf36df739e31/widgetsnbextension-4.0.15-py3-none-any.whl", hash = "sha256:8156704e4346a571d9ce73b84bee86a29906c9abfd7223b7228a28899ccf3366", size = 2196503, upload-time = "2025-11-01T21:15:53.565Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/5e/2c199e70e636ecfd217cde0bc7469f4511e1d03d0685eb92bfdfce391430/wandb-0.27.0-py3-none-macosx_12_0_arm64.whl", hash = "sha256:c156be4851485f3c4160cb6eb2e8991b4cdeffbccefc5636d33cf5e254847365", size = 24886476, upload-time = "2026-05-14T03:43:27.569Z" },
+    { url = "https://files.pythonhosted.org/packages/0b/cd/a617c871cd304a9804e56a7ec2ec2c65685bf0091a2b9f91910175a149e2/wandb-0.27.0-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:20179f38afb0158859a4141d29ac650d3fdbd0cf801a74ce25565c934f03776c", size = 26045779, upload-time = "2026-05-14T03:43:31.999Z" },
+    { url = "https://files.pythonhosted.org/packages/10/0a/d3f159a201530b84b72ca5f98c68d1f351c2d9a1864558ed76c811407fae/wandb-0.27.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:626497d7975fa898d0a4a239da7a510483495ca3514510dbe75004a25963af4d", size = 25480764, upload-time = "2026-05-14T03:43:35.922Z" },
+    { url = "https://files.pythonhosted.org/packages/5f/6a/8721fcdf71d42639191040a77a585d2982402b1754700cb2ecfc2ca1470a/wandb-0.27.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:f772da7005cc26a2a32b729a16982a583dc68b3d493df6a09d0aa5c5ca5a2060", size = 27256204, upload-time = "2026-05-14T03:43:39.765Z" },
+    { url = "https://files.pythonhosted.org/packages/00/5e/279d167ba79fb7a8a43401c9f25efd0f6663ee9bd1eaf5a8578530198888/wandb-0.27.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:63acfc5b994e4a90e4a2fbdee6d45e664da3dd865bb1419942c8995c06c41cf1", size = 25647469, upload-time = "2026-05-14T03:43:44.817Z" },
+    { url = "https://files.pythonhosted.org/packages/94/51/a69ac59300e3c813939d0764348959ed2a21e14c668cb1cebcb04010da6a/wandb-0.27.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:17aae6e4a88cd05c00ea8f546220918e3ebb6f8c1c36b70ef04a5ac75f0d7160", size = 27599005, upload-time = "2026-05-14T03:43:50.926Z" },
+    { url = "https://files.pythonhosted.org/packages/5f/40/bf510c8758727df020f83b717ebc1fcc1739ed7f6ae1796ebef60bf6f592/wandb-0.27.0-py3-none-win32.whl", hash = "sha256:0bd5659417e386bf6538b5e2ffe6885774c6197f0e4853bfed517d5b0db457f1", size = 25036164, upload-time = "2026-05-14T03:43:54.839Z" },
+    { url = "https://files.pythonhosted.org/packages/54/ff/69f88e7d90c22b79bcb911143c13e59742ee192080b21015ff83a5a1f60a/wandb-0.27.0-py3-none-win_amd64.whl", hash = "sha256:89d584b73166eecee96fb446f18d0e45b1aa45aba6a3696296f3f06d7454516b", size = 25036170, upload-time = "2026-05-14T03:43:59.227Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/38/f7efd7a87297a55c7e9a331a1dbb5b19e54aeacc11fe6f43f8636a73987c/wandb-0.27.0-py3-none-win_arm64.whl", hash = "sha256:a6c129c311edf210a2b4f2f4acc557eff522628125f5f28ed27df19c16c07079", size = 22972710, upload-time = "2026-05-14T03:44:03.275Z" },
 ]
 
 [[package]]
@@ -4362,15 +3076,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/2d/82/f56956041adef78f849db6b289b282e72b55ab8045a75abad81898c28d19/wrapt-1.17.2-py3-none-any.whl", hash = "sha256:b18f2d1533a71f069c7f82d524a52599053d4c7166e9dd374ae2136b7f40f7c8", size = 23594, upload-time = "2025-01-14T10:35:44.018Z" },
 ]
 
-[[package]]
-name = "xmltodict"
-version = "1.0.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/6a/aa/917ceeed4dbb80d2f04dbd0c784b7ee7bba8ae5a54837ef0e5e062cd3cfb/xmltodict-1.0.2.tar.gz", hash = "sha256:54306780b7c2175a3967cad1db92f218207e5bc1aba697d887807c0fb68b7649", size = 25725, upload-time = "2025-09-17T21:59:26.459Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c0/20/69a0e6058bc5ea74892d089d64dfc3a62ba78917ec5e2cfa70f7c92ba3a5/xmltodict-1.0.2-py3-none-any.whl", hash = "sha256:62d0fddb0dcbc9f642745d8bbf4d81fd17d6dfaec5a15b5c1876300aad92af0d", size = 13893, upload-time = "2025-09-17T21:59:24.859Z" },
-]
-
 [[package]]
 name = "xxhash"
 version = "3.5.0"
@@ -4508,17 +3213,46 @@ wheels = [
 
 [[package]]
 name = "zarr"
-version = "2.18.7"
+version = "3.1.6"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version < '3.12' and sys_platform == 'linux'",
+    "python_full_version < '3.12' and sys_platform != 'linux'",
+]
+dependencies = [
+    { name = "donfig", marker = "python_full_version < '3.12'" },
+    { name = "google-crc32c", marker = "python_full_version < '3.12'" },
+    { name = "numcodecs", marker = "python_full_version < '3.12'" },
+    { name = "numpy", marker = "python_full_version < '3.12'" },
+    { name = "packaging", marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/31/5a/b8a0cf39a14c770c30bd1f2d120c54000c8cd9e84e8e79f38d9a7ce58071/zarr-3.1.6.tar.gz", hash = "sha256:d95e72cbea4b90e9a70679468b8266400331756232576ae2b43400ac5108d0eb", size = 386531, upload-time = "2026-03-23T17:25:18.748Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/de/7c/ba8ca8cbe9dbef8e83a95fc208fed8e6686c98b4719aaa0aa7f3d31fe390/zarr-3.1.6-py3-none-any.whl", hash = "sha256:b5a82c5079d1c3d4ee8f06746fa3b9a98a7d804300fa3f4be154362a33e1207e", size = 295655, upload-time = "2026-03-23T17:25:17.189Z" },
+]
+
+[[package]]
+name = "zarr"
+version = "3.2.1"
 source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version >= '3.13' and sys_platform == 'linux'",
+    "python_full_version >= '3.13' and sys_platform != 'linux'",
+    "python_full_version == '3.12.*' and sys_platform == 'linux'",
+    "python_full_version == '3.12.*' and sys_platform != 'linux'",
+]
 dependencies = [
-    { name = "asciitree" },
-    { name = "fasteners", marker = "sys_platform != 'emscripten'" },
-    { name = "numcodecs" },
-    { name = "numpy" },
+    { name = "donfig", marker = "python_full_version >= '3.12'" },
+    { name = "google-crc32c", marker = "python_full_version >= '3.12'" },
+    { name = "numcodecs", marker = "python_full_version >= '3.12'" },
+    { name = "numpy", marker = "python_full_version >= '3.12'" },
+    { name = "packaging", marker = "python_full_version >= '3.12'" },
+    { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/da/1d/01cf9e3ab2d85190278efc3fca9f68563de35ae30ee59e7640e3af98abe3/zarr-2.18.7.tar.gz", hash = "sha256:b2b8f66f14dac4af66b180d2338819981b981f70e196c9a66e6bfaa9e59572f5", size = 3604558, upload-time = "2025-04-09T07:59:28.482Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/93/8d/aeb164004f87543b06ef54f885d02c342c31ceb274e2bbec470a98927621/zarr-3.2.1.tar.gz", hash = "sha256:71565b738a0e7e8ed226f0516eba8c6bb53440ad7669a8c48ebb3534a161d035", size = 675161, upload-time = "2026-05-05T12:37:22.383Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/5e/d8/9ffd8c237b3559945bb52103cf0eed64ea098f7b7f573f8d2962ef27b4b2/zarr-2.18.7-py3-none-any.whl", hash = "sha256:ac3dc4033e9ae4e9d7b5e27c97ea3eaf1003cc0a07f010bd83d5134bf8c4b223", size = 211273, upload-time = "2025-04-09T07:59:27.039Z" },
+    { url = "https://files.pythonhosted.org/packages/88/0a/469e2bd01be1490336e6c8707386845655d59261543315778a3ccc7e8019/zarr-3.2.1-py3-none-any.whl", hash = "sha256:f78cdd3d9687ad0e9f9cba2c5683b64f0c52589c19f685eeabe872e93cc0d2c7", size = 319617, upload-time = "2026-05-05T12:37:20.66Z" },
 ]
 
 [[package]]