Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.py[cod]
+.ipynb_checkpoints/
+*.npy

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,101 @@
 ---
+language:
+- en
+library_name: pytorch
+pipeline_tag: feature-extraction
+tags:
+- single-cell
+- transcriptomics
+- scRNA-seq
+- cell-embedding
+- pytorch
 license: apache-2.0
 ---
+
+# scRep
+
+scRep is a PyTorch model for extracting L2-normalized cell embeddings from
+single-cell RNA-seq AnnData (`.h5ad`) inputs. This repository is a standalone
+Hugging Face release bundle: it includes checkpoint weights, the exact paired
+gene vocabulary, the inference implementation, and runnable examples.
+
+## Checkpoints
+
+| Checkpoint | Directory | Vocabulary size | Required paired tokenizer |
+| --- | --- | ---: | --- |
+| 2026-06-25 / 30M | `checkpoints/scRep_20260625_30M/` | 19,240 | `checkpoints/scRep_20260625_30M/tokenizer/` |
+| 2026-07-17 / 3M | `checkpoints/scRep_20260717_3M/` | 19,239 | `checkpoints/scRep_20260717_3M/tokenizer/` |
+
+`model.safetensors`, `config.json`, and `tokenizer/gene_vocab.json` from the
+same checkpoint directory must always be used together. The checkpoint names
+are retained as release identifiers; they do not by themselves specify a
+parameter count or a benchmark claim.
+
+## Installation
+
+Python 3.10--3.12 is supported. Install a PyTorch build appropriate for the
+target CPU/CUDA platform first, then install the remaining requirements:
+
+```bash
+pip install -r requirements.txt
+```
+
+## Upload to Hugging Face
+
+Upload the contents of this directory as a model repository. The included
+`.gitattributes` marks `*.safetensors` for Git LFS; run `git lfs install`
+before committing the two weight files.
+
+## Quick start
+
+From the root of this model repository:
+
+```bash
+python examples/inference.py \
+  --input_h5ad /path/to/input.h5ad \
+  --output embeddings.npy \
+  --checkpoint checkpoints/scRep_20260625_30M
+```
+
+The output is a float32 NumPy matrix with one 768-dimensional, L2-normalized
+embedding per retained input cell. Input genes are matched by exact names to
+the bundled vocabulary. Cells with no matched non-zero genes are skipped.
+By default the script reads `adata.X`; add `--use_raw` to use `adata.raw.X`.
+
+The same workflow is available interactively in
+`examples/embedding_example.ipynb`.
+
+## Files
+
+```text
+scRep_release/
+├── README.md
+├── LICENSE                         
+├── requirements.txt
+├── scRep_inference.py
+├── scRep_pretrain/                 # Inference architecture implementation
+├── examples/
+│   ├── inference.py
+│   └── embedding_example.ipynb
+└── checkpoints/
+    ├── scRep_20260625_30M/
+    │   ├── config.json
+    │   ├── model.safetensors
+    │   └── tokenizer/{gene_vocab.json,gene_names.txt}
+    └── scRep_20260717_3M/
+        ├── config.json
+        ├── model.safetensors
+        └── tokenizer/{gene_vocab.json,gene_names.txt}
+```
+
+## Limitations and responsible use
+
+The embeddings are research outputs, not clinical measurements or diagnostic
+predictions. Validate performance, preprocessing, gene naming conventions,
+and potential batch/demographic bias for every downstream dataset before use.
+Do not use this model as the sole basis for medical decisions.
+
+## License
+
+The code and checkpoint weights in this release are distributed under the
+Apache License 2.0. See `LICENSE` for the complete terms.

checkpoints/scRep_20260625_30M/config.json

@@ -0,0 +1,8 @@
+{
+  "d_model": 768, "n_heads": 12, "dropout": 0.0, "encoder_layers": 12,
+  "gradient_checkpointing": false, "gene_vocab_size": 19240, "bin_vocab_size": 53,
+  "gene_pad_id": 0, "bin_pad_id": 0, "bin_mask_id": 1, "bin_offset": 3,
+  "gene_cls_id": 1, "expr_embedding_mode": "bin", "expr_rank_mlp_hidden_dim": 256,
+  "dino_out_dim": 4096, "dino_hidden_dim": 2048, "dino_bottleneck_dim": 256,
+  "inference_use_teacher": true
+}

checkpoints/scRep_20260625_30M/meta_vocab.json

@@ -0,0 +1,4 @@
+{
+  "meta_vocab_sizes": {},
+  "meta_maps": {}
+}

checkpoints/scRep_20260625_30M/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abc7eb419f785b6f545d27d0ecf5770663a505cdee36b657b72ac8c0f7138ec4
+size 915673940

checkpoints/scRep_20260717_3M/config.json

@@ -0,0 +1,8 @@
+{
+  "d_model": 768, "n_heads": 12, "dropout": 0.0, "encoder_layers": 12,
+  "gradient_checkpointing": false, "gene_vocab_size": 19239, "bin_vocab_size": 53,
+  "gene_pad_id": 0, "bin_pad_id": 0, "bin_mask_id": 1, "bin_offset": 3,
+  "gene_cls_id": 1, "expr_embedding_mode": "bin", "expr_rank_mlp_hidden_dim": 256,
+  "dino_out_dim": 4096, "dino_hidden_dim": 2048, "dino_bottleneck_dim": 256,
+  "inference_use_teacher": true
+}

checkpoints/scRep_20260717_3M/meta_vocab.json

@@ -0,0 +1,4 @@
+{
+  "meta_vocab_sizes": {},
+  "meta_maps": {}
+}

checkpoints/scRep_20260717_3M/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9204967d6e99165b69864b0b5b4c89cf14cf39407276d8e4e1312dadd807cf36
+size 915667796

examples/embedding_example.ipynb

@@ -0,0 +1,54 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# scRep embedding example\n",
+    "\n",
+    "Run this notebook from the root of the downloaded scRep model repository."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -r requirements.txt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "from types import SimpleNamespace\n",
+    "import numpy as np\n",
+    "from scRep_inference import encode_embeddings, load_examples, load_scRep_bundle\n",
+    "from scRep_pretrain.vocab import gene_vocab_to_map\n",
+    "\n",
+    "checkpoint = Path('checkpoints/scRep_20260625_30M')\n",
+    "input_h5ad = Path('/path/to/input.h5ad')  # change this\n",
+    "args = SimpleNamespace(\n",
+    "    model_dir=str(checkpoint), asset_dir='', device='cuda',\n",
+    "    batch_size=128, n_bins=50, max_input_genes=2048, backbone='auto',\n",
+    ")\n",
+    "model, gene_vocab, *_ = load_scRep_bundle(checkpoint, args)\n",
+    "gene_name_to_id = gene_vocab_to_map(gene_vocab)\n",
+    "examples, label_key = load_examples(input_h5ad, gene_name_to_id)\n",
+    "embeddings = encode_embeddings(model, examples, args)\n",
+    "np.save('embeddings.npy', embeddings)\n",
+    "print(embeddings.shape, label_key)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
+  "language_info": {"name": "python", "version": "3.10"}
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/inference.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+"""Extract L2-normalized scRep cell embeddings from an .h5ad file."""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+
+RELEASE_ROOT = Path(__file__).resolve().parents[1]
+if str(RELEASE_ROOT) not in sys.path:
+    sys.path.insert(0, str(RELEASE_ROOT))
+from scRep_inference import encode_embeddings, load_examples, load_scRep_bundle
+from scRep_pretrain.vocab import gene_vocab_to_map
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input_h5ad", required=True, help="Input AnnData (.h5ad) file.")
+    parser.add_argument("--output", default="embeddings.npy", help="Output .npy embedding matrix.")
+    parser.add_argument("--checkpoint", default="checkpoints/scRep_20260625_30M")
+    parser.add_argument("--device", default="cuda" if __import__("torch").cuda.is_available() else "cpu")
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--max_input_genes", type=int, default=2048)
+    parser.add_argument("--n_bins", type=int, default=50)
+    parser.add_argument("--max_cells", type=int, default=0, help="0 means all cells.")
+    parser.add_argument("--use_raw", action="store_true", help="Use adata.raw instead of adata.X.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    checkpoint = Path(args.checkpoint)
+    if not checkpoint.is_absolute():
+        checkpoint = RELEASE_ROOT / checkpoint
+    model, gene_vocab, *_ = load_scRep_bundle(checkpoint, SimpleNamespace(model_dir=str(checkpoint), asset_dir="", device=args.device))
+    gene_name_to_id = gene_vocab_to_map(gene_vocab)
+    examples, _ = load_examples(args.input_h5ad, gene_name_to_id, max_total_cells=args.max_cells, use_raw=args.use_raw)
+    if not examples:
+        raise ValueError("No input cells had genes present in the checkpoint vocabulary.")
+    embeddings = encode_embeddings(model, examples, args)
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.save(output, embeddings)
+    print(f"Saved {embeddings.shape[0]} x {embeddings.shape[1]} embeddings to {output}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,8 @@
+# Minimal dependencies for checkpoint inference examples (Python 3.10--3.12).
+# Install the PyTorch build appropriate for your CPU/CUDA platform first:
+# https://pytorch.org/get-started/locally/
+torch>=2.2,<2.10
+safetensors>=0.4,<1.0
+anndata>=0.10,<0.12
+numpy>=1.26,<2.0
+scipy>=1.11,<1.16

scRep_inference.py

@@ -0,0 +1,193 @@
+"""Standalone inference helpers for the public scRep release.
+
+This is the small, public equivalent of scDINO's experiment-side inference
+adapter.  It deliberately supports a checkpoint containing only weights by
+inferring the architecture and taking vocab assets from ``assets/``.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Sequence
+
+import anndata as ad
+import numpy as np
+import torch
+
+from scRep_pretrain.model import CellContrastivePredictModel, ContrastiveSplitCellCollator, FoundationConfig
+from scRep_pretrain.vocab import VocabSpec, gene_vocab_to_map
+
+ROOT = Path(__file__).resolve().parent
+LABEL_CANDIDATES = ["cell_type", "author_cell_type", "cell_type_ontology_term_id", "time_label"]
+BATCH_CANDIDATES = ["extended_batch_id", "batch_id", "batch", "sample_id", "sample", "donor_id"]
+
+
+@dataclass
+class EvalExample:
+    gene_ids: list[int]
+    expr: list[float]
+    label: str | None = None
+    source_file: str = ""
+    cell_idx: int = -1
+    batch_id: str | None = None
+
+
+def safe_name(value: Any) -> str:
+    result = "".join(c if c.isalnum() or c in "._-" else "-" for c in str(value)).strip("-")
+    return result or "item"
+
+
+def infer_model_cache_name(model_dir: str | Path) -> str:
+    path = Path(model_dir).expanduser().resolve()
+    return safe_name(f"{path.parent.name}_{path.name}" if path.name.startswith("checkpoint-") else path.name)
+
+
+def choose_label_key(explicit_key: str, available: Sequence[str]) -> str | None:
+    if explicit_key:
+        return explicit_key if explicit_key in set(available) else None
+    return next((key for key in LABEL_CANDIDATES if key in set(available)), None)
+
+
+def choose_batch_key(available: Sequence[str]) -> str | None:
+    return next((key for key in BATCH_CANDIDATES if key in set(available)), None)
+
+
+def list_data_files(data_dir: str | Path, data_files: Sequence[str | Path] = ()) -> list[Path]:
+    root = Path(data_dir).expanduser().resolve()
+    if data_files:
+        paths = [Path(p).expanduser() for p in data_files]
+        return [(p if p.is_absolute() else root / p).resolve() for p in paths]
+    paths = sorted(root.glob("*.h5ad")) if root.is_dir() else [root]
+    if not paths or not all(p.is_file() for p in paths):
+        raise FileNotFoundError(f"No .h5ad files found in: {root}")
+    return paths
+
+
+def _extract_row(matrix: Any, row: int) -> tuple[np.ndarray, np.ndarray]:
+    value = matrix[row]
+    if hasattr(value, "indices"):
+        return np.asarray(value.indices), np.asarray(value.data, dtype=np.float32)
+    dense = np.asarray(value).reshape(-1)
+    idx = np.flatnonzero(dense)
+    return idx, dense[idx].astype(np.float32)
+
+
+def _process_expr(values: np.ndarray, use_raw: bool, cp_target_sum: float) -> np.ndarray:
+    values = np.asarray(values, dtype=np.float32)
+    if not use_raw or not values.size:
+        return values
+    total = float(values.sum())
+    return np.log1p(values * (float(cp_target_sum) / total)) if total > 0 else np.empty(0, dtype=np.float32)
+
+
+def load_examples(data_dir: str | Path, gene_name_to_id: dict[str, int], *, max_cells_per_file: int = 0,
+                  max_total_cells: int = 0, seed: int = 42, backed: str = "r", use_raw: bool = False,
+                  cp_target_sum: float = 1e4, label_key: str = "") -> tuple[list[EvalExample], str | None]:
+    rng, result, chosen_label = np.random.default_rng(seed), [], None
+    for path in list_data_files(data_dir):
+        adata = ad.read_h5ad(path, backed=backed)
+        try:
+            matrix, names = (adata.raw.X, adata.raw.var_names) if use_raw else (adata.X, adata.var_names)
+            if matrix is None:
+                raise ValueError(f"Missing expression matrix: {path}")
+            if use_raw and adata.raw is None:
+                raise ValueError(f"--use_raw requested but raw is missing: {path}")
+            if chosen_label is None:
+                chosen_label = choose_label_key(label_key, list(adata.obs.columns))
+            local_to_global = np.fromiter((gene_name_to_id.get(str(g), 0) for g in names), dtype=np.int64, count=len(names))
+            indices = np.arange(adata.n_obs); rng.shuffle(indices)
+            if max_cells_per_file > 0:
+                indices = indices[:max_cells_per_file]
+            for row in indices:
+                idx, values = _extract_row(matrix, int(row)); values = _process_expr(values, use_raw, cp_target_sum)
+                gids = local_to_global[idx]; valid = gids > 0
+                if not np.any(valid) or not values.size:
+                    continue
+                label = str(adata.obs.iloc[int(row)][chosen_label]) if chosen_label else None
+                result.append(EvalExample(gids[valid].astype(int).tolist(), values[valid].astype(float).tolist(), label,
+                                          path.name, int(row), None))
+                if max_total_cells > 0 and len(result) >= max_total_cells:
+                    return result, chosen_label
+        finally:
+            if getattr(adata, "file", None) is not None:
+                adata.file.close()
+    return result, chosen_label
+
+
+def _config_from_state(state: dict[str, torch.Tensor], vocab: VocabSpec) -> FoundationConfig:
+    emb = state["student_backbone.tokenizer.gene_emb.weight"]
+    layers = len({key.split(".")[2] for key in state if key.startswith("student_backbone.encoder.") and key.endswith("norm1.scale")})
+    return FoundationConfig(d_model=int(emb.shape[1]), n_heads=12, encoder_layers=layers,
+        gene_vocab_size=int(emb.shape[0]), bin_vocab_size=int(state["student_backbone.tokenizer.bin_emb.weight"].shape[0]),
+        gene_pad_id=vocab.gene_pad_id, bin_pad_id=vocab.bin_pad_id, gene_cls_id=vocab.gene_cls_id,
+        dino_out_dim=int(state["student_head.prototypes.weight"].shape[0]),
+        dino_hidden_dim=int(state["student_head.mlp.0.weight"].shape[0]),
+        dino_bottleneck_dim=int(state["student_head.mlp.4.weight"].shape[0]))
+
+
+def load_scRep_bundle(model_dir_or_args: str | Path | Any, args: Any | None = None):
+    if args is None:
+        args, model_dir = model_dir_or_args, getattr(model_dir_or_args, "model_dir")
+    else:
+        model_dir = model_dir_or_args
+    weights_dir = Path(model_dir).expanduser().resolve()
+    from safetensors.torch import load_file
+    state = load_file(str(weights_dir / "model.safetensors"), device="cpu")
+    asset_arg = Path(str(getattr(args, "asset_dir", "") or ROOT / "assets")).expanduser()
+    candidates = [weights_dir, weights_dir / "tokenizer", weights_dir.parent, asset_arg, ROOT / "assets"]
+    def asset(name: str) -> Path:
+        found = next((p / name for p in candidates if (p / name).is_file()), None)
+        if found is None: raise FileNotFoundError(f"Missing {name}; pass --asset_dir")
+        return found
+    gene_vocab = json.loads(asset("gene_vocab.json").read_text())
+    meta_path, manifest_path = next((p / "meta_vocab.json" for p in candidates if (p / "meta_vocab.json").is_file()), None), next((p / "manifest.json" for p in candidates if (p / "manifest.json").is_file()), None)
+    vocab = VocabSpec()
+    config_path = next((candidate for p in candidates for candidate in (p / "config.json", p / "model_config.json") if candidate.is_file()), None)
+    cfg = FoundationConfig(**json.loads(config_path.read_text())) if config_path else _config_from_state(state, vocab)
+    model = CellContrastivePredictModel(cfg); model.load_state_dict(state, strict=True)
+    device = str(getattr(args, "device", "cpu")); model.to(device).eval()
+    return model, gene_vocab, json.loads(meta_path.read_text()) if meta_path else None, json.loads(manifest_path.read_text()) if manifest_path else None, {"model_family": "model", "weights_dir": str(weights_dir), "config_path": str(config_path or "")}
+
+
+def _select_input_genes(gene_ids: np.ndarray, expr: np.ndarray, max_input_genes: int):
+    if max_input_genes <= 0 or len(gene_ids) <= max_input_genes: return gene_ids, expr
+    order = np.argsort(-expr, kind="stable")[:max_input_genes]
+    return gene_ids[order], expr[order]
+
+
+def _make_eval_collator(vocab: VocabSpec, gene_vocab_size: int, n_bins: int, max_input_genes: int, model_family: str = "model", expr_embedding_mode: str = "bin"):
+    limit = max_input_genes if max_input_genes > 0 else gene_vocab_size
+    return ContrastiveSplitCellCollator(vocab=vocab, gene_vocab_size=gene_vocab_size, n_bins=n_bins, max_input_genes=max_input_genes,
+        shuffle_genes=False, teacher_top_genes=limit, teacher_num_views=1, teacher_global_ratio=1.0,
+        student_global_ratio_min=1.0, student_global_ratio_max=1.0, student_global_min_genes=limit,
+        student_local_min_genes=limit, student_local_max_genes=limit, student_num_local_views=0,
+        student_global_dropout_prob=0.0, student_local_dropout_prob=0.0, student_global_bin_mask_prob=0.0,
+        expr_embedding_mode=expr_embedding_mode)
+
+
+def build_eval_batch(examples: Sequence[EvalExample], *, vocab: VocabSpec, collator: Any, max_input_genes: int, use_bin: bool):
+    rows = []; width = 1
+    for ex in examples:
+        gids, expr = _select_input_genes(np.asarray(ex.gene_ids, dtype=np.int64), np.asarray(ex.expr, dtype=np.float32), max_input_genes)
+        order = np.argsort(-expr, kind="stable") if use_bin else np.argsort(gids, kind="stable")
+        gids, expr = gids[order], expr[order]
+        bins = collator._encode_expression_ids(expr) if use_bin else np.full(len(gids), vocab.bin_pad_id, dtype=np.int64)
+        rows.append((gids, bins)); width = max(width, len(gids) + 1)
+    gene = torch.full((len(rows), width), vocab.gene_pad_id, dtype=torch.long); bins = torch.full_like(gene, vocab.bin_pad_id); mask = torch.zeros_like(gene, dtype=torch.bool)
+    gene[:, 0] = vocab.gene_cls_id; mask[:, 0] = True
+    for i, (gids, vals) in enumerate(rows):
+        gene[i, 1:len(gids)+1] = torch.from_numpy(gids); bins[i, 1:len(gids)+1] = torch.from_numpy(vals); mask[i, 1:len(gids)+1] = True
+    return {"gene_ids": gene, "bin_ids": bins, "attention_mask": mask}
+
+
+def encode_embeddings(model: Any, examples: Sequence[EvalExample], args: Any, *, use_bin: bool = True, model_family: str = "model") -> np.ndarray:
+    vocab = VocabSpec(); collator = _make_eval_collator(vocab, model.cfg.gene_vocab_size, int(args.n_bins), int(args.max_input_genes), model_family, model.cfg.expr_embedding_mode)
+    backbone = None if getattr(args, "backbone", "auto") == "auto" else args.backbone; pieces = []
+    with torch.inference_mode():
+        for start in range(0, len(examples), int(args.batch_size)):
+            batch = build_eval_batch(examples[start:start+int(args.batch_size)], vocab=vocab, collator=collator, max_input_genes=int(args.max_input_genes), use_bin=use_bin)
+            states = model._encode_tokens(batch["gene_ids"].to(args.device), batch["bin_ids"].to(args.device), batch["attention_mask"].to(args.device), use_bin_embeddings=use_bin, backbone=backbone)
+            pieces.append(torch.nn.functional.normalize(states[:, 0], dim=-1).cpu())
+    return torch.cat(pieces).numpy() if pieces else np.empty((0, model.cfg.d_model), dtype=np.float32)

scRep_pretrain/__init__.py

@@ -0,0 +1,15 @@
+from .manifest import build_manifest_from_dir, save_manifest_json
+from .vocab import VocabSpec, build_gene_vocab_from_manifest, build_meta_vocab_from_manifest, gene_vocab_to_map
+from .dataset import MultiH5ADCellStream
+from . import model
+
+__all__ = [
+    "build_manifest_from_dir",
+    "save_manifest_json",
+    "VocabSpec",
+    "build_gene_vocab_from_manifest",
+    "build_meta_vocab_from_manifest",
+    "gene_vocab_to_map",
+    "MultiH5ADCellStream",
+    "model",
+]

scRep_pretrain/bin_cache.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import numpy as np
+
+
+def cache_key_for_h5ad(path: str) -> str:
+    resolved = str(Path(path).expanduser().resolve())
+    digest = hashlib.sha1(resolved.encode("utf-8")).hexdigest()[:12]
+    return f"{Path(path).stem}.{digest}"
+
+
+def value_bin_cache_paths(cache_dir: str, h5ad_path: str) -> Dict[str, Path]:
+    root = Path(cache_dir).expanduser().resolve()
+    key = cache_key_for_h5ad(h5ad_path)
+    return {
+        "indptr": root / f"{key}.indptr.npy",
+        "bin_ids": root / f"{key}.bin_ids.npy",
+        "meta": root / f"{key}.meta.json",
+    }
+
+
+def load_value_bin_cache(
+    cache_dir: str,
+    h5ad_path: str,
+    *,
+    expected_n_cells: Optional[int] = None,
+    mmap_mode: str = "r",
+) -> Optional[Dict[str, Any]]:
+    if not cache_dir:
+        return None
+    paths = value_bin_cache_paths(cache_dir, h5ad_path)
+    if not paths["indptr"].exists() or not paths["bin_ids"].exists():
+        return None
+
+    indptr = np.load(paths["indptr"], mmap_mode=mmap_mode)
+    bin_ids = np.load(paths["bin_ids"], mmap_mode=mmap_mode)
+    if expected_n_cells is not None and int(indptr.shape[0]) != int(expected_n_cells) + 1:
+        raise ValueError(
+            f"bin cache n_cells mismatch for {h5ad_path}: "
+            f"indptr_len={indptr.shape[0]} expected={int(expected_n_cells) + 1}"
+        )
+
+    meta: Dict[str, Any] = {}
+    if paths["meta"].exists():
+        with open(paths["meta"], "r", encoding="utf-8") as f:
+            meta = json.load(f)
+    return {"indptr": indptr, "bin_ids": bin_ids, "meta": meta, "paths": paths}

scRep_pretrain/dataset.py

@@ -0,0 +1,725 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from collections import OrderedDict
+import os
+
+import torch
+import anndata as ad
+import numpy as np
+import scipy.sparse as sp
+from torch.utils.data import IterableDataset, get_worker_info
+
+from .bin_cache import load_value_bin_cache
+
+
+# =============================================================================
+# Debug / Logging
+# =============================================================================
+DEBUG_PRINT: bool = False
+
+
+def _log(msg: str) -> None:
+    if DEBUG_PRINT:
+        print(f"[sc-token] {msg}")
+
+
+def _safe_close_anndata(adata_obj: Any) -> None:
+    f = getattr(adata_obj, "file", None)
+    if f is not None:
+        try:
+            f.close()
+        except Exception:
+            pass
+
+
+# =============================================================================
+# Per-file materialized reader
+# =============================================================================
+class H5ADFileReader:
+    """
+    Per-file reader.
+
+    Matrix semantics:
+      - use_raw=True  : adata.raw.X counts -> log1p(CP10K)
+      - use_raw=False : adata.X is assumed to already be log1p(CP10K)
+
+    Design notes:
+      - This reader always returns ALL observed nonzero genes for a cell.
+      - No top-k truncation happens here anymore.
+      - Top-k is now applied later inside the collator.
+      - Gene order is preserved here; any shuffling happens in the collator.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        dataset_id: int,
+        *,
+        gene_name_to_id: Dict[str, int],
+        meta_maps: Dict[str, Dict[str, int]],
+        meta_keys: List[str],
+        use_raw: bool,
+        cp_target_sum: float = 1e4,
+        backed: Optional[str] = "r",
+        prefer_csr: bool = True,
+        strict_gene: bool = True,
+        filter_missing_genes: bool = False,
+        value_bin_cache_dir: str = "",
+        require_value_bin_cache: bool = False,
+    ) -> None:
+        self.path = path
+        self.dataset_id = int(dataset_id)
+
+        self.gene_name_to_id = gene_name_to_id
+        self.meta_maps = meta_maps
+        self.meta_keys = list(meta_keys)
+
+        self.use_raw = bool(use_raw)
+        self.cp_target_sum = float(cp_target_sum)
+        self.prefer_csr = bool(prefer_csr)
+        self.filter_missing_genes = bool(filter_missing_genes)
+        self.value_bin_cache_dir = str(value_bin_cache_dir or "")
+        self.require_value_bin_cache = bool(require_value_bin_cache)
+        self._bin_cache_indptr = None
+        self._bin_cache_ids = None
+
+        _log(f"open: {path} (use_raw={self.use_raw}, backed={backed})")
+        self.adata = ad.read_h5ad(path, backed=backed)
+        self.n_cells = int(self.adata.n_obs)
+
+        if self.use_raw:
+            if (self.adata.raw is None) or (self.adata.raw.X is None):
+                raise ValueError(f"[{path}] use_raw=True but adata.raw.X missing.")
+            self.X = self.adata.raw.X
+            self.var_names = self.adata.raw.var_names
+        else:
+            if self.adata.X is None:
+                raise ValueError(f"[{path}] use_raw=False but adata.X missing.")
+            self.X = self.adata.X
+            self.var_names = self.adata.var_names
+
+        self.local_to_global = self._build_local_to_global(strict=strict_gene)
+        self.meta_int_cols = self._build_meta_int_cols()
+        self._load_value_bin_cache()
+
+        self._is_csr = sp.isspmatrix_csr(self.X)
+        if self._is_csr and self.prefer_csr:
+            self._csr_indptr = self.X.indptr
+            self._csr_indices = self.X.indices
+            self._csr_data = self.X.data
+
+        _log(f"[{path}] n_cells={self.n_cells}")
+
+    def close(self) -> None:
+        _safe_close_anndata(self.adata)
+
+    def _load_value_bin_cache(self) -> None:
+        if not self.value_bin_cache_dir:
+            return
+        cache = load_value_bin_cache(
+            self.value_bin_cache_dir,
+            self.path,
+            expected_n_cells=self.n_cells,
+            mmap_mode="r",
+        )
+        if cache is None:
+            if self.require_value_bin_cache:
+                raise FileNotFoundError(
+                    f"Missing value bin cache for {self.path} under {self.value_bin_cache_dir}"
+                )
+            return
+        self._bin_cache_indptr = cache["indptr"]
+        self._bin_cache_ids = cache["bin_ids"]
+
+    def has_value_bin_cache(self) -> bool:
+        return self._bin_cache_indptr is not None and self._bin_cache_ids is not None
+
+    def get_cached_bin_ids(self, cell_idx: int) -> Optional[np.ndarray]:
+        if not self.has_value_bin_cache():
+            return None
+        i = int(cell_idx)
+        start = int(self._bin_cache_indptr[i])
+        end = int(self._bin_cache_indptr[i + 1])
+        return np.asarray(self._bin_cache_ids[start:end], dtype=np.int64)
+
+    def _build_local_to_global(self, *, strict: bool) -> np.ndarray:
+        ids = np.fromiter(
+            (self.gene_name_to_id.get(str(g), 0) for g in list(self.var_names)),
+            dtype=np.int64,
+            count=len(self.var_names),
+        )
+        if strict and int(np.sum(ids == 0)) > 0:
+            missing = int(np.sum(ids == 0))
+            raise ValueError(
+                f"[{self.path}] {missing} genes missing in global gene vocab. "
+                "Rebuild vocab with full manifest or check gene naming consistency."
+            )
+        return ids
+
+    def _build_meta_int_cols(self) -> Dict[str, np.ndarray]:
+        out: Dict[str, np.ndarray] = {}
+        for k in self.meta_keys:
+            if k not in self.adata.obs.columns:
+                raise ValueError(f"[{self.path}] obs has no column '{k}'")
+            if k not in self.meta_maps:
+                raise ValueError(f"meta_maps missing key='{k}'. Build meta vocabs first.")
+
+            gmap = self.meta_maps[k]
+            col = self.adata.obs[k].astype("category")
+            cats = col.cat.categories.astype(str).tolist()
+            codes = col.cat.codes.to_numpy()
+
+            cat_ids = np.array([gmap.get(c, 0) for c in cats], dtype=np.int32)
+            ids = np.zeros((self.n_cells,), dtype=np.int32)
+            valid = codes >= 0
+            ids[valid] = cat_ids[codes[valid]]
+            out[k] = ids
+        return out
+
+    def _csr_row_indices_data(self, cell_idx: int) -> Tuple[np.ndarray, np.ndarray]:
+        i = int(cell_idx)
+        start = int(self._csr_indptr[i])
+        end = int(self._csr_indptr[i + 1])
+        return self._csr_indices[start:end], self._csr_data[start:end]
+
+    def _row_indices_data(self, cell_idx: int) -> Tuple[np.ndarray, np.ndarray]:
+        if self._is_csr and self.prefer_csr:
+            idx, val = self._csr_row_indices_data(cell_idx)
+        else:
+            row = self.X[int(cell_idx)]
+            if sp.issparse(row):
+                if not sp.isspmatrix_csr(row):
+                    row = row.tocsr()
+                idx, val = row.indices, row.data
+            else:
+                dense = np.asarray(row).ravel()
+                idx = np.flatnonzero(dense)
+                val = dense[idx]
+
+        if val.size > 0:
+            keep = val != 0
+            if not np.all(keep):
+                idx = idx[keep]
+                val = val[keep]
+
+        return idx, val
+
+    def get_cell_tokens(self, cell_idx: int) -> Tuple[List[int], List[float]]:
+        idx, val = self._row_indices_data(cell_idx)
+
+        if idx.size == 0:
+            return [], []
+
+        row_sum = float(val.sum())
+        if self.use_raw:
+            if row_sum <= 0:
+                return [], []
+            expr = np.log1p(val.astype(np.float32) * (self.cp_target_sum / row_sum))
+        else:
+            expr = val.astype(np.float32)
+
+        gene_ids = self.local_to_global[idx.astype(np.int64)].astype(np.int64)
+        if self.filter_missing_genes:
+            keep = gene_ids > 0
+            if not np.all(keep):
+                gene_ids = gene_ids[keep]
+                expr = expr[keep]
+            if gene_ids.size == 0:
+                return [], []
+        return gene_ids.tolist(), expr.tolist()
+
+    def get_cell_observation_stats(self, cell_idx: int) -> Tuple[float, int]:
+        idx, val = self._row_indices_data(cell_idx)
+        detected_genes = int(idx.size)
+        if self.use_raw:
+            return float(val.sum()), detected_genes
+
+        obs = self.adata.obs
+        row = obs.iloc[int(cell_idx)]
+
+        total_counts = None
+        for key in ("total_counts", "n_counts"):
+            if key not in obs.columns:
+                continue
+            candidate = float(row[key])
+            if np.isfinite(candidate):
+                total_counts = candidate
+                break
+        if total_counts is None:
+            total_counts = float(detected_genes)
+
+        for key in ("n_genes_by_counts", "n_genes"):
+            if key not in obs.columns:
+                continue
+            candidate = float(row[key])
+            if np.isfinite(candidate):
+                detected_genes = int(candidate)
+                break
+        return total_counts, detected_genes
+
+    def get_meta_batch(self, cell_ids: np.ndarray, *, include_dataset_id: bool = False) -> Dict[str, List[int]]:
+        ids = cell_ids.astype(np.int64)
+        out: Dict[str, List[int]] = {}
+
+        if include_dataset_id:
+            out["dataset_id"] = [self.dataset_id + 1] * int(ids.size)
+
+        for k in self.meta_keys:
+            out[k] = self.meta_int_cols[k][ids].astype(np.int64).tolist()
+
+        return out
+
+
+# =============================================================================
+# Global pseudo-batch style index metadata
+# =============================================================================
+@dataclass
+class _PseudoBatchGroupIndex:
+    """
+    One synthetic group inside one file.
+
+    For scRep we keep one group containing all cells in the file.
+    """
+
+    raw_cell_ids: np.ndarray   # int32
+    n_batches: int
+    usable_cells: int
+
+
+@dataclass
+class _FileIndex:
+    path: str
+    dataset_id: int
+    groups: List[_PseudoBatchGroupIndex]
+
+
+# =============================================================================
+# Small LRU reader cache
+# =============================================================================
+class _ReaderCache:
+    """
+    Small LRU cache of open H5ADFileReader objects.
+
+    Each worker / rank gets its own cache inside __iter__.
+    """
+
+    def __init__(
+        self,
+        *,
+        max_open_files: int,
+        gene_name_to_id: Dict[str, int],
+        meta_maps: Dict[str, Dict[str, int]],
+        meta_keys: List[str],
+        use_raw: bool,
+        cp_target_sum: float,
+        backed: Optional[str],
+        prefer_csr: bool,
+        strict_gene: bool,
+        filter_missing_genes: bool,
+        value_bin_cache_dir: str,
+        require_value_bin_cache: bool,
+    ) -> None:
+        self.max_open_files = int(max_open_files)
+        self.gene_name_to_id = gene_name_to_id
+        self.meta_maps = meta_maps
+        self.meta_keys = meta_keys
+        self.use_raw = bool(use_raw)
+        self.cp_target_sum = float(cp_target_sum)
+        self.backed = backed
+        self.prefer_csr = bool(prefer_csr)
+        self.strict_gene = bool(strict_gene)
+        self.filter_missing_genes = bool(filter_missing_genes)
+        self.value_bin_cache_dir = str(value_bin_cache_dir or "")
+        self.require_value_bin_cache = bool(require_value_bin_cache)
+
+        self._cache: "OrderedDict[int, H5ADFileReader]" = OrderedDict()
+
+    def get(self, file_id: int, file_meta: _FileIndex) -> H5ADFileReader:
+        if file_id in self._cache:
+            reader = self._cache.pop(file_id)
+            self._cache[file_id] = reader
+            return reader
+
+        reader = H5ADFileReader(
+            path=file_meta.path,
+            dataset_id=file_meta.dataset_id,
+            gene_name_to_id=self.gene_name_to_id,
+            meta_maps=self.meta_maps,
+            meta_keys=self.meta_keys,
+            use_raw=self.use_raw,
+            cp_target_sum=self.cp_target_sum,
+            backed=self.backed,
+            prefer_csr=self.prefer_csr,
+            strict_gene=self.strict_gene,
+            filter_missing_genes=self.filter_missing_genes,
+            value_bin_cache_dir=self.value_bin_cache_dir,
+            require_value_bin_cache=self.require_value_bin_cache,
+        )
+
+        self._cache[file_id] = reader
+
+        while len(self._cache) > self.max_open_files:
+            _, old_reader = self._cache.popitem(last=False)
+            try:
+                old_reader.close()
+            except Exception:
+                pass
+
+        return reader
+
+    def close_all(self) -> None:
+        for _, reader in self._cache.items():
+            try:
+                reader.close()
+            except Exception:
+                pass
+        self._cache.clear()
+
+
+# =============================================================================
+# Multi-file single-cell stream
+# =============================================================================
+class MultiH5ADCellStream(IterableDataset):
+    """
+    Globally shuffled single-cell stream.
+
+    Changes from the old implementation:
+    - Build a static global cell-slot index once.
+    - Sample slots by a global permutation, so cells are drawn with weight
+      proportional to file size (number of cells), not by equal-file coin flips.
+    - Still materializes one cell per yielded example.
+    """
+
+    def __init__(
+        self,
+        manifest: List[Dict[str, Any]],
+        *,
+        gene_name_to_id: Dict[str, int],
+        meta_maps: Dict[str, Dict[str, int]],
+        meta_keys: Optional[List[str]] = None,
+        buffer_size: int = 16,
+        seed: int = 42,
+        use_raw: bool = True,
+        cp_target_sum: float = 1e4,
+        backed: Optional[str] = "r",
+        prefer_csr: bool = True,
+        sort_cell_ids_for_io: bool = False,
+        include_dataset_id: bool = False,
+        skip_bad_files: bool = False,
+        filter_missing_genes: bool = False,
+        value_bin_cache_dir: str = "",
+        require_value_bin_cache: bool = False,
+    ) -> None:
+        self.manifest = list(manifest)
+        self.gene_name_to_id = gene_name_to_id
+        self.meta_maps = meta_maps
+
+        self.meta_keys = list(meta_keys or [])
+        self.full_meta_keys = list(self.meta_keys)
+
+        # Kept for compatibility with existing call sites / configs.
+        self.buffer_size = int(buffer_size)
+        self.seed = int(seed)
+
+        # Equivalent to sampling one-cell pseudo-batches.
+        self.N = 1
+
+        self.use_raw = bool(use_raw)
+        self.cp_target_sum = float(cp_target_sum)
+        self.backed = backed
+        self.prefer_csr = bool(prefer_csr)
+        self.filter_missing_genes = bool(filter_missing_genes)
+        self.strict_gene = not self.filter_missing_genes
+
+        self.sort_cell_ids_for_io = bool(sort_cell_ids_for_io)
+        self.include_dataset_id = bool(include_dataset_id)
+        self.skip_bad_files = bool(skip_bad_files)
+        self.value_bin_cache_dir = str(value_bin_cache_dir or "")
+        self.require_value_bin_cache = bool(require_value_bin_cache)
+
+        # reader cache
+        self.max_open_files = 1024
+
+        self.cycle = True
+        self._iter_id = 0
+
+        if len(self.manifest) == 0:
+            raise RuntimeError("Empty manifest.")
+
+        # Built once
+        self.file_index: List[_FileIndex] = []
+        # shape: (num_global_pseudo_batches, 3), columns = [file_id, group_id, batch_id]
+        self.batch_refs: np.ndarray = np.empty((0, 3), dtype=np.int32)
+
+        self._build_global_index()
+
+    # -------------------------------------------------------------------------
+    # Global index construction
+    # -------------------------------------------------------------------------
+    def _scan_one_file(self, item: Dict[str, Any], file_id: int) -> Optional[_FileIndex]:
+        path = item["path"]
+        dataset_id = int(item.get("dataset_id", 0))
+
+        adata = None
+        try:
+            adata = ad.read_h5ad(path, backed=self.backed)
+            n_cells = int(adata.n_obs)
+            if n_cells < self.N:
+                _log(f"skip file (n_cells < N): {path}")
+                return None
+
+            # Keep failure mode aligned with old stream: invalid/missing columns fail at scan
+            # rather than mid-iteration.
+            for k in self.full_meta_keys:
+                if k not in adata.obs.columns:
+                    raise ValueError(f"[{path}] obs has no required column '{k}'.")
+
+            raw_cells = np.arange(n_cells, dtype=np.int32)
+            n_batches = int(raw_cells.size // self.N)
+            if n_batches <= 0:
+                _log(f"skip file (no full pseudo-batches): {path}")
+                return None
+
+            usable_cells = int(n_batches * self.N)
+            group = _PseudoBatchGroupIndex(
+                raw_cell_ids=raw_cells,
+                n_batches=n_batches,
+                usable_cells=usable_cells,
+            )
+            return _FileIndex(path=path, dataset_id=dataset_id, groups=[group])
+        finally:
+            if adata is not None:
+                _safe_close_anndata(adata)
+
+    def _build_global_index(self) -> None:
+        file_index: List[_FileIndex] = []
+        batch_refs: List[Tuple[int, int, int]] = []
+
+        for src_file_id, item in enumerate(self.manifest):
+            try:
+                fidx = self._scan_one_file(item, src_file_id)
+            except Exception as e:
+                if self.skip_bad_files:
+                    _log(f"skip bad file: {item.get('path')} | {e}")
+                    continue
+                raise
+
+            if fidx is None:
+                continue
+
+            new_file_id = len(file_index)
+            file_index.append(fidx)
+
+            for group_id, g in enumerate(fidx.groups):
+                for batch_id in range(int(g.n_batches)):
+                    batch_refs.append((new_file_id, group_id, batch_id))
+
+        if len(file_index) == 0 or len(batch_refs) == 0:
+            raise RuntimeError(
+                "No eligible cells found. "
+                "Check manifest, requested meta keys, and skip_bad_files setting."
+            )
+
+        self.file_index = file_index
+        self.batch_refs = np.asarray(batch_refs, dtype=np.int32)
+
+        total_files = len(self.file_index)
+        total_groups = int(sum(len(f.groups) for f in self.file_index))
+        total_batches = int(self.batch_refs.shape[0])
+        total_cells_usable = int(sum(g.usable_cells for f in self.file_index for g in f.groups))
+
+        _log(
+            f"global index built: files={total_files}, groups={total_groups}, "
+            f"slots={total_batches}, usable_cells={total_cells_usable}"
+        )
+
+    # -------------------------------------------------------------------------
+    # Introspection helpers
+    # -------------------------------------------------------------------------
+    @property
+    def num_global_pseudo_batches(self) -> int:
+        return int(self.batch_refs.shape[0])
+
+    def estimate_batches_per_shard(self, num_shards: int) -> int:
+        if num_shards <= 0:
+            return self.num_global_pseudo_batches
+        return (self.num_global_pseudo_batches + num_shards - 1) // num_shards
+
+    # -------------------------------------------------------------------------
+    # Per-epoch group reshuffle
+    # -------------------------------------------------------------------------
+    def _build_epoch_group_batches(self, epoch_seed: int) -> List[List[np.ndarray]]:
+        """
+        For the current epoch/pass, rebuild each file/group by:
+          - reshuffling cell ids
+          - dropping tail
+          - reshaping to (n_batches, N)
+
+        Returns:
+          epoch_group_batches[file_id][group_id] = np.ndarray of shape (n_batches, N)
+        """
+        epoch_group_batches: List[List[np.ndarray]] = []
+
+        for file_id, fidx in enumerate(self.file_index):
+            file_batches: List[np.ndarray] = []
+
+            for group_id, gidx in enumerate(fidx.groups):
+                group_seed = (
+                    int(epoch_seed)
+                    + 1000003 * int(file_id + 1)
+                    + 9176 * int(group_id + 1)
+                ) % (2**63 - 1)
+
+                rng = np.random.default_rng(group_seed)
+
+                raw = gidx.raw_cell_ids
+                perm = rng.permutation(raw.size)
+                shuffled = raw[perm]
+
+                usable = shuffled[: gidx.usable_cells]
+                batch_cells = usable.reshape(gidx.n_batches, self.N).astype(np.int32, copy=False)
+                file_batches.append(batch_cells)
+
+            epoch_group_batches.append(file_batches)
+
+        return epoch_group_batches
+
+    # -------------------------------------------------------------------------
+    # Materialization
+    # -------------------------------------------------------------------------
+    def _materialize_one_batch(
+        self,
+        *,
+        file_id: int,
+        group_id: int,
+        batch_id: int,
+        epoch_group_batches: List[List[np.ndarray]],
+        cache: _ReaderCache,
+    ) -> Dict[str, Any]:
+        fmeta = self.file_index[file_id]
+        cell_ids = epoch_group_batches[file_id][group_id][batch_id]
+
+        if self.sort_cell_ids_for_io:
+            cell_ids = np.sort(cell_ids)
+
+        reader = cache.get(file_id, fmeta)
+
+        meta_acc = reader.get_meta_batch(
+            cell_ids,
+            include_dataset_id=self.include_dataset_id,
+        )
+
+        cid = int(cell_ids[0])
+        gene_ids, expr = reader.get_cell_tokens(cid)
+        total_counts, n_genes_by_counts = reader.get_cell_observation_stats(cid)
+        out = {
+            "gene_ids": gene_ids,
+            "expr": expr,
+            "meta": meta_acc,
+            "total_counts": float(total_counts),
+            "n_genes_by_counts": int(n_genes_by_counts),
+        }
+        cached_bin_ids = reader.get_cached_bin_ids(cid)
+        if cached_bin_ids is not None:
+            if int(cached_bin_ids.shape[0]) != len(gene_ids):
+                raise ValueError(
+                    f"Cached bin length mismatch for {reader.path} cell={cid}: "
+                    f"bin_ids={cached_bin_ids.shape[0]} genes={len(gene_ids)}"
+                )
+            out["bin_ids"] = cached_bin_ids.astype(np.int64, copy=False).tolist()
+        return out
+
+    # -------------------------------------------------------------------------
+    # Iteration
+    # -------------------------------------------------------------------------
+    def __iter__(self):
+        # -------------------------
+        # rank / world
+        # -------------------------
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            rank = torch.distributed.get_rank()
+            world_size = torch.distributed.get_world_size()
+        else:
+            rank = int(os.environ.get("RANK", "0"))
+            world_size = int(os.environ.get("WORLD_SIZE", "1"))
+
+        # -------------------------
+        # worker info
+        # -------------------------
+        info = get_worker_info()
+        if info is None:
+            worker_id, num_workers = 0, 1
+        else:
+            worker_id, num_workers = info.id, info.num_workers
+
+        shard_id = rank * num_workers + worker_id
+        num_shards = world_size * num_workers
+
+        iter_id = getattr(self, "_iter_id", 0)
+        self._iter_id = iter_id + 1
+
+        base_seed = int(self.seed)
+        mix = (base_seed + 10000019 * iter_id) % (2**63 - 1)
+        rng = np.random.default_rng(mix)
+
+        if DEBUG_PRINT:
+            _log(
+                f"rank={rank}/{world_size} worker={worker_id}/{num_workers} "
+                f"iter_id={iter_id} seed={mix} shard={shard_id}/{num_shards} "
+                f"global_pseudo_batches={self.num_global_pseudo_batches}"
+            )
+
+        cache = _ReaderCache(
+            max_open_files=self.max_open_files,
+            gene_name_to_id=self.gene_name_to_id,
+            meta_maps=self.meta_maps,
+            meta_keys=self.full_meta_keys,
+            use_raw=self.use_raw,
+            cp_target_sum=self.cp_target_sum,
+            backed=self.backed,
+            prefer_csr=self.prefer_csr,
+            strict_gene=self.strict_gene,
+            filter_missing_genes=self.filter_missing_genes,
+            value_bin_cache_dir=self.value_bin_cache_dir,
+            require_value_bin_cache=self.require_value_bin_cache,
+        )
+
+        try:
+            epoch_idx = 0
+            while True:
+                epoch_seed = (mix + 1000003 * epoch_idx) % (2**63 - 1)
+
+                # 1) Re-shuffle every group for this epoch
+                epoch_group_batches = self._build_epoch_group_batches(epoch_seed)
+
+                # 2) Global shuffle over static slot table
+                global_order = rng.permutation(self.num_global_pseudo_batches)
+
+                # 3) Shard by rank * worker
+                local_order = global_order[shard_id::num_shards]
+
+                if DEBUG_PRINT:
+                    _log(
+                        f"new epoch/pass: epoch_idx={epoch_idx}, total={self.num_global_pseudo_batches}, "
+                        f"local={len(local_order)}, shard={shard_id}/{num_shards}"
+                    )
+
+                for ref_idx in local_order.tolist():
+                    file_id, group_id, batch_id = self.batch_refs[int(ref_idx)].tolist()
+                    yield self._materialize_one_batch(
+                        file_id=int(file_id),
+                        group_id=int(group_id),
+                        batch_id=int(batch_id),
+                        epoch_group_batches=epoch_group_batches,
+                        cache=cache,
+                    )
+
+                epoch_idx += 1
+                if not self.cycle:
+                    break
+
+        finally:
+            cache.close_all()

scRep_pretrain/manifest.py

@@ -0,0 +1,197 @@
+import json
+import time
+import logging
+import anndata as ad
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+
+# =============================================================================
+# Debug / Logging
+# =============================================================================
+DEBUG_PRINT: bool = False
+logger = logging.getLogger("scRep_train")
+
+
+def _log(msg: str) -> None:
+    if DEBUG_PRINT:
+        print(f"[manifest] {msg}")
+
+
+# =============================================================================
+# Utils
+# =============================================================================
+def _safe_close_anndata(adata_obj: Any) -> None:
+    """Close backed AnnData file handle if present."""
+    f = getattr(adata_obj, "file", None)
+    if f is not None:
+        try:
+            f.close()
+        except Exception:
+            pass
+
+
+def save_manifest_json(manifest: List[Dict[str, Any]], out_path: str) -> None:
+    """Save manifest to json."""
+    _log(f"Save -> {out_path}")
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump(manifest, f, ensure_ascii=False, indent=2)
+    _log(f"Saved items: {len(manifest)}")
+
+
+def _obs_nunique_and_topk(
+    adata_obj: Any,
+    col_name: str,
+    *,
+    topk: int = 0,
+) -> Tuple[int, Optional[List[Tuple[str, int]]]]:
+    """
+    Compute nunique (dropna) for one obs column, optionally return top-k value counts.
+    """
+    if col_name not in adata_obj.obs.columns:
+        return 0, None
+
+    s = adata_obj.obs[col_name]
+
+    try:
+        n_uniq = int(s.nunique(dropna=True))
+    except Exception:
+        s_cat = s.astype("category")
+        n_uniq = int(len(s_cat.cat.categories))
+
+    top_items: Optional[List[Tuple[str, int]]] = None
+    if topk and topk > 0:
+        try:
+            vc = s.value_counts(dropna=True)
+            top_items = [(str(idx), int(cnt)) for idx, cnt in vc.head(int(topk)).items()]
+        except Exception:
+            top_items = None
+
+    return n_uniq, top_items
+
+
+# =============================================================================
+# Manifest builder
+# =============================================================================
+def build_manifest_from_dir(
+    dir_path: str,
+    *,
+    pattern: str = "*.h5ad",
+    backed: Optional[str] = "r",
+    use_raw: bool = True,
+    required_obs_cols: Optional[List[str]] = None,
+    obs_count_cols: Optional[List[str]] = None,
+    include_topk: int = 0,
+    min_cells: int = 1,
+    max_files: Optional[int] = None,
+    print_every: int = 50,
+) -> List[Dict[str, Any]]:
+    """
+    Scan a directory and build a stable manifest.
+
+    Stored per file:
+      - path, name, dataset_id, n_cells
+      - obs_nunique: Dict[col, nunique]
+      - obs_topk: Dict[col, List[(value, count)]]   # optional
+    """
+    required_obs_cols = list(required_obs_cols or [])
+    obs_count_cols = list(obs_count_cols or [])
+
+    root = Path(dir_path)
+    files = sorted(root.glob(pattern))
+    if max_files is not None:
+        files = files[: int(max_files)]
+
+    _log(f"Scan dir: {dir_path}")
+    _log(f"Found files: {len(files)} (pattern={pattern})")
+    _log(f"Checks: use_raw={use_raw}, min_cells={min_cells}")
+    _log(f"Required obs cols: {required_obs_cols}")
+    _log(f"Count obs cols: {obs_count_cols} (topk={include_topk})")
+    logger.info(
+        "manifest scan start: dir=%s files=%s use_raw=%s min_cells=%s required_obs_cols=%s",
+        dir_path, len(files), use_raw, min_cells, required_obs_cols,
+    )
+
+    manifest: List[Dict[str, Any]] = []
+    n_skipped = 0
+    n_failed = 0
+    start_time = time.time()
+
+    for i, p in enumerate(files):
+        path = str(p)
+
+        if (i + 1) % max(1, int(print_every)) == 0:
+            elapsed = time.time() - start_time
+            _log(
+                f"Progress: {i + 1}/{len(files)} | "
+                f"kept={len(manifest)} skipped={n_skipped} failed={n_failed}"
+            )
+            logger.info(
+                "manifest scan progress: %s/%s kept=%s skipped=%s failed=%s elapsed=%.1fs",
+                i + 1, len(files), len(manifest), n_skipped, n_failed, elapsed,
+            )
+
+        adata_obj = None
+        try:
+            adata_obj = ad.read_h5ad(path, backed=backed)
+
+            n_obs = int(adata_obj.n_obs)
+            if n_obs < int(min_cells):
+                n_skipped += 1
+                continue
+
+            if use_raw:
+                has_raw = (adata_obj.raw is not None) and (getattr(adata_obj.raw, "X", None) is not None)
+                if not has_raw:
+                    n_skipped += 1
+                    continue
+            else:
+                has_x = getattr(adata_obj, "X", None) is not None
+                if not has_x:
+                    n_skipped += 1
+                    continue
+
+            obs_cols = set(adata_obj.obs.columns.tolist())
+            if any(c not in obs_cols for c in required_obs_cols):
+                n_skipped += 1
+                continue
+
+            obs_nunique: Dict[str, int] = {}
+            obs_topk: Dict[str, List[Tuple[str, int]]] = {}
+
+            for col in obs_count_cols:
+                n_uniq, top_items = _obs_nunique_and_topk(adata_obj, col, topk=include_topk)
+                obs_nunique[col] = int(n_uniq)
+                if include_topk and top_items is not None:
+                    obs_topk[col] = top_items
+
+            dataset_id = len(manifest)
+
+            item: Dict[str, Any] = {
+                "path": path,
+                "name": p.name,
+                "dataset_id": int(dataset_id),
+                "n_cells": int(n_obs),
+                "obs_nunique": obs_nunique,
+            }
+            if include_topk and len(obs_topk) > 0:
+                item["obs_topk"] = obs_topk
+
+            manifest.append(item)
+
+        except Exception as e:
+            n_failed += 1
+            _log(f"Failed: {path} | {type(e).__name__}: {e}")
+        finally:
+            if adata_obj is not None:
+                _safe_close_anndata(adata_obj)
+
+    elapsed = time.time() - start_time
+    _log(f"Done. Kept={len(manifest)}, skipped={n_skipped}, failed={n_failed}")
+    logger.info(
+        "manifest scan done: kept=%s skipped=%s failed=%s elapsed=%.1fs",
+        len(manifest), n_skipped, n_failed, elapsed,
+    )
+    if len(manifest) > 0:
+        _log(f"Example item:\n{json.dumps(manifest[0], ensure_ascii=False, indent=2)[:2000]}")
+    return manifest

scRep_pretrain/model/__init__.py

@@ -0,0 +1,10 @@
+from .foundation_model import CellContrastivePredictModel, FoundationConfig
+from .contrastive_collator import ContrastiveSplitCellCollator
+from .contrastive_trainer import ContrastiveCellTrainer
+
+__all__ = [
+    "CellContrastivePredictModel",
+    "FoundationConfig",
+    "ContrastiveSplitCellCollator",
+    "ContrastiveCellTrainer",
+]

scRep_pretrain/model/contrastive_collator.py

@@ -0,0 +1,981 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch.utils.data import get_worker_info
+from typing import Any, Dict, List, Optional, Tuple
+
+
+def _distributed_rank_world() -> Tuple[int, int]:
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        return int(torch.distributed.get_rank()), int(torch.distributed.get_world_size())
+    try:
+        rank = int(__import__("os").environ.get("RANK", "0"))
+        world = int(__import__("os").environ.get("WORLD_SIZE", "1"))
+    except Exception:
+        return 0, 1
+    return max(0, rank), max(1, world)
+
+
+PlanStep = tuple[str, int, int, int]
+
+
+def _split_target_count(count: int, target: float, bins_left: int, reserve: int, split_factor: float) -> int:
+    max_split = min(int(count), int(bins_left) - int(reserve))
+    if max_split < 2 or float(count) < float(split_factor) * float(target):
+        return 1
+    denom = max(float(target), 1.0)
+    return max(2, min(int(round(float(count) / denom)), int(max_split)))
+
+
+def _plan_value_aware_lookahead_bins(
+    counts: np.ndarray,
+    n_bins: int,
+    split_factor: float = 1.2,
+) -> list[PlanStep]:
+    counts_i = np.asarray(counts, dtype=np.int64)
+    num_values = int(counts_i.shape[0])
+    total = int(counts_i.sum())
+    if num_values == 0 or total <= 0:
+        return []
+
+    k_eff = max(1, min(int(n_bins), total))
+    prefix = np.zeros(num_values + 1, dtype=np.int64)
+    prefix[1:] = np.cumsum(counts_i, dtype=np.int64)
+
+    plan: list[PlanStep] = []
+    i = 0
+    bins_left = int(k_eff)
+
+    while i < num_values and bins_left > 0:
+        remaining_values = int(num_values - i)
+        remaining_total = int(prefix[num_values] - prefix[i])
+        if remaining_total <= 0:
+            break
+
+        if bins_left == 1:
+            plan.append(("merge", int(i), int(num_values - 1), 1))
+            break
+
+        if remaining_values == 1:
+            count = int(counts_i[i])
+            bins_used = min(int(bins_left), count)
+            if bins_used >= 2:
+                plan.append(("split", int(i), int(i), int(bins_used)))
+            else:
+                plan.append(("merge", int(i), int(i), 1))
+            break
+
+        if bins_left > remaining_values:
+            count = int(counts_i[i])
+            needed_extra = int(bins_left - remaining_values)
+            bins_used = min(count, needed_extra + 1)
+            if bins_used >= 2:
+                plan.append(("split", int(i), int(i), int(bins_used)))
+                bins_left -= int(bins_used)
+                i += 1
+                continue
+
+        target = float(remaining_total) / float(bins_left)
+        reserve = 1 if (i + 1) < num_values else 0
+        split_bins = _split_target_count(int(counts_i[i]), target, bins_left, reserve, float(split_factor))
+        if split_bins >= 2:
+            plan.append(("split", int(i), int(i), int(split_bins)))
+            bins_left -= int(split_bins)
+            i += 1
+            continue
+
+        max_end = int(num_values - max(1, bins_left - 1) - 1)
+        max_end = max(int(i), max_end)
+        best_end = int(i)
+        best_cost = (float("inf"), float("inf"), float("inf"))
+        current_size = 0
+        worse_after_target = 0
+
+        for end in range(int(i), max_end + 1):
+            current_size += int(counts_i[end])
+            rest = int(prefix[num_values] - prefix[end + 1])
+            rest_bins = max(1, int(bins_left) - 1)
+            rest_target = float(rest) / float(rest_bins)
+            current_dev = abs(float(current_size) - target)
+            rest_dev = abs(rest_target - target)
+            cost = (max(current_dev, rest_dev), current_dev, rest_dev)
+            if cost < best_cost:
+                best_cost = cost
+                best_end = int(end)
+                worse_after_target = 0
+            elif current_size >= target:
+                worse_after_target += 1
+                if worse_after_target >= 2:
+                    break
+
+        plan.append(("merge", int(i), int(best_end), 1))
+        bins_left -= 1
+        i = int(best_end) + 1
+
+    return plan
+
+
+class ContrastiveSplitCellCollator:
+    """
+    DINO-like scRNA batch builder.
+
+    Per cell:
+      - teacher views:
+          1) global gene+equal-bin view sampled at teacher_global_ratio from top teacher_top_genes genes
+          2) another global gene+equal-bin view sampled at teacher_global_ratio from top teacher_top_genes genes
+      - student views:
+          1) global gene+equal-bin subset with masked bin ids
+          2) global gene+equal-bin subset
+          3) local gene+equal-bin subset
+          4) local gene+equal-bin subset
+    """
+
+    def __init__(
+        self,
+        *,
+        vocab: Any,
+        gene_vocab_size: int,
+        n_bins: int = 50,
+        max_input_genes: int = 5120,
+        shuffle_genes: bool = True,
+        seed: int = 42,
+        teacher_top_genes: int = 2048,
+        teacher_num_views: int = 2,
+        teacher_global_ratio: float = 0.95,
+        student_global_ratio_min: float = 0.60,
+        student_global_ratio_max: float = 0.85,
+        student_global_min_genes: int = 1000,
+        student_local_min_genes: int = 512,
+        student_local_max_genes: int = 768,
+        student_num_local_views: int = 2,
+        student_global_dropout_prob: float = 0.10,
+        student_local_dropout_prob: float = 0.20,
+        student_global_bin_mask_prob: float = 0.35,
+        student_global_bin_mask_before_binning: bool = False,
+        student_high_expr_global_view: str = "none",
+        student_high_expr_local_view: str = "none",
+        student_gene_only_local_view: str = "none",
+        student_global_mask_mode: str = "current",
+        student_global_model_masking: bool = False,
+        student_global_mask_sample_probability: float = 0.5,
+        student_global_mask_ratio_min: float = 0.1,
+        student_global_mask_ratio_max: float = 0.5,
+        include_teacher_cls_decoder_targets: bool = False,
+        teacher_cls_decoder_target_mode: str = "top_other_zero",
+        teacher_cls_decoder_other_expr_genes: int = 512,
+        teacher_cls_decoder_zero_genes: int = 512,
+        plan_mode: str = "lookahead_split_1.2",
+        expr_embedding_mode: str = "bin",
+    ) -> None:
+        self.vocab = vocab
+        self.gene_vocab_size = int(gene_vocab_size)
+        self.n_bins = int(n_bins)
+        self.shuffle_genes = bool(shuffle_genes)
+        self.base_seed = int(seed)
+        self.num_cls_tokens = 1
+        self.expr_embedding_mode = str(expr_embedding_mode).strip().lower()
+
+        if self.gene_vocab_size <= int(getattr(self.vocab, "gene_offset", 0)):
+            raise ValueError("gene_vocab_size must be greater than vocab.gene_offset")
+        if self.n_bins <= 0:
+            raise ValueError("n_bins must be > 0")
+        if self.expr_embedding_mode not in {"bin", "rank"}:
+            raise ValueError("expr_embedding_mode must be one of {'bin', 'rank'}")
+
+        teacher_top_genes = int(teacher_top_genes)
+        if teacher_top_genes <= 0:
+            raise ValueError("teacher_top_genes must be > 0")
+        if int(max_input_genes) > 0:
+            teacher_top_genes = min(teacher_top_genes, int(max_input_genes))
+        self.teacher_top_genes = int(teacher_top_genes)
+        teacher_num_views = int(teacher_num_views)
+        if teacher_num_views <= 0:
+            raise ValueError("teacher_num_views must be > 0")
+        self.teacher_num_views = int(teacher_num_views)
+        if not (0.0 < float(teacher_global_ratio) <= 1.0):
+            raise ValueError("teacher_global_ratio must be in (0, 1]")
+        self.teacher_global_ratio = float(teacher_global_ratio)
+
+        if not (0.0 < float(student_global_ratio_min) <= 1.0):
+            raise ValueError("student_global_ratio_min must be in (0, 1]")
+        if not (0.0 < float(student_global_ratio_max) <= 1.0):
+            raise ValueError("student_global_ratio_max must be in (0, 1]")
+        if float(student_global_ratio_min) > float(student_global_ratio_max):
+            raise ValueError("student_global_ratio_min must be <= student_global_ratio_max")
+        self.student_global_ratio_min = float(student_global_ratio_min)
+        self.student_global_ratio_max = float(student_global_ratio_max)
+        self.student_global_min_genes = int(student_global_min_genes)
+        self.student_local_min_genes = int(student_local_min_genes)
+        self.student_local_max_genes = int(student_local_max_genes)
+        self.student_num_local_views = int(student_num_local_views)
+        self.student_global_dropout_prob = float(student_global_dropout_prob)
+        self.student_local_dropout_prob = float(student_local_dropout_prob)
+        self.student_global_bin_mask_prob = float(student_global_bin_mask_prob)
+        self.student_global_bin_mask_before_binning = bool(student_global_bin_mask_before_binning)
+        high_expr_global_view = str(student_high_expr_global_view or "none").strip().lower()
+        if high_expr_global_view not in {"none", "global1", "global2"}:
+            raise ValueError("student_high_expr_global_view must be one of {'none', 'global1', 'global2'}")
+        self.student_high_expr_global_view = high_expr_global_view
+        high_expr_local_view = str(student_high_expr_local_view or "none").strip().lower()
+        if high_expr_local_view != "none":
+            if not (high_expr_local_view.startswith("local") and high_expr_local_view[5:].isdigit()):
+                raise ValueError("student_high_expr_local_view must be 'none' or localN, e.g. local1")
+        self.student_high_expr_local_view = high_expr_local_view
+        gene_only_local_view = str(student_gene_only_local_view or "none").strip().lower()
+        if gene_only_local_view != "none":
+            if not (gene_only_local_view.startswith("local") and gene_only_local_view[5:].isdigit()):
+                raise ValueError("student_gene_only_local_view must be 'none' or localN, e.g. local1")
+        self.student_gene_only_local_view = gene_only_local_view
+        global_mask_mode = str(student_global_mask_mode or "current").strip().lower().replace("-", "_")
+        if global_mask_mode not in {"current", "global1", "global2", "both"}:
+            raise ValueError("student_global_mask_mode must be one of {'current', 'global1', 'global2', 'both'}")
+        self.student_global_mask_mode = "global1" if global_mask_mode == "current" else global_mask_mode
+        self.student_global_model_masking = bool(student_global_model_masking)
+        self.student_global_mask_sample_probability = float(student_global_mask_sample_probability)
+        self.student_global_mask_ratio_min = float(student_global_mask_ratio_min)
+        self.student_global_mask_ratio_max = float(student_global_mask_ratio_max)
+        self.include_teacher_cls_decoder_targets = bool(include_teacher_cls_decoder_targets)
+        target_mode = str(teacher_cls_decoder_target_mode or "top_other_zero").strip().lower().replace("-", "_")
+        if target_mode not in {"teacher_view", "top_other_zero"}:
+            raise ValueError("teacher_cls_decoder_target_mode must be one of {'teacher_view', 'top_other_zero'}")
+        self.teacher_cls_decoder_target_mode = target_mode
+        self.teacher_cls_decoder_other_expr_genes = int(teacher_cls_decoder_other_expr_genes)
+        self.teacher_cls_decoder_zero_genes = int(teacher_cls_decoder_zero_genes)
+        if self.teacher_cls_decoder_other_expr_genes < 0:
+            raise ValueError("teacher_cls_decoder_other_expr_genes must be >= 0")
+        if self.teacher_cls_decoder_zero_genes < 0:
+            raise ValueError("teacher_cls_decoder_zero_genes must be >= 0")
+        if self.student_global_min_genes <= 0:
+            raise ValueError("student_global_min_genes must be > 0")
+        if self.student_local_min_genes <= 0:
+            raise ValueError("student_local_min_genes must be > 0")
+        if self.student_local_max_genes < self.student_local_min_genes:
+            raise ValueError("student_local_max_genes must be >= student_local_min_genes")
+        if self.student_num_local_views < 0:
+            raise ValueError("student_num_local_views must be >= 0")
+        if self.student_high_expr_local_view != "none":
+            high_expr_local_idx = int(self.student_high_expr_local_view[5:])
+            if high_expr_local_idx < 1 or high_expr_local_idx > self.student_num_local_views:
+                raise ValueError("student_high_expr_local_view must refer to an existing local view")
+        if self.student_gene_only_local_view != "none":
+            gene_only_local_idx = int(self.student_gene_only_local_view[5:])
+            if gene_only_local_idx < 1 or gene_only_local_idx > self.student_num_local_views:
+                raise ValueError("student_gene_only_local_view must refer to an existing local view")
+        if not (0.0 <= self.student_global_dropout_prob < 1.0):
+            raise ValueError("student_global_dropout_prob must be in [0, 1)")
+        if not (0.0 <= self.student_local_dropout_prob < 1.0):
+            raise ValueError("student_local_dropout_prob must be in [0, 1)")
+        if not (0.0 <= self.student_global_bin_mask_prob < 1.0):
+            raise ValueError("student_global_bin_mask_prob must be in [0, 1)")
+        if not (0.0 <= self.student_global_mask_sample_probability <= 1.0):
+            raise ValueError("student_global_mask_sample_probability must be in [0, 1]")
+        if not (0.0 <= self.student_global_mask_ratio_min < 1.0):
+            raise ValueError("student_global_mask_ratio_min must be in [0, 1)")
+        if not (0.0 <= self.student_global_mask_ratio_max < 1.0):
+            raise ValueError("student_global_mask_ratio_max must be in [0, 1)")
+        if self.student_global_mask_ratio_min > self.student_global_mask_ratio_max:
+            raise ValueError("student_global_mask_ratio_min must be <= student_global_mask_ratio_max")
+
+        mode = str(plan_mode or "lookahead_split_1.2").strip().lower().replace("-", "_")
+        if mode not in {"lookahead_split_1.2", "lookahead_split_1p2", "lookahead", "greedy_lookahead"}:
+            raise ValueError("plan_mode must resolve to lookahead_split_1.2")
+        self.plan_mode = "lookahead_split_1.2"
+        self._rng: Optional[np.random.Generator] = None
+        self._global_mask_plan_calls = 0
+
+        self.teacher_view_use_bin = torch.ones((self.teacher_num_views,), dtype=torch.bool)
+        self.teacher_view_is_global = torch.ones((self.teacher_num_views,), dtype=torch.bool)
+        self.student_view_use_bin = torch.ones((2 + self.student_num_local_views,), dtype=torch.bool)
+        if self.student_gene_only_local_view != "none":
+            self.student_view_use_bin[2 + int(self.student_gene_only_local_view[5:]) - 1] = False
+        self.student_view_is_global = torch.tensor(
+            [True, True] + [False] * self.student_num_local_views,
+            dtype=torch.bool,
+        )
+
+    def _get_rng(self) -> np.random.Generator:
+        if self._rng is None:
+            info = get_worker_info()
+            seed = self.base_seed if info is None else (self.base_seed ^ int(info.seed))
+            self._rng = np.random.default_rng(seed % (2**63 - 1))
+        return self._rng
+
+    def _top_teacher_genes(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        nnz = int(gene_ids.shape[0])
+        if nnz == 0:
+            return gene_ids, expr
+        if nnz <= self.teacher_top_genes:
+            return gene_ids, expr
+        top_idx = np.argpartition(expr, -self.teacher_top_genes)[-self.teacher_top_genes :]
+        return gene_ids[top_idx], expr[top_idx]
+
+    def _maybe_shuffle(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        if not self.shuffle_genes or int(gene_ids.shape[0]) <= 1:
+            return gene_ids, expr
+        perm = rng.permutation(int(gene_ids.shape[0]))
+        return gene_ids[perm], expr[perm]
+
+    def _sample_subset(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+        sample_size: int,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        n = int(gene_ids.shape[0])
+        if n == 0:
+            return gene_ids, expr
+        k = min(max(int(sample_size), 0), n)
+        if k == n:
+            return self._maybe_shuffle(gene_ids, expr, rng)
+        idx = rng.choice(n, size=k, replace=False)
+        return self._maybe_shuffle(gene_ids[idx], expr[idx], rng)
+
+    def _select_high_expr_subset(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+        sample_size: int,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        n = int(gene_ids.shape[0])
+        if n == 0:
+            return gene_ids, expr
+        k = min(max(int(sample_size), 0), n)
+        if k == n:
+            return self._maybe_shuffle(gene_ids, expr, rng)
+        top_idx = np.argpartition(expr, -k)[-k:]
+        return self._maybe_shuffle(gene_ids[top_idx], expr[top_idx], rng)
+
+    def _sample_gene_ids_only(
+        self,
+        gene_ids: np.ndarray,
+        sample_size: int,
+        rng: np.random.Generator,
+        *,
+        replace_if_needed: bool = False,
+    ) -> np.ndarray:
+        gene_ids = np.asarray(gene_ids, dtype=np.int64).reshape(-1)
+        n = int(gene_ids.shape[0])
+        k = int(sample_size)
+        if n == 0 or k <= 0:
+            return np.empty((0,), dtype=np.int64)
+        replace = bool(replace_if_needed and n < k)
+        size = k if replace else min(k, n)
+        idx = rng.choice(n, size=size, replace=replace)
+        return np.asarray(gene_ids[idx], dtype=np.int64).reshape(-1)
+
+    def _build_top_other_zero_decoder_targets(
+        self,
+        *,
+        all_gene_ids: np.ndarray,
+        all_expr: np.ndarray,
+        top_gene_ids: np.ndarray,
+        top_expr: np.ndarray,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        top_gene_ids = np.asarray(top_gene_ids, dtype=np.int64).reshape(-1)
+        top_expr = np.asarray(top_expr, dtype=np.float32).reshape(-1)
+        all_gene_ids = np.asarray(all_gene_ids, dtype=np.int64).reshape(-1)
+        all_expr = np.asarray(all_expr, dtype=np.float32).reshape(-1)
+
+        if int(top_gene_ids.shape[0]) > 0:
+            top_set = set(int(x) for x in top_gene_ids.tolist())
+            other_mask = np.asarray([int(g) not in top_set for g in all_gene_ids], dtype=np.bool_)
+        else:
+            other_mask = np.ones((int(all_gene_ids.shape[0]),), dtype=np.bool_)
+        other_gene_pool = all_gene_ids[other_mask]
+        other_expr_pool = all_expr[other_mask]
+        if int(other_gene_pool.shape[0]) > 0 and self.teacher_cls_decoder_other_expr_genes > 0:
+            other_gene_ids, other_expr = self._sample_subset(
+                other_gene_pool,
+                other_expr_pool,
+                self.teacher_cls_decoder_other_expr_genes,
+                rng=rng,
+            )
+        else:
+            other_gene_ids = np.empty((0,), dtype=np.int64)
+            other_expr = np.empty((0,), dtype=np.float32)
+
+        zero_gene_ids = np.empty((0,), dtype=np.int64)
+        zero_expr = np.empty((0,), dtype=np.float32)
+        if self.teacher_cls_decoder_zero_genes > 0:
+            valid_gene_ids = np.arange(
+                int(getattr(self.vocab, "gene_offset", 0)),
+                int(self.gene_vocab_size),
+                dtype=np.int64,
+            )
+            if int(valid_gene_ids.shape[0]) > 0:
+                observed = np.unique(all_gene_ids.astype(np.int64, copy=False))
+                zero_pool = np.setdiff1d(valid_gene_ids, observed, assume_unique=False)
+                zero_gene_ids = self._sample_gene_ids_only(
+                    zero_pool,
+                    self.teacher_cls_decoder_zero_genes,
+                    rng,
+                    replace_if_needed=False,
+                )
+                zero_expr = np.zeros((int(zero_gene_ids.shape[0]),), dtype=np.float32)
+
+        query_gene_ids = np.concatenate([top_gene_ids, other_gene_ids, zero_gene_ids]).astype(np.int64, copy=False)
+        query_expr = np.concatenate([top_expr, other_expr.astype(np.float32, copy=False), zero_expr]).astype(
+            np.float32, copy=False
+        )
+        if self.shuffle_genes and int(query_gene_ids.shape[0]) > 1:
+            perm = rng.permutation(int(query_gene_ids.shape[0]))
+            query_gene_ids = query_gene_ids[perm]
+            query_expr = query_expr[perm]
+        return query_gene_ids, query_expr
+
+    def _choose_student_subset(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+        sample_size: int,
+        *,
+        high_expr_priority: bool,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        if high_expr_priority:
+            return self._select_high_expr_subset(gene_ids, expr, sample_size, rng)
+        return self._sample_subset(gene_ids, expr, sample_size, rng)
+
+    def _apply_gene_dropout(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+        *,
+        dropout_prob: float,
+        min_keep: int,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        n = int(gene_ids.shape[0])
+        if n == 0 or dropout_prob <= 0.0:
+            return gene_ids, expr
+
+        min_keep = max(1, min(int(min_keep), n))
+        if n <= min_keep:
+            return gene_ids, expr
+
+        keep_mask = rng.random(n) >= float(dropout_prob)
+        keep_count = int(keep_mask.sum())
+        if keep_count < min_keep:
+            restore_count = min_keep - keep_count
+            dropped_idx = np.flatnonzero(~keep_mask)
+            if restore_count > 0 and dropped_idx.size > 0:
+                restore_idx = rng.choice(dropped_idx, size=restore_count, replace=False)
+                keep_mask[restore_idx] = True
+        if bool(keep_mask.all()):
+            return gene_ids, expr
+        return gene_ids[keep_mask], expr[keep_mask]
+
+    def _global_sample_size(self, available: int, rng: np.random.Generator) -> int:
+        if available <= 0:
+            return 0
+        if available <= self.student_global_min_genes:
+            return int(available)
+        ratio = float(rng.uniform(self.student_global_ratio_min, self.student_global_ratio_max))
+        target = int(round(ratio * float(available)))
+        target = max(int(self.student_global_min_genes), target)
+        return min(int(available), target)
+
+    def _local_sample_size(self, available: int, rng: np.random.Generator) -> int:
+        if available <= 0:
+            return 0
+        if available <= self.student_local_min_genes:
+            return int(available)
+        low = min(int(available), int(self.student_local_min_genes))
+        high = min(int(available), int(self.student_local_max_genes))
+        if high <= low:
+            return int(high)
+        return int(rng.integers(low, high + 1))
+
+    def _bucketize_observed_expr(self, expr: np.ndarray) -> np.ndarray:
+        expr_f = np.asarray(expr, dtype=np.float32)
+        n = int(expr_f.shape[0])
+        if n == 0:
+            return np.empty((0,), dtype=np.int64)
+
+        _, inverse, counts = np.unique(expr_f, return_inverse=True, return_counts=True)
+        value_order = np.argsort(inverse, kind="stable")
+        value_offsets = np.zeros(int(counts.shape[0]) + 1, dtype=np.int64)
+        value_offsets[1:] = np.cumsum(counts, dtype=np.int64)
+        plan = _plan_value_aware_lookahead_bins(counts, int(self.n_bins), split_factor=1.2)
+
+        bin_ids = np.empty(n, dtype=np.int64)
+        next_bin = 0
+        for action, start, end, bins_used in plan:
+            if action == "merge":
+                positions = value_order[int(value_offsets[start]) : int(value_offsets[end + 1])]
+                bin_ids[positions] = next_bin + int(self.vocab.bin_offset)
+                next_bin += 1
+                continue
+            if action == "split":
+                positions = value_order[int(value_offsets[start]) : int(value_offsets[start + 1])]
+                for local_bin, group_idx in enumerate(np.array_split(positions, int(bins_used))):
+                    if group_idx.size == 0:
+                        continue
+                    bin_ids[group_idx] = next_bin + local_bin + int(self.vocab.bin_offset)
+                next_bin += int(bins_used)
+                continue
+            raise RuntimeError(f"Unknown bin action: {action}")
+        return bin_ids
+
+    def _rank_observed_expr(self, expr: np.ndarray) -> np.ndarray:
+        expr_f = np.asarray(expr, dtype=np.float32)
+        n = int(expr_f.shape[0])
+        if n == 0:
+            return np.empty((0,), dtype=np.int64)
+        order = np.argsort(-expr_f, kind="stable")
+        rank_ids = np.empty((n,), dtype=np.int64)
+        rank_ids[order] = np.arange(int(getattr(self.vocab, "bin_offset", 3)), int(getattr(self.vocab, "bin_offset", 3)) + n, dtype=np.int64)
+        return rank_ids
+
+    def _encode_expression_ids(self, expr: np.ndarray) -> np.ndarray:
+        if self.expr_embedding_mode == "rank":
+            return self._rank_observed_expr(expr)
+        return self._bucketize_observed_expr(expr)
+
+    def _build_view(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+        *,
+        use_bin: bool,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        gene_ids, expr = self._maybe_shuffle(gene_ids, expr, rng)
+        if not use_bin:
+            pad_id = int(self.vocab.bin_pad_id)
+            return gene_ids, np.full((int(gene_ids.shape[0]),), pad_id, dtype=np.int64)
+        if int(expr.shape[0]) == 0:
+            return gene_ids, np.empty((0,), dtype=np.int64)
+        return gene_ids, self._encode_expression_ids(expr)
+
+    def _mask_bin_ids(
+        self,
+        bin_ids: np.ndarray,
+        *,
+        mask_prob: float,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        n = int(bin_ids.shape[0])
+        mask = self._sample_bin_mask(n=n, mask_prob=mask_prob, rng=rng)
+        if n == 0 or not bool(mask.any()):
+            return bin_ids, mask
+
+        masked_bin_ids = np.array(bin_ids, copy=True)
+        masked_bin_ids[mask] = int(self.vocab.bin_mask_id)
+        return masked_bin_ids, mask
+
+    def _sample_bin_mask(
+        self,
+        *,
+        n: int,
+        mask_prob: float,
+        rng: np.random.Generator,
+    ) -> np.ndarray:
+        n = int(n)
+        if n <= 0 or mask_prob <= 0.0:
+            return np.zeros((max(n, 0),), dtype=np.bool_)
+
+        mask = rng.random(n) < float(mask_prob)
+        if not bool(mask.any()):
+            mask[int(rng.integers(0, n))] = True
+        return mask
+
+    def _build_masked_bin_view(
+        self,
+        gene_ids: np.ndarray,
+        expr: np.ndarray,
+        *,
+        mask_prob: float,
+        rng: np.random.Generator,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        gene_ids, expr = self._maybe_shuffle(gene_ids, expr, rng)
+        n = int(gene_ids.shape[0])
+        if n == 0:
+            empty = np.empty((0,), dtype=np.int64)
+            return gene_ids, empty, np.zeros((0,), dtype=np.bool_)
+
+        mask = self._sample_bin_mask(n=n, mask_prob=mask_prob, rng=rng)
+        if not self.student_global_bin_mask_before_binning:
+            bin_ids = self._encode_expression_ids(expr)
+            masked_bin_ids = np.array(bin_ids, copy=True)
+            masked_bin_ids[mask] = int(self.vocab.bin_mask_id)
+            return gene_ids, masked_bin_ids, mask
+
+        masked_bin_ids = np.full((n,), int(self.vocab.bin_mask_id), dtype=np.int64)
+        keep_mask = ~mask
+        if bool(keep_mask.any()):
+            masked_bin_ids[keep_mask] = self._encode_expression_ids(expr[keep_mask])
+        return gene_ids, masked_bin_ids, mask
+
+    def _build_model_global_mask_plan(
+        self,
+        *,
+        batch_size: int,
+        rng: np.random.Generator,
+    ) -> List[Tuple[bool, float]]:
+        del rng
+        local_global_views = int(batch_size) * 2
+        rank, world = _distributed_rank_world()
+        total_global_views = int(local_global_views) * int(world)
+        n_masked = int(total_global_views * self.student_global_mask_sample_probability)
+        n_masked = max(0, min(n_masked, total_global_views))
+        if local_global_views <= 0 or n_masked == 0:
+            return [(False, 0.0) for _ in range(local_global_views)]
+
+        plan_seed = int(self.base_seed) + int(self._global_mask_plan_calls)
+        self._global_mask_plan_calls += 1
+        plan_rng = np.random.default_rng(plan_seed % (2**63 - 1))
+        masked_global_indices = set(int(i) for i in plan_rng.permutation(total_global_views)[:n_masked].tolist())
+        if self.student_global_mask_ratio_min == self.student_global_mask_ratio_max:
+            global_ratios = np.full((total_global_views,), self.student_global_mask_ratio_min, dtype=np.float64)
+        else:
+            global_ratios = plan_rng.uniform(
+                self.student_global_mask_ratio_min,
+                self.student_global_mask_ratio_max,
+                size=total_global_views,
+            )
+
+        start = int(rank) * int(local_global_views)
+        local_plan: List[Tuple[bool, float]] = []
+        for local_idx in range(local_global_views):
+            global_idx = start + int(local_idx)
+            is_masked = global_idx in masked_global_indices
+            local_plan.append((bool(is_masked), float(global_ratios[global_idx]) if is_masked else 0.0))
+        return local_plan
+
+    def _pack_views(
+        self,
+        view_gene_ids: List[List[np.ndarray]],
+        view_bin_ids: List[List[np.ndarray]],
+        view_use_bin: torch.Tensor,
+        view_token_masks: Optional[List[List[np.ndarray]]] = None,
+    ) -> Dict[str, torch.Tensor]:
+        def _as_1d_int_array(arr: Any) -> np.ndarray:
+            np_arr = np.asarray(arr, dtype=np.int64)
+            if np_arr.ndim == 0:
+                np_arr = np_arr.reshape(-1)
+            elif np_arr.ndim != 1:
+                np_arr = np_arr.reshape(-1)
+            return np_arr
+
+        batch_size = len(view_gene_ids)
+        num_views = len(view_gene_ids[0]) if batch_size > 0 else int(view_use_bin.numel())
+        max_tokens = self.num_cls_tokens
+        for sample_views in view_gene_ids:
+            for gene_ids in sample_views:
+                gene_ids_1d = _as_1d_int_array(gene_ids)
+                max_tokens = max(max_tokens, self.num_cls_tokens + int(gene_ids_1d.shape[0]))
+
+        device = torch.device("cpu")
+        gene_pad_id = int(self.vocab.gene_pad_id)
+        bin_pad_id = int(self.vocab.bin_pad_id)
+        gene_cls_id = int(self.vocab.gene_cls_id)
+
+        gene_tensor = torch.full((batch_size, num_views, max_tokens), gene_pad_id, dtype=torch.long, device=device)
+        bin_tensor = torch.full((batch_size, num_views, max_tokens), bin_pad_id, dtype=torch.long, device=device)
+        attn_tensor = torch.zeros((batch_size, num_views, max_tokens), dtype=torch.bool, device=device)
+        token_mask_tensor = None
+        if view_token_masks is not None:
+            token_mask_tensor = torch.zeros((batch_size, num_views, max_tokens), dtype=torch.bool, device=device)
+
+        for m in range(batch_size):
+            for v in range(num_views):
+                gene_tensor[m, v, 0] = gene_cls_id
+                bin_tensor[m, v, 0] = bin_pad_id
+                attn_tensor[m, v, 0] = True
+
+                gene_ids = view_gene_ids[m][v]
+                bin_ids = view_bin_ids[m][v]
+                gene_ids = _as_1d_int_array(gene_ids)
+                bin_ids = _as_1d_int_array(bin_ids)
+                if int(gene_ids.shape[0]) == 0:
+                    continue
+                start = self.num_cls_tokens
+                end = start + int(gene_ids.shape[0])
+                gene_tensor[m, v, start:end] = torch.from_numpy(gene_ids).to(dtype=torch.long)
+                bin_tensor[m, v, start:end] = torch.from_numpy(bin_ids).to(dtype=torch.long)
+                attn_tensor[m, v, start:end] = True
+                if token_mask_tensor is not None:
+                    token_mask_tensor[m, v, start:end] = torch.from_numpy(view_token_masks[m][v]).to(dtype=torch.bool)
+
+        packed = {
+            "gene_ids": gene_tensor,
+            "bin_ids": bin_tensor,
+            "attention_mask": attn_tensor,
+        }
+        if token_mask_tensor is not None:
+            packed["token_mask"] = token_mask_tensor
+        return packed
+
+    def _pack_decoder_targets(
+        self,
+        target_gene_ids: List[np.ndarray],
+        target_expr: List[np.ndarray],
+    ) -> Dict[str, torch.Tensor]:
+        batch_size = len(target_gene_ids)
+        max_targets = 0
+        for gene_ids in target_gene_ids:
+            max_targets = max(max_targets, int(np.asarray(gene_ids, dtype=np.int64).reshape(-1).shape[0]))
+
+        device = torch.device("cpu")
+        gene_pad_id = int(self.vocab.gene_pad_id)
+        gene_tensor = torch.full((batch_size, max_targets), gene_pad_id, dtype=torch.long, device=device)
+        expr_tensor = torch.zeros((batch_size, max_targets), dtype=torch.float32, device=device)
+        mask_tensor = torch.zeros((batch_size, max_targets), dtype=torch.bool, device=device)
+
+        for idx, (gene_ids, expr) in enumerate(zip(target_gene_ids, target_expr)):
+            gene_arr = np.asarray(gene_ids, dtype=np.int64).reshape(-1)
+            expr_arr = np.asarray(expr, dtype=np.float32).reshape(-1)
+            if int(gene_arr.shape[0]) == 0:
+                continue
+            end = int(gene_arr.shape[0])
+            gene_tensor[idx, :end] = torch.from_numpy(gene_arr).to(dtype=torch.long)
+            expr_tensor[idx, :end] = torch.from_numpy(expr_arr).to(dtype=torch.float32)
+            mask_tensor[idx, :end] = True
+
+        return {
+            "gene_ids": gene_tensor,
+            "expr": expr_tensor,
+            "mask": mask_tensor,
+        }
+
+    def _pack_decoder_obs_features(self, obs_features: List[np.ndarray]) -> torch.Tensor:
+        batch_size = len(obs_features)
+        device = torch.device("cpu")
+        obs_tensor = torch.zeros((batch_size, 2), dtype=torch.float32, device=device)
+        for idx, features in enumerate(obs_features):
+            feat_arr = np.asarray(features, dtype=np.float32).reshape(-1)
+            if int(feat_arr.shape[0]) != 2:
+                raise ValueError(f"Expected 2 decoder observation features, got shape={feat_arr.shape!r}")
+            obs_tensor[idx] = torch.from_numpy(feat_arr).to(dtype=torch.float32)
+        return obs_tensor
+
+    def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        batch_size = len(examples)
+        if batch_size == 0:
+            raise ValueError("Empty batch passed to collator.")
+
+        rng = self._get_rng()
+        teacher_gene_views: List[List[np.ndarray]] = []
+        teacher_bin_views: List[List[np.ndarray]] = []
+        student_gene_views: List[List[np.ndarray]] = []
+        student_bin_views: List[List[np.ndarray]] = []
+        student_token_masks: List[List[np.ndarray]] = []
+        decoder_target_gene_ids: List[np.ndarray] = []
+        decoder_target_expr: List[np.ndarray] = []
+        decoder_target_obs_features: List[np.ndarray] = []
+        model_mask_plan = (
+            self._build_model_global_mask_plan(batch_size=batch_size, rng=rng)
+            if self.student_global_model_masking
+            else None
+        )
+        model_mask_plan_idx = 0
+
+        for ex in examples:
+            gene_ids = np.asarray(ex["gene_ids"], dtype=np.int64)
+            expr = np.asarray(ex["expr"], dtype=np.float32)
+
+            teacher_gene_ids, teacher_expr = self._top_teacher_genes(gene_ids, expr)
+            teacher_gene_ids, teacher_expr = self._maybe_shuffle(teacher_gene_ids, teacher_expr, rng)
+
+            available = int(teacher_gene_ids.shape[0])
+            teacher_sample_size = int(round(self.teacher_global_ratio * float(available)))
+            teacher_sample_size = max(1, teacher_sample_size) if available > 0 else 0
+
+            teacher_views: List[Tuple[np.ndarray, np.ndarray]] = []
+            first_teacher_target_gene_ids: Optional[np.ndarray] = None
+            first_teacher_target_expr: Optional[np.ndarray] = None
+            for view_idx in range(self.teacher_num_views):
+                sampled_teacher_gene_ids, sampled_teacher_expr = self._sample_subset(
+                    teacher_gene_ids,
+                    teacher_expr,
+                    teacher_sample_size,
+                    rng=rng,
+                )
+                if view_idx == 0 and self.include_teacher_cls_decoder_targets:
+                    first_teacher_target_gene_ids, first_teacher_target_expr = self._build_top_other_zero_decoder_targets(
+                        all_gene_ids=gene_ids,
+                        all_expr=expr,
+                        top_gene_ids=teacher_gene_ids,
+                        top_expr=teacher_expr,
+                        rng=rng,
+                    )
+                t_bin_gene, t_bin_bin = self._build_view(
+                    sampled_teacher_gene_ids,
+                    sampled_teacher_expr,
+                    use_bin=True,
+                    rng=rng,
+                )
+                teacher_views.append((t_bin_gene, t_bin_bin))
+            global_bin_size = self._global_sample_size(available, rng)
+            global_gene_size = self._global_sample_size(available, rng)
+            local_sizes = [self._local_sample_size(available, rng) for _ in range(self.student_num_local_views)]
+
+            sgb_gene_ids, sgb_expr = self._choose_student_subset(
+                teacher_gene_ids,
+                teacher_expr,
+                global_bin_size,
+                high_expr_priority=self.student_high_expr_global_view == "global1",
+                rng=rng,
+            )
+            sgg_gene_ids, sgg_expr = self._choose_student_subset(
+                teacher_gene_ids,
+                teacher_expr,
+                global_gene_size,
+                high_expr_priority=self.student_high_expr_global_view == "global2",
+                rng=rng,
+            )
+            local_subsets = [
+                self._choose_student_subset(
+                    teacher_gene_ids,
+                    teacher_expr,
+                    local_size,
+                    high_expr_priority=self.student_high_expr_local_view == f"local{local_idx + 1}",
+                    rng=rng,
+                )
+                for local_idx, local_size in enumerate(local_sizes)
+            ]
+
+            sgb_gene_ids, sgb_expr = self._apply_gene_dropout(
+                sgb_gene_ids,
+                sgb_expr,
+                dropout_prob=self.student_global_dropout_prob,
+                min_keep=self.student_global_min_genes,
+                rng=rng,
+            )
+            sgg_gene_ids, sgg_expr = self._apply_gene_dropout(
+                sgg_gene_ids,
+                sgg_expr,
+                dropout_prob=self.student_global_dropout_prob,
+                min_keep=self.student_global_min_genes,
+                rng=rng,
+            )
+            local_subsets = [
+                self._apply_gene_dropout(
+                    local_gene_ids,
+                    local_expr,
+                    dropout_prob=self.student_local_dropout_prob,
+                    min_keep=self.student_local_min_genes,
+                    rng=rng,
+                )
+                for local_gene_ids, local_expr in local_subsets
+            ]
+
+            if model_mask_plan is not None:
+                mask_global1, global1_mask_prob = model_mask_plan[model_mask_plan_idx]
+                mask_global2, global2_mask_prob = model_mask_plan[model_mask_plan_idx + 1]
+                model_mask_plan_idx += 2
+            else:
+                mask_global1 = self.student_global_mask_mode in {"global1", "both"}
+                mask_global2 = self.student_global_mask_mode in {"global2", "both"}
+                global1_mask_prob = self.student_global_bin_mask_prob
+                global2_mask_prob = self.student_global_bin_mask_prob
+            if mask_global1:
+                sgb_gene, sgb_bin, sgb_mask = self._build_masked_bin_view(
+                    sgb_gene_ids,
+                    sgb_expr,
+                    mask_prob=global1_mask_prob,
+                    rng=rng,
+                )
+            else:
+                sgb_gene, sgb_bin = self._build_view(sgb_gene_ids, sgb_expr, use_bin=True, rng=rng)
+                sgb_mask_value = False if model_mask_plan is not None else True
+                sgb_mask = np.full((int(sgb_gene.shape[0]),), sgb_mask_value, dtype=np.bool_)
+            if mask_global2:
+                sgg_gene, sgg_bin, sgg_mask = self._build_masked_bin_view(
+                    sgg_gene_ids,
+                    sgg_expr,
+                    mask_prob=global2_mask_prob,
+                    rng=rng,
+                )
+            else:
+                sgg_gene, sgg_bin = self._build_view(sgg_gene_ids, sgg_expr, use_bin=True, rng=rng)
+                sgg_mask_value = False if model_mask_plan is not None else True
+                sgg_mask = np.full((int(sgg_gene.shape[0]),), sgg_mask_value, dtype=np.bool_)
+            local_views = [
+                self._build_view(
+                    local_gene_ids,
+                    local_expr,
+                    use_bin=self.student_gene_only_local_view != f"local{local_idx + 1}",
+                    rng=rng,
+                )
+                for local_idx, (local_gene_ids, local_expr) in enumerate(local_subsets)
+            ]
+            local_genes = [view_gene for view_gene, _ in local_views]
+            local_bins = [view_bin for _, view_bin in local_views]
+            local_masks = [np.zeros((int(view_gene.shape[0]),), dtype=np.bool_) for view_gene in local_genes]
+
+            teacher_gene_views.append([view_gene for view_gene, _ in teacher_views])
+            teacher_bin_views.append([view_bin for _, view_bin in teacher_views])
+            student_gene_views.append([sgb_gene, sgg_gene] + local_genes)
+            student_bin_views.append([sgb_bin, sgg_bin] + local_bins)
+            student_token_masks.append([sgb_mask, sgg_mask] + local_masks)
+            if self.include_teacher_cls_decoder_targets:
+                decoder_target_gene_ids.append(
+                    np.empty((0,), dtype=np.int64)
+                    if first_teacher_target_gene_ids is None
+                    else first_teacher_target_gene_ids
+                )
+                decoder_target_expr.append(
+                    np.empty((0,), dtype=np.float32)
+                    if first_teacher_target_expr is None
+                    else first_teacher_target_expr
+                )
+                total_counts = float(ex.get("total_counts", gene_ids.shape[0]))
+                n_genes_by_counts = float(ex.get("n_genes_by_counts", gene_ids.shape[0]))
+                decoder_target_obs_features.append(
+                    np.asarray(
+                        [
+                            np.log1p(max(total_counts, 0.0)),
+                            np.log1p(max(n_genes_by_counts, 0.0)),
+                        ],
+                        dtype=np.float32,
+                    )
+                )
+
+        teacher_packed = self._pack_views(teacher_gene_views, teacher_bin_views, self.teacher_view_use_bin)
+        student_packed = self._pack_views(
+            student_gene_views,
+            student_bin_views,
+            self.student_view_use_bin,
+            view_token_masks=student_token_masks,
+        )
+
+        batch = {
+            "teacher_gene_ids": teacher_packed["gene_ids"],
+            "teacher_bin_ids": teacher_packed["bin_ids"],
+            "teacher_attention_mask": teacher_packed["attention_mask"],
+            "teacher_view_use_bin": self.teacher_view_use_bin.clone(),
+            "teacher_view_is_global": self.teacher_view_is_global.clone(),
+            "student_gene_ids": student_packed["gene_ids"],
+            "student_bin_ids": student_packed["bin_ids"],
+            "student_attention_mask": student_packed["attention_mask"],
+            "student_ibot_mask": student_packed["token_mask"],
+            "student_view_use_bin": self.student_view_use_bin.clone(),
+            "student_view_is_global": self.student_view_is_global.clone(),
+        }
+        if self.include_teacher_cls_decoder_targets:
+            decoder_targets = self._pack_decoder_targets(decoder_target_gene_ids, decoder_target_expr)
+            batch["teacher_cls_decoder_target_gene_ids"] = decoder_targets["gene_ids"]
+            batch["teacher_cls_decoder_target_expr"] = decoder_targets["expr"]
+            batch["teacher_cls_decoder_target_mask"] = decoder_targets["mask"]
+            batch["teacher_cls_decoder_obs_features"] = self._pack_decoder_obs_features(
+                decoder_target_obs_features
+            )
+        return batch
+
+
+__all__ = ["ContrastiveSplitCellCollator"]

scRep_pretrain/model/contrastive_trainer.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from transformers import Trainer, TrainerCallback
+
+
+def _unwrap_model(model):
+    current = model
+    while hasattr(current, "module"):
+        current = current.module
+    return current
+
+
+class _TeacherMomentumCallback(TrainerCallback):
+    def on_optimizer_step(self, args, state, control, model=None, **kwargs):
+        base_model = _unwrap_model(model)
+        if hasattr(base_model, "update_teacher"):
+            current_step = int(state.global_step) + 1
+            max_steps = int(getattr(args, "max_steps", 0) or 0)
+            base_model.update_teacher(global_step=current_step, max_steps=max_steps)
+        return control
+
+
+class ContrastiveCellTrainer(Trainer):
+    """
+    Trainer for DINO-like student/teacher models.
+    """
+
+    def __init__(
+        self,
+        *args,
+        loss_scale: float = 1.0,
+        steps_per_epoch: Optional[int] = None,
+        planned_num_epochs: Optional[float] = None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.loss_scale = float(loss_scale)
+        self.steps_per_epoch = None if steps_per_epoch is None else int(steps_per_epoch)
+        self.planned_num_epochs = None if planned_num_epochs is None else float(planned_num_epochs)
+        self._latest_train_metrics = {}
+        self.add_callback(_TeacherMomentumCallback())
+
+    def _current_data_epoch(self) -> Optional[float]:
+        if self.steps_per_epoch is None or self.steps_per_epoch <= 0:
+            return None
+        return float(self.state.global_step + 1) / float(self.steps_per_epoch)
+
+    def _current_train_epoch(self) -> Optional[float]:
+        max_steps = int(getattr(self.args, "max_steps", 0) or 0)
+        if max_steps <= 0:
+            return None
+        planned_num_epochs = self.planned_num_epochs
+        if planned_num_epochs is None or planned_num_epochs <= 0.0:
+            planned_num_epochs = 1.0
+        return float(self.state.global_step + 1) / float(max_steps) * float(planned_num_epochs)
+
+    def _current_train_progress(self) -> Optional[float]:
+        max_steps = int(getattr(self.args, "max_steps", 0) or 0)
+        if max_steps <= 0:
+            return None
+        return float(self.state.global_step + 1) / float(max_steps)
+
+    def compute_loss(self, model, inputs, return_outputs: bool = False, **kwargs):
+        batch = inputs["batch"]
+        outputs = model(batch)
+        loss_unscaled = outputs["loss"]
+        loss = loss_unscaled * self.loss_scale
+        if return_outputs:
+            return loss, outputs
+        return loss
+
+    def training_step(self, model, inputs, num_items_in_batch=None):
+        model.train()
+        inputs = self._prepare_inputs(inputs)
+
+        with self.compute_loss_context_manager():
+            loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+
+        if self.args.n_gpu > 1:
+            loss = loss.mean()
+
+        self.accelerator.backward(loss)
+
+        data_epoch = self._current_data_epoch()
+        train_epoch = self._current_train_epoch()
+        train_progress = self._current_train_progress()
+
+        log_dict = {}
+        for key in (
+            "loss_total",
+            "loss_dino",
+            "loss_dino_same_mode",
+            "loss_dino_cross_mode",
+            "loss_token",
+            "loss_token_same_mode",
+            "loss_token_cross_mode",
+            "loss_decoder",
+            "teacher_momentum",
+        ):
+            if key in outputs:
+                value = outputs[key]
+                if hasattr(value, "detach"):
+                    value = value.detach().float().item()
+                log_dict[key] = float(value)
+        if data_epoch is not None:
+            log_dict["data_epoch"] = float(data_epoch)
+        if train_epoch is not None:
+            log_dict["train_epoch"] = float(train_epoch)
+        if train_progress is not None:
+            log_dict["train_progress"] = float(train_progress)
+        if log_dict:
+            self._latest_train_metrics = log_dict
+
+        return loss.detach()
+
+    def log(self, logs: dict, start_time=None) -> None:
+        if self._latest_train_metrics:
+            for key, value in self._latest_train_metrics.items():
+                logs.setdefault(key, value)
+        return super().log(logs, start_time=start_time)
+
+
+__all__ = ["ContrastiveCellTrainer"]

scRep_pretrain/model/foundation_model.py

@@ -0,0 +1,1411 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, fields as dataclass_fields
+from contextlib import ExitStack
+from typing import Any, Dict, List, Optional, Tuple
+
+import json
+import math
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+try:
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+except Exception:  # pragma: no cover - FSDP availability depends on runtime
+    FSDP = None
+
+
+def _dist_is_initialized() -> bool:
+    return torch.distributed.is_available() and torch.distributed.is_initialized()
+
+
+def _leaf_fsdp_modules(module: nn.Module) -> List[nn.Module]:
+    if FSDP is None:
+        return []
+
+    fsdp_modules: List[nn.Module] = []
+    for child in module.modules():
+        if not isinstance(child, FSDP):
+            continue
+        has_fsdp_descendant = any(isinstance(grandchild, FSDP) for grandchild in child.children())
+        if not has_fsdp_descendant:
+            fsdp_modules.append(child)
+    return fsdp_modules
+
+
+def _json_dump(obj: Any, path: str) -> None:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(obj, f, ensure_ascii=False, indent=2)
+
+
+class RmsNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = float(eps)
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        inv = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
+        return x * inv * self.scale
+
+
+class SdpaSelfAttention(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, dropout: float):
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+        self.d_model = int(d_model)
+        self.n_heads = int(n_heads)
+        self.head_dim = self.d_model // self.n_heads
+        self.dropout = float(dropout)
+
+        self.q_proj = nn.Linear(self.d_model, self.d_model, bias=False)
+        self.k_proj = nn.Linear(self.d_model, self.d_model, bias=False)
+        self.v_proj = nn.Linear(self.d_model, self.d_model, bias=False)
+        self.o_proj = nn.Linear(self.d_model, self.d_model, bias=False)
+
+    def forward(self, x: torch.Tensor, *, attn_mask_valid: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if attn_mask_valid.shape != (batch_size, seq_len):
+            raise ValueError(
+                "attn_mask_valid must have shape [batch_size, seq_len], "
+                f"got {tuple(attn_mask_valid.shape)} for x={tuple(x.shape)}"
+            )
+        # Expression-masked genes remain valid tokens. This mask only excludes
+        # right-padding positions from the attention key/value softmax.
+        key_mask = attn_mask_valid.to(device=x.device, dtype=torch.bool)[:, None, None, :]
+        drop_p = self.dropout if self.training else 0.0
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=key_mask,
+            dropout_p=drop_p,
+            is_causal=False,
+        )
+        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
+        return self.o_proj(out)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, dropout: float):
+        super().__init__()
+        self.norm1 = RmsNorm(d_model)
+        self.attn = SdpaSelfAttention(d_model, n_heads, dropout)
+        self.drop = nn.Dropout(dropout)
+        self.norm2 = RmsNorm(d_model)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, 4 * d_model, bias=False),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(4 * d_model, d_model, bias=False),
+        )
+
+    def forward(self, x: torch.Tensor, *, attn_mask_valid: torch.Tensor) -> torch.Tensor:
+        x = x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype)
+        h = self.norm1(x)
+        x = x + self.drop(self.attn(h, attn_mask_valid=attn_mask_valid))
+        x = x + self.drop(self.ffn(self.norm2(x)))
+        return x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype)
+
+
+@dataclass
+class FoundationConfig:
+    d_model: int = 768
+    n_heads: int = 12
+    dropout: float = 0.0
+    encoder_layers: int = 12
+    gradient_checkpointing: bool = False
+
+    gene_vocab_size: int = 0
+    bin_vocab_size: int = 0
+    gene_pad_id: int = 0
+    bin_pad_id: int = 0
+    bin_mask_id: int = 1
+    bin_offset: int = 3
+    gene_cls_id: int = 1
+    expr_embedding_mode: str = "bin"
+    expr_rank_mlp_hidden_dim: int = 256
+
+    dino_out_dim: int = 4096
+    dino_hidden_dim: int = 2048
+    dino_bottleneck_dim: int = 256
+    student_temperature: float = 0.10
+    teacher_temperature: float = 0.04
+    teacher_momentum: float = 0.996
+    teacher_momentum_end: float = 1.0
+    center_momentum: float = 0.9
+    token_center_momentum: float = 0.9
+    cross_modal_dino_weight: float = 0.25
+    cross_modal_pairing_mode: str = "disabled"
+    token_loss_weight: float = 1.0
+    token_unmasked_loss_weight: float = 0.0
+    inference_use_teacher: bool = True
+    teacher_cls_decoder_loss_weight: float = 0.0
+    teacher_cls_decoder_hidden_dim: int = 1024
+    teacher_cls_decoder_condition_on_obs: bool = False
+    teacher_cls_decoder_rank_loss_alpha: float = 0.05
+    teacher_cls_decoder_rank_num_pairs: int = 4096
+    teacher_student_decoder_distill_weight: float = 0.0
+    teacher_student_decoder_loss_mode: str = "distribution"
+    teacher_student_decoder_rank_alpha: float = 0.05
+    teacher_student_decoder_raw_mmd_weight: float = 0.1
+    teacher_student_decoder_raw_rank_weight: float = 0.01
+    teacher_student_decoder_raw_mse_weight: float = 1.0
+
+
+class GeneBinTokenizer(nn.Module):
+    def __init__(self, cfg: FoundationConfig):
+        super().__init__()
+        self.gene_emb = nn.Embedding(cfg.gene_vocab_size, cfg.d_model, padding_idx=cfg.gene_pad_id)
+        self.bin_emb = nn.Embedding(cfg.bin_vocab_size, cfg.d_model, padding_idx=cfg.bin_pad_id)
+        self.expr_embedding_mode = str(getattr(cfg, "expr_embedding_mode", "bin")).strip().lower()
+        if self.expr_embedding_mode not in {"bin", "rank"}:
+            raise ValueError(f"Unsupported expr_embedding_mode={self.expr_embedding_mode!r}")
+        self.bin_mask_id = int(getattr(cfg, "bin_mask_id", 1))
+        self.bin_offset = int(getattr(cfg, "bin_offset", 3))
+        rank_hidden_dim = int(getattr(cfg, "expr_rank_mlp_hidden_dim", 256))
+        if rank_hidden_dim <= 0:
+            raise ValueError("expr_rank_mlp_hidden_dim must be > 0")
+        self.rank_mlp: Optional[nn.Sequential]
+        self.rank_mask_emb: Optional[nn.Parameter]
+        if self.expr_embedding_mode == "rank":
+            self.rank_mlp = nn.Sequential(
+                nn.Linear(1, rank_hidden_dim, bias=True),
+                nn.SiLU(),
+                nn.Linear(rank_hidden_dim, cfg.d_model, bias=False),
+            )
+            self.rank_mask_emb = nn.Parameter(torch.empty(cfg.d_model))
+            nn.init.normal_(self.rank_mask_emb, mean=0.0, std=0.02)
+        else:
+            self.rank_mlp = None
+            self.rank_mask_emb = None
+
+    def _rank_expression_emb(self, expr_ids: torch.Tensor) -> torch.Tensor:
+        if self.rank_mlp is None:
+            raise RuntimeError("rank_mlp is not initialized for non-rank expr_embedding_mode")
+        out = self.gene_emb.weight.new_zeros((*expr_ids.shape, self.gene_emb.embedding_dim))
+        positive_mask = expr_ids >= int(self.bin_offset)
+        if bool(positive_mask.any().item()):
+            rank_values = (expr_ids[positive_mask] - int(self.bin_offset) + 1).to(dtype=out.dtype).unsqueeze(-1)
+            rank_emb = self.rank_mlp(rank_values)
+            if rank_emb.dtype != out.dtype:
+                rank_emb = rank_emb.to(dtype=out.dtype)
+            out[positive_mask] = rank_emb
+        mask_mask = expr_ids == int(self.bin_mask_id)
+        if bool(mask_mask.any().item()):
+            if self.rank_mask_emb is None:
+                raise RuntimeError("rank_mask_emb is not initialized for rank expr_embedding_mode")
+            out[mask_mask] = self.rank_mask_emb.to(dtype=out.dtype)
+        return out
+
+    def forward(
+        self,
+        gene_ids: torch.Tensor,
+        bin_ids: torch.Tensor,
+        attn_mask_valid: torch.Tensor,
+        *,
+        use_bin_embeddings: bool,
+    ) -> torch.Tensor:
+        x = self.gene_emb(gene_ids)
+        if bool(use_bin_embeddings):
+            if self.expr_embedding_mode == "rank":
+                x = x + self._rank_expression_emb(bin_ids)
+            else:
+                x = x + self.bin_emb(bin_ids)
+        return x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype)
+
+
+class EncoderBackbone(nn.Module):
+    def __init__(self, cfg: FoundationConfig):
+        super().__init__()
+        self.gradient_checkpointing = bool(getattr(cfg, "gradient_checkpointing", False))
+        self.tokenizer = GeneBinTokenizer(cfg)
+        self.encoder = nn.ModuleList(
+            [EncoderBlock(cfg.d_model, cfg.n_heads, cfg.dropout) for _ in range(cfg.encoder_layers)]
+        )
+        self.encoder_norm = RmsNorm(cfg.d_model)
+
+    def forward(
+        self,
+        gene_ids: torch.Tensor,
+        bin_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        *,
+        use_bin_embeddings: bool,
+    ) -> torch.Tensor:
+        x = self.tokenizer(
+            gene_ids,
+            bin_ids,
+            attention_mask,
+            use_bin_embeddings=use_bin_embeddings,
+        )
+        for blk in self.encoder:
+            if self.gradient_checkpointing and self.training and torch.is_grad_enabled():
+                def _block_forward(hidden: torch.Tensor, *, block: EncoderBlock = blk) -> torch.Tensor:
+                    return block(hidden, attn_mask_valid=attention_mask)
+
+                x = checkpoint(_block_forward, x, use_reentrant=False)
+            else:
+                x = blk(x, attn_mask_valid=attention_mask)
+        x = self.encoder_norm(x)
+        return x * attention_mask.unsqueeze(-1).to(dtype=x.dtype)
+
+
+class DinoHead(nn.Module):
+    def __init__(self, cfg: FoundationConfig):
+        super().__init__()
+        hidden_dim = int(cfg.dino_hidden_dim)
+        bottleneck_dim = int(cfg.dino_bottleneck_dim)
+        out_dim = int(cfg.dino_out_dim)
+        if hidden_dim <= 0 or bottleneck_dim <= 0 or out_dim <= 0:
+            raise ValueError("DINO head dimensions must all be > 0")
+        self.mlp = nn.Sequential(
+            nn.Linear(cfg.d_model, hidden_dim, bias=False),
+            nn.GELU(),
+            nn.Linear(hidden_dim, hidden_dim, bias=False),
+            nn.GELU(),
+            nn.Linear(hidden_dim, bottleneck_dim, bias=False),
+        )
+        self.norm = nn.LayerNorm(bottleneck_dim)
+        self.prototypes = nn.Linear(bottleneck_dim, out_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.mlp(x)
+        x = self.norm(x)
+        x = F.normalize(x, dim=-1)
+        return self.prototypes(x)
+
+
+class TeacherCLSExpressionDecoder(nn.Module):
+    def __init__(self, cfg: FoundationConfig):
+        super().__init__()
+        hidden_dim = int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024))
+        if hidden_dim <= 0:
+            raise ValueError("teacher_cls_decoder_hidden_dim must be > 0")
+        self.condition_on_obs = bool(getattr(cfg, "teacher_cls_decoder_condition_on_obs", False))
+        self.cls_proj = nn.Sequential(
+            nn.Linear(cfg.d_model, hidden_dim, bias=False),
+            nn.LayerNorm(hidden_dim),
+            nn.SiLU(),
+        )
+        self.obs_proj: Optional[nn.Sequential]
+        if self.condition_on_obs:
+            self.obs_proj = nn.Sequential(
+                nn.Linear(2, hidden_dim, bias=False),
+                nn.LayerNorm(hidden_dim),
+                nn.SiLU(),
+            )
+        else:
+            self.obs_proj = None
+        self.gene_proj = nn.Sequential(
+            nn.Linear(cfg.d_model, hidden_dim, bias=False),
+            nn.LayerNorm(hidden_dim),
+            nn.SiLU(),
+        )
+        self.out_proj = nn.Sequential(
+            nn.Linear(2 * hidden_dim, hidden_dim, bias=False),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, 1, bias=True),
+        )
+
+    def forward(
+        self,
+        cls_state: torch.Tensor,
+        gene_query_states: torch.Tensor,
+        obs_features: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        cls_hidden = self.cls_proj(cls_state)
+        if self.condition_on_obs:
+            if obs_features is None:
+                raise ValueError("teacher_cls_decoder_condition_on_obs=True requires obs_features")
+            cls_hidden = cls_hidden + self.obs_proj(
+                obs_features.to(device=cls_hidden.device, dtype=cls_hidden.dtype)
+            )
+        cls_hidden = cls_hidden.unsqueeze(1).expand(-1, gene_query_states.shape[1], -1)
+        gene_hidden = self.gene_proj(gene_query_states)
+        return self.out_proj(torch.cat([cls_hidden, gene_hidden], dim=-1)).squeeze(-1)
+
+
+class CellContrastivePredictModel(nn.Module):
+    def __init__(self, cfg: FoundationConfig):
+        super().__init__()
+        self.cfg = cfg
+        if cfg.encoder_layers < 1:
+            raise ValueError("encoder_layers must be >= 1")
+        if cfg.n_heads < 1:
+            raise ValueError("n_heads must be >= 1")
+        if cfg.d_model <= 0:
+            raise ValueError("d_model must be > 0")
+        if cfg.gene_vocab_size <= cfg.gene_cls_id:
+            raise ValueError("gene_vocab_size too small for gene_cls_id")
+        if cfg.bin_vocab_size <= 0:
+            raise ValueError("bin_vocab_size must be > 0")
+        if not (0.0 < float(cfg.student_temperature)):
+            raise ValueError("student_temperature must be > 0")
+        if not (0.0 < float(cfg.teacher_temperature)):
+            raise ValueError("teacher_temperature must be > 0")
+        if not (0.0 <= float(cfg.token_loss_weight)):
+            raise ValueError("token_loss_weight must be >= 0")
+        if float(getattr(cfg, "token_unmasked_loss_weight", 0.0)) < 0.0:
+            raise ValueError("token_unmasked_loss_weight must be >= 0")
+        if not (0.0 <= float(getattr(cfg, "teacher_cls_decoder_loss_weight", 0.0))):
+            raise ValueError("teacher_cls_decoder_loss_weight must be >= 0")
+        if int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024)) <= 0:
+            raise ValueError("teacher_cls_decoder_hidden_dim must be > 0")
+        if not isinstance(getattr(cfg, "teacher_cls_decoder_condition_on_obs", False), bool):
+            raise ValueError("teacher_cls_decoder_condition_on_obs must be a bool")
+        if float(getattr(cfg, "teacher_cls_decoder_rank_loss_alpha", 0.05)) < 0.0:
+            raise ValueError("teacher_cls_decoder_rank_loss_alpha must be >= 0")
+        if int(getattr(cfg, "teacher_cls_decoder_rank_num_pairs", 4096)) < 0:
+            raise ValueError("teacher_cls_decoder_rank_num_pairs must be >= 0")
+        if float(getattr(cfg, "teacher_student_decoder_distill_weight", 0.0)) < 0.0:
+            raise ValueError("teacher_student_decoder_distill_weight must be >= 0")
+        loss_mode = str(getattr(cfg, "teacher_student_decoder_loss_mode", "distribution")).strip().lower()
+        if loss_mode not in {"distribution", "mse"}:
+            raise ValueError("teacher_student_decoder_loss_mode must be one of {'distribution', 'mse'}")
+        cfg.teacher_student_decoder_loss_mode = loss_mode
+        if float(getattr(cfg, "teacher_student_decoder_rank_alpha", 0.05)) < 0.0:
+            raise ValueError("teacher_student_decoder_rank_alpha must be >= 0")
+        if float(getattr(cfg, "teacher_student_decoder_raw_mmd_weight", 0.1)) < 0.0:
+            raise ValueError("teacher_student_decoder_raw_mmd_weight must be >= 0")
+        if float(getattr(cfg, "teacher_student_decoder_raw_rank_weight", 0.01)) < 0.0:
+            raise ValueError("teacher_student_decoder_raw_rank_weight must be >= 0")
+        if float(getattr(cfg, "teacher_student_decoder_raw_mse_weight", 1.0)) < 0.0:
+            raise ValueError("teacher_student_decoder_raw_mse_weight must be >= 0")
+        if str(getattr(cfg, "expr_embedding_mode", "bin")).strip().lower() not in {"bin", "rank"}:
+            raise ValueError("expr_embedding_mode must be one of {'bin', 'rank'}")
+        if int(getattr(cfg, "expr_rank_mlp_hidden_dim", 256)) <= 0:
+            raise ValueError("expr_rank_mlp_hidden_dim must be > 0")
+        cross_modal_pairing_mode = str(getattr(cfg, "cross_modal_pairing_mode", "disabled")).strip().lower()
+        if cross_modal_pairing_mode not in {"disabled", "symmetric", "bin_to_gene_only"}:
+            raise ValueError("cross_modal_pairing_mode must be one of {'disabled', 'symmetric', 'bin_to_gene_only'}")
+        cfg.cross_modal_pairing_mode = cross_modal_pairing_mode
+
+        self.student_backbone = EncoderBackbone(cfg)
+        self.teacher_backbone = EncoderBackbone(cfg)
+        self.student_head = DinoHead(cfg)
+        self.teacher_head = DinoHead(cfg)
+        self.student_ibot_head = DinoHead(cfg)
+        self.teacher_ibot_head = DinoHead(cfg)
+        self.student_cls_expression_decoder: Optional[TeacherCLSExpressionDecoder] = None
+        self.teacher_cls_expression_decoder: Optional[TeacherCLSExpressionDecoder] = None
+        self._configure_teacher_cls_expression_decoder(
+            loss_weight=float(getattr(cfg, "teacher_cls_decoder_loss_weight", 0.0)),
+            hidden_dim=int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024)),
+        )
+
+        self.register_buffer("center", torch.zeros(1, int(cfg.dino_out_dim)))
+        self.register_buffer("token_center", torch.zeros(1, int(cfg.dino_out_dim)))
+        self.register_buffer("last_teacher_momentum", torch.tensor(float(cfg.teacher_momentum), dtype=torch.float32))
+
+        self._sync_teacher_from_student()
+        self._freeze_teacher()
+
+    def _configure_teacher_cls_expression_decoder(
+        self,
+        *,
+        loss_weight: float,
+        hidden_dim: int,
+        condition_on_obs: Optional[bool] = None,
+    ) -> None:
+        current_hidden = int(getattr(self.cfg, "teacher_cls_decoder_hidden_dim", hidden_dim))
+        current_condition = bool(getattr(self.cfg, "teacher_cls_decoder_condition_on_obs", False))
+        next_condition = current_condition if condition_on_obs is None else bool(condition_on_obs)
+        distill_enabled = float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0
+        self.cfg.teacher_cls_decoder_loss_weight = float(loss_weight)
+        self.cfg.teacher_cls_decoder_hidden_dim = int(hidden_dim)
+        self.cfg.teacher_cls_decoder_condition_on_obs = bool(next_condition)
+        if float(loss_weight) <= 0.0:
+            self.student_cls_expression_decoder = None
+            self.teacher_cls_expression_decoder = None
+            return
+        if distill_enabled:
+            if (
+                self.student_cls_expression_decoder is None
+                or int(current_hidden) != int(hidden_dim)
+                or current_condition != bool(next_condition)
+            ):
+                self.student_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg)
+            if (
+                self.teacher_cls_expression_decoder is None
+                or int(current_hidden) != int(hidden_dim)
+                or current_condition != bool(next_condition)
+            ):
+                self.teacher_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg)
+            self.teacher_cls_expression_decoder.load_state_dict(self.student_cls_expression_decoder.state_dict())
+            for param in self.teacher_cls_expression_decoder.parameters():
+                param.requires_grad = False
+            return
+        self.student_cls_expression_decoder = None
+        if (
+            self.teacher_cls_expression_decoder is None
+            or int(current_hidden) != int(hidden_dim)
+            or current_condition != bool(next_condition)
+        ):
+            self.teacher_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg)
+
+    def configure_teacher_cls_expression_decoder(
+        self,
+        *,
+        loss_weight: float,
+        hidden_dim: int,
+        condition_on_obs: Optional[bool] = None,
+    ) -> None:
+        self._configure_teacher_cls_expression_decoder(
+            loss_weight=float(loss_weight),
+            hidden_dim=int(hidden_dim),
+            condition_on_obs=condition_on_obs,
+        )
+
+    def _freeze_teacher(self) -> None:
+        for param in self.teacher_backbone.parameters():
+            param.requires_grad = False
+        for param in self.teacher_head.parameters():
+            param.requires_grad = False
+        for param in self.teacher_ibot_head.parameters():
+            param.requires_grad = False
+        if (
+            float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0
+            and self.teacher_cls_expression_decoder is not None
+        ):
+            for param in self.teacher_cls_expression_decoder.parameters():
+                param.requires_grad = False
+
+    def _sync_teacher_from_student(self) -> None:
+        self.teacher_backbone.load_state_dict(self.student_backbone.state_dict())
+        self.teacher_head.load_state_dict(self.student_head.state_dict())
+        self.teacher_ibot_head.load_state_dict(self.student_ibot_head.state_dict())
+        if self.student_cls_expression_decoder is not None and self.teacher_cls_expression_decoder is not None:
+            self.teacher_cls_expression_decoder.load_state_dict(self.student_cls_expression_decoder.state_dict())
+
+    def _pool_cell_repr(self, token_states: torch.Tensor) -> torch.Tensor:
+        return token_states[:, 0, :]
+
+    def _pool_instance_repr(self, token_states: torch.Tensor) -> torch.Tensor:
+        return self._pool_cell_repr(token_states)
+
+    def _select_backbone(self, backbone: Optional[str] = None) -> EncoderBackbone:
+        chosen = backbone
+        if chosen is None:
+            chosen = "teacher" if bool(getattr(self.cfg, "inference_use_teacher", True)) else "student"
+        if chosen == "teacher":
+            return self.teacher_backbone
+        if chosen == "student":
+            return self.student_backbone
+        raise ValueError(f"Unknown backbone={backbone!r}")
+
+    def _encode_tokens(
+        self,
+        gene_ids: torch.Tensor,
+        bin_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        *,
+        use_bin_embeddings: bool = True,
+        backbone: Optional[str] = None,
+    ) -> torch.Tensor:
+        module = self._select_backbone(backbone)
+        return module(
+            gene_ids,
+            bin_ids,
+            attention_mask,
+            use_bin_embeddings=use_bin_embeddings,
+        )
+
+    def _encode(
+        self,
+        gene_ids: torch.Tensor,
+        bin_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        *,
+        use_bin_embeddings: bool = True,
+        backbone: Optional[str] = None,
+    ) -> torch.Tensor:
+        return self._pool_cell_repr(
+            self._encode_tokens(
+                gene_ids,
+                bin_ids,
+                attention_mask,
+                use_bin_embeddings=use_bin_embeddings,
+                backbone=backbone,
+            )
+        )
+
+    def _encode_view_stack(
+        self,
+        gene_ids: torch.Tensor,
+        bin_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        view_use_bin: torch.Tensor,
+        *,
+        backbone_name: str,
+        head: nn.Module,
+        capture_token_views: Optional[Tuple[int, ...]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[int, torch.Tensor]]:
+        num_views = int(gene_ids.shape[1])
+        pooled_views = []
+        logits_views = []
+        token_view_set = set(capture_token_views or ())
+        captured_token_states: Dict[int, torch.Tensor] = {}
+        for view_idx in range(num_views):
+            use_bin = bool(view_use_bin[view_idx].item())
+            view_attention_mask = attention_mask[:, view_idx, :]
+            valid_positions = torch.nonzero(view_attention_mask, as_tuple=False)
+            if valid_positions.numel() == 0:
+                raise ValueError(f"view {view_idx} contains no valid tokens")
+            max_valid_len = int(valid_positions[:, -1].max().item()) + 1
+            tokens = self._encode_tokens(
+                gene_ids[:, view_idx, :max_valid_len],
+                bin_ids[:, view_idx, :max_valid_len],
+                view_attention_mask[:, :max_valid_len],
+                use_bin_embeddings=use_bin,
+                backbone=backbone_name,
+            )
+            pooled = self._pool_cell_repr(tokens)
+            pooled_views.append(pooled)
+            logits_views.append(head(pooled))
+            if view_idx in token_view_set:
+                captured_token_states[view_idx] = tokens
+        return torch.stack(pooled_views, dim=1), torch.stack(logits_views, dim=1), captured_token_states
+
+    def _distributed_batch_mean(self, x: torch.Tensor) -> torch.Tensor:
+        mean = x.mean(dim=0, keepdim=True)
+        if _dist_is_initialized():
+            mean = mean.clone()
+            torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM)
+            mean /= float(torch.distributed.get_world_size())
+        return mean
+
+    def _distributed_optional_rank_mean(
+        self,
+        x: Optional[torch.Tensor],
+        *,
+        output_dim: int,
+        device: torch.device,
+    ) -> Optional[torch.Tensor]:
+        if x is None or int(x.numel()) == 0 or int(x.shape[0]) == 0:
+            mean = torch.zeros((1, int(output_dim)), device=device, dtype=torch.float32)
+            active = torch.zeros([], device=device, dtype=torch.float32)
+        else:
+            mean = x.reshape(-1, x.shape[-1]).detach().float().mean(dim=0, keepdim=True)
+            active = torch.ones([], device=mean.device, dtype=torch.float32)
+
+        if _dist_is_initialized():
+            torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM)
+            torch.distributed.all_reduce(active, op=torch.distributed.ReduceOp.SUM)
+
+        if float(active.item()) <= 0.0:
+            return None
+        return mean / active.clamp_min(1.0)
+
+    def _update_center(self, teacher_logits: torch.Tensor) -> None:
+        batch_center = self._distributed_batch_mean(teacher_logits.reshape(-1, teacher_logits.shape[-1]).detach())
+        batch_center = batch_center.to(device=self.center.device, dtype=self.center.dtype)
+        momentum = float(getattr(self.cfg, "center_momentum", 0.9))
+        self.center.mul_(momentum).add_(batch_center, alpha=1.0 - momentum)
+
+    def _update_token_center(self, teacher_logits: Optional[torch.Tensor]) -> None:
+        batch_center = self._distributed_optional_rank_mean(
+            teacher_logits,
+            output_dim=int(self.token_center.shape[-1]),
+            device=self.token_center.device,
+        )
+        if batch_center is None:
+            return
+        batch_center = batch_center.to(device=self.token_center.device, dtype=self.token_center.dtype)
+        momentum = float(getattr(self.cfg, "token_center_momentum", getattr(self.cfg, "center_momentum", 0.9)))
+        self.token_center.mul_(momentum).add_(batch_center, alpha=1.0 - momentum)
+
+    def _teacher_probs(self, logits: torch.Tensor) -> torch.Tensor:
+        temperature = float(getattr(self.cfg, "teacher_temperature", 0.04))
+        logits_f = logits.float()
+        center_f = self.center.to(device=logits_f.device, dtype=logits_f.dtype)
+        return F.softmax((logits_f - center_f) / temperature, dim=-1)
+
+    def _teacher_token_probs(self, logits: torch.Tensor) -> torch.Tensor:
+        temperature = float(getattr(self.cfg, "teacher_temperature", 0.04))
+        logits_f = logits.float()
+        center_f = self.token_center.to(device=logits_f.device, dtype=logits_f.dtype)
+        return F.softmax((logits_f - center_f) / temperature, dim=-1)
+
+    def _student_log_probs(self, logits: torch.Tensor) -> torch.Tensor:
+        temperature = float(getattr(self.cfg, "student_temperature", 0.10))
+        return F.log_softmax(logits.float() / temperature, dim=-1)
+
+    def _cross_entropy_target(self, student_logits: torch.Tensor, teacher_probs: torch.Tensor) -> torch.Tensor:
+        student_log_probs = self._student_log_probs(student_logits)
+        return (-teacher_probs * student_log_probs).sum(dim=-1).mean()
+
+    @staticmethod
+    def _mean_pairwise_abs_1d(values: torch.Tensor) -> torch.Tensor:
+        n = int(values.numel())
+        if n <= 1:
+            return values.new_zeros(())
+        sorted_values = values.sort().values
+        coeff = torch.arange(n, device=values.device, dtype=values.dtype) * 2 - (n - 1)
+        return (2.0 * (coeff * sorted_values).sum()) / float(n * n)
+
+    @staticmethod
+    def _mean_cross_abs_1d(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        nx = int(x.numel())
+        ny = int(y.numel())
+        if nx == 0 or ny == 0:
+            return x.new_zeros(())
+        y_sorted = y.sort().values
+        prefix = torch.cat([y_sorted.new_zeros((1,)), y_sorted.cumsum(dim=0)])
+        right = torch.searchsorted(y_sorted, x, right=False)
+        left_sum = prefix[right]
+        total_sum = prefix[-1]
+        right_f = right.to(dtype=x.dtype)
+        ny_f = float(ny)
+        abs_sum = x * right_f - left_sum + (total_sum - left_sum) - x * (ny_f - right_f)
+        return abs_sum.mean() / ny_f
+
+    def _gene_mmd_loss(self, pred_expr: torch.Tensor, target_expr: torch.Tensor, target_mask: torch.Tensor) -> torch.Tensor:
+        mask = target_mask.to(dtype=torch.bool)
+        losses = []
+        for batch_idx in range(int(pred_expr.shape[0])):
+            valid = mask[batch_idx]
+            if not bool(valid.any().item()):
+                continue
+            pred = pred_expr[batch_idx, valid].float()
+            target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype)
+            if int(pred.numel()) <= 1:
+                losses.append(F.smooth_l1_loss(pred, target, reduction="mean"))
+                continue
+            cross = self._mean_cross_abs_1d(pred, target)
+            pred_pair = self._mean_pairwise_abs_1d(pred)
+            target_pair = self._mean_pairwise_abs_1d(target)
+            losses.append((2.0 * cross - pred_pair - target_pair).clamp_min(0.0))
+        if not losses:
+            return pred_expr.new_zeros(())
+        return torch.stack(losses).mean()
+
+    def _rank_loss(self, pred_expr: torch.Tensor, target_expr: torch.Tensor, target_mask: torch.Tensor) -> torch.Tensor:
+        num_pairs = int(getattr(self.cfg, "teacher_cls_decoder_rank_num_pairs", 4096))
+        if num_pairs <= 0:
+            return pred_expr.new_zeros(())
+        mask = target_mask.to(dtype=torch.bool)
+        losses = []
+        for batch_idx in range(int(pred_expr.shape[0])):
+            valid = mask[batch_idx]
+            if int(valid.sum().item()) <= 1:
+                continue
+            pred = pred_expr[batch_idx, valid].float()
+            target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype)
+            n = int(pred.numel())
+            device = pred.device
+            expressed = torch.nonzero(target > 0, as_tuple=False).flatten()
+            zero = torch.nonzero(target <= 0, as_tuple=False).flatten()
+            pair_i = []
+            pair_j = []
+            ez_pairs = num_pairs // 2 if int(expressed.numel()) > 0 and int(zero.numel()) > 0 else 0
+            if ez_pairs > 0:
+                pair_i.append(expressed[torch.randint(int(expressed.numel()), (ez_pairs,), device=device)])
+                pair_j.append(zero[torch.randint(int(zero.numel()), (ez_pairs,), device=device)])
+            random_pairs = num_pairs - ez_pairs
+            if random_pairs > 0:
+                pair_i.append(torch.randint(n, (random_pairs,), device=device))
+                pair_j.append(torch.randint(n, (random_pairs,), device=device))
+            i = torch.cat(pair_i, dim=0)
+            j = torch.cat(pair_j, dim=0)
+            true_diff = target[i] - target[j]
+            true_sign = torch.sign(true_diff)
+            valid_pair = true_sign != 0
+            if not bool(valid_pair.any().item()):
+                continue
+            pred_diff = pred[i] - pred[j]
+            losses.append(F.softplus(-true_sign[valid_pair] * pred_diff[valid_pair]).mean())
+        if not losses:
+            return pred_expr.new_zeros(())
+        return torch.stack(losses).mean()
+
+    def _rank_loss_random_pairs(
+        self,
+        pred_expr: torch.Tensor,
+        target_expr: torch.Tensor,
+        target_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        num_pairs = int(getattr(self.cfg, "teacher_cls_decoder_rank_num_pairs", 4096))
+        if num_pairs <= 0:
+            return pred_expr.new_zeros(())
+        mask = target_mask.to(dtype=torch.bool)
+        losses = []
+        for batch_idx in range(int(pred_expr.shape[0])):
+            valid = mask[batch_idx]
+            if int(valid.sum().item()) <= 1:
+                continue
+            pred = pred_expr[batch_idx, valid].float()
+            target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype)
+            n = int(pred.numel())
+            device = pred.device
+            i = torch.randint(n, (num_pairs,), device=device)
+            j = torch.randint(n, (num_pairs,), device=device)
+            true_sign = torch.sign(target[i] - target[j])
+            valid_pair = true_sign != 0
+            if not bool(valid_pair.any().item()):
+                continue
+            pred_diff = pred[i] - pred[j]
+            losses.append(F.softplus(-true_sign[valid_pair] * pred_diff[valid_pair]).mean())
+        if not losses:
+            return pred_expr.new_zeros(())
+        return torch.stack(losses).mean()
+
+    def _masked_mse_loss(
+        self,
+        pred_expr: torch.Tensor,
+        target_expr: torch.Tensor,
+        target_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        mask = target_mask.to(device=pred_expr.device, dtype=torch.bool)
+        if not bool(mask.any().item()):
+            return pred_expr.new_zeros(())
+        target = target_expr.to(device=pred_expr.device, dtype=pred_expr.dtype)
+        return F.mse_loss(pred_expr[mask].float(), target[mask].float(), reduction="mean")
+
+    def _teacher_cls_decoder_loss(
+        self,
+        *,
+        teacher_gene_ids: torch.Tensor,
+        teacher_bin_ids: torch.Tensor,
+        teacher_attention_mask: torch.Tensor,
+        teacher_view_use_bin: torch.Tensor,
+        target_gene_ids: torch.Tensor,
+        target_expr: torch.Tensor,
+        target_mask: torch.Tensor,
+        obs_features: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.teacher_cls_expression_decoder is None:
+            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)
+        if target_gene_ids.numel() == 0 or not bool(target_mask.any().item()):
+            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)
+
+        decoder_tokens = self._encode_tokens(
+            teacher_gene_ids[:, 0, :],
+            teacher_bin_ids[:, 0, :],
+            teacher_attention_mask[:, 0, :],
+            use_bin_embeddings=bool(teacher_view_use_bin[0].item()),
+            backbone="student",
+        )
+        decoder_cls = self._pool_cell_repr(decoder_tokens)
+        gene_queries = self.student_backbone.tokenizer.gene_emb(target_gene_ids)
+        pred_expr = self.teacher_cls_expression_decoder(decoder_cls, gene_queries, obs_features=obs_features)
+
+        gene_mmd = self._gene_mmd_loss(pred_expr, target_expr, target_mask)
+        rank_alpha = float(getattr(self.cfg, "teacher_cls_decoder_rank_loss_alpha", 0.05))
+        if rank_alpha <= 0.0:
+            return gene_mmd
+        return gene_mmd + rank_alpha * self._rank_loss(pred_expr, target_expr, target_mask)
+
+    def _teacher_student_cls_decoder_loss(
+        self,
+        *,
+        teacher_gene_ids: torch.Tensor,
+        teacher_bin_ids: torch.Tensor,
+        teacher_attention_mask: torch.Tensor,
+        teacher_view_use_bin: torch.Tensor,
+        student_gene_ids: torch.Tensor,
+        student_bin_ids: torch.Tensor,
+        student_attention_mask: torch.Tensor,
+        student_view_use_bin: torch.Tensor,
+        target_gene_ids: torch.Tensor,
+        target_expr: torch.Tensor,
+        target_mask: torch.Tensor,
+        obs_features: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.student_cls_expression_decoder is None or self.teacher_cls_expression_decoder is None:
+            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)
+        if target_gene_ids.numel() == 0 or not bool(target_mask.any().item()):
+            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)
+
+        student_tokens = self._encode_tokens(
+            student_gene_ids[:, 0, :],
+            student_bin_ids[:, 0, :],
+            student_attention_mask[:, 0, :],
+            use_bin_embeddings=bool(student_view_use_bin[0].item()),
+            backbone="student",
+        )
+        student_cls = self._pool_cell_repr(student_tokens)
+        student_gene_queries = self.student_backbone.tokenizer.gene_emb(target_gene_ids)
+        student_scores = self.student_cls_expression_decoder(
+            student_cls,
+            student_gene_queries,
+            obs_features=obs_features,
+        )
+
+        with torch.no_grad():
+            teacher_tokens = self._encode_tokens(
+                teacher_gene_ids[:, 0, :],
+                teacher_bin_ids[:, 0, :],
+                teacher_attention_mask[:, 0, :],
+                use_bin_embeddings=bool(teacher_view_use_bin[0].item()),
+                backbone="teacher",
+            )
+            teacher_cls = self._pool_cell_repr(teacher_tokens)
+            teacher_gene_queries = self.teacher_backbone.tokenizer.gene_emb(target_gene_ids)
+            teacher_scores = self.teacher_cls_expression_decoder(
+                teacher_cls,
+                teacher_gene_queries,
+                obs_features=obs_features,
+            )
+
+        distill_weight = float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0))
+        loss_mode = str(getattr(self.cfg, "teacher_student_decoder_loss_mode", "distribution")).strip().lower()
+        if loss_mode == "mse":
+            raw_mse_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_mse_weight", 1.0))
+            loss = student_scores.new_zeros(())
+            if distill_weight > 0.0:
+                loss = loss + distill_weight * self._masked_mse_loss(student_scores, teacher_scores, target_mask)
+            if raw_mse_weight > 0.0:
+                loss = loss + raw_mse_weight * self._masked_mse_loss(student_scores, target_expr, target_mask)
+            return loss
+        if loss_mode != "distribution":
+            raise ValueError(f"Unsupported teacher_student_decoder_loss_mode={loss_mode!r}")
+
+        rank_alpha = float(getattr(self.cfg, "teacher_student_decoder_rank_alpha", 0.05))
+        raw_mmd_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_mmd_weight", 0.1))
+        raw_rank_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_rank_weight", 0.01))
+
+        loss = student_scores.new_zeros(())
+        if distill_weight > 0.0:
+            loss = loss + distill_weight * self._gene_mmd_loss(student_scores, teacher_scores, target_mask)
+            if rank_alpha > 0.0:
+                loss = loss + distill_weight * rank_alpha * self._rank_loss_random_pairs(student_scores, teacher_scores, target_mask)
+        if raw_mmd_weight > 0.0:
+            loss = loss + raw_mmd_weight * self._gene_mmd_loss(student_scores, target_expr, target_mask)
+        if raw_rank_weight > 0.0:
+            loss = loss + raw_rank_weight * self._rank_loss(student_scores, target_expr, target_mask)
+        return loss
+
+    def _allow_cross_mode_pair(self, *, teacher_is_bin: bool, student_is_bin: bool) -> bool:
+        mode = str(getattr(self.cfg, "cross_modal_pairing_mode", "bin_to_gene_only"))
+        if mode == "disabled":
+            return False
+        if mode == "symmetric":
+            return teacher_is_bin != student_is_bin
+        if mode == "bin_to_gene_only":
+            return bool(teacher_is_bin) and (not bool(student_is_bin))
+        raise ValueError(f"Unknown cross_modal_pairing_mode={mode!r}")
+
+    def _pairwise_dino_loss(
+        self,
+        student_logits: torch.Tensor,
+        student_view_use_bin: torch.Tensor,
+        teacher_logits: torch.Tensor,
+        teacher_view_use_bin: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        same_losses = []
+        cross_losses = []
+        cross_weight = float(getattr(self.cfg, "cross_modal_dino_weight", 0.0))
+
+        num_teacher_views = int(teacher_logits.shape[1])
+        num_student_views = int(student_logits.shape[1])
+
+        for teacher_idx in range(num_teacher_views):
+            teacher_probs = self._teacher_probs(teacher_logits[:, teacher_idx, :]).detach()
+            teacher_is_bin = bool(teacher_view_use_bin[teacher_idx].item())
+            for student_idx in range(num_student_views):
+                student_is_bin = bool(student_view_use_bin[student_idx].item())
+                pair_loss = self._cross_entropy_target(
+                    student_logits[:, student_idx, :],
+                    teacher_probs,
+                )
+                if student_is_bin == teacher_is_bin:
+                    same_losses.append(pair_loss)
+                elif cross_weight > 0.0 and self._allow_cross_mode_pair(
+                    teacher_is_bin=teacher_is_bin,
+                    student_is_bin=student_is_bin,
+                ):
+                    cross_losses.append(pair_loss)
+
+        device = student_logits.device
+        same_loss = (
+            torch.stack(same_losses).mean()
+            if same_losses
+            else torch.zeros([], device=device, dtype=torch.float32)
+        )
+        cross_loss = (
+            torch.stack(cross_losses).mean()
+            if cross_losses
+            else torch.zeros([], device=device, dtype=torch.float32)
+        )
+        total = same_loss + float(cross_weight) * cross_loss
+        return total, same_loss, cross_loss
+
+    def _matched_token_states_for_pair(
+        self,
+        *,
+        teacher_tokens: torch.Tensor,
+        teacher_gene_ids: torch.Tensor,
+        teacher_attention_mask: torch.Tensor,
+        student_tokens: torch.Tensor,
+        student_gene_ids: torch.Tensor,
+        student_attention_mask: torch.Tensor,
+        student_ibot_mask: torch.Tensor,
+        use_masked_tokens: bool = True,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        teacher_states: List[torch.Tensor] = []
+        student_states: List[torch.Tensor] = []
+
+        batch_size = int(teacher_tokens.shape[0])
+        hidden_dim = int(teacher_tokens.shape[-1])
+        for batch_idx in range(batch_size):
+            teacher_valid = teacher_attention_mask[batch_idx].bool().clone()
+            student_attention_valid = student_attention_mask[batch_idx].bool()
+            student_mask = student_ibot_mask[batch_idx].bool()
+            if use_masked_tokens:
+                student_valid = (student_attention_valid & student_mask).clone()
+            else:
+                student_valid = (student_attention_valid & (~student_mask)).clone()
+
+            if teacher_valid.numel() <= 1 or student_valid.numel() <= 1:
+                continue
+
+            teacher_valid[0] = False
+            student_valid[0] = False
+            if not teacher_valid.any().item() or not student_valid.any().item():
+                continue
+
+            teacher_positions = torch.nonzero(teacher_valid, as_tuple=False).squeeze(-1)
+            student_positions = torch.nonzero(student_valid, as_tuple=False).squeeze(-1)
+            teacher_genes = teacher_gene_ids[batch_idx, teacher_positions]
+            student_genes = student_gene_ids[batch_idx, student_positions]
+
+            matches = student_genes.unsqueeze(1) == teacher_genes.unsqueeze(0)
+            matched_student = matches.any(dim=1)
+            if not matched_student.any().item():
+                continue
+
+            matched_student_positions = student_positions[matched_student]
+            matched_teacher_rel_pos = matches[matched_student].to(dtype=torch.int64).argmax(dim=1)
+            matched_teacher_positions = teacher_positions[matched_teacher_rel_pos]
+
+            teacher_states.append(teacher_tokens[batch_idx, matched_teacher_positions, :])
+            student_states.append(student_tokens[batch_idx, matched_student_positions, :])
+
+        if teacher_states:
+            return torch.cat(teacher_states, dim=0), torch.cat(student_states, dim=0)
+
+        empty = teacher_tokens.new_empty((0, hidden_dim))
+        return empty, empty
+
+    def _pairwise_token_loss(
+        self,
+        *,
+        student_token_views: Dict[int, torch.Tensor],
+        student_gene_ids: torch.Tensor,
+        student_attention_mask: torch.Tensor,
+        student_ibot_mask: torch.Tensor,
+        student_view_use_bin: torch.Tensor,
+        student_view_is_global: torch.Tensor,
+        teacher_token_views: Dict[int, torch.Tensor],
+        teacher_gene_ids: torch.Tensor,
+        teacher_attention_mask: torch.Tensor,
+        teacher_view_use_bin: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        same_losses = []
+        cross_losses = []
+        unmasked_same_losses = []
+        unmasked_cross_losses = []
+        teacher_logits_for_center: List[torch.Tensor] = []
+        cross_weight = float(getattr(self.cfg, "cross_modal_dino_weight", 0.0))
+        unmasked_weight = float(getattr(self.cfg, "token_unmasked_loss_weight", 0.0))
+
+        for teacher_idx, teacher_tokens in teacher_token_views.items():
+            teacher_is_bin = bool(teacher_view_use_bin[teacher_idx].item())
+            for student_idx, student_tokens in student_token_views.items():
+                if not bool(student_view_is_global[student_idx].item()):
+                    continue
+                if not student_ibot_mask[:, student_idx, :].any().item():
+                    continue
+
+                student_is_bin = bool(student_view_use_bin[student_idx].item())
+                if student_is_bin == teacher_is_bin:
+                    allow_pair = True
+                elif cross_weight > 0.0:
+                    allow_pair = self._allow_cross_mode_pair(
+                        teacher_is_bin=teacher_is_bin,
+                        student_is_bin=student_is_bin,
+                    )
+                else:
+                    allow_pair = False
+                if not allow_pair:
+                    continue
+
+                teacher_states, student_states = self._matched_token_states_for_pair(
+                    teacher_tokens=teacher_tokens,
+                    teacher_gene_ids=teacher_gene_ids[:, teacher_idx, :],
+                    teacher_attention_mask=teacher_attention_mask[:, teacher_idx, :],
+                    student_tokens=student_tokens,
+                    student_gene_ids=student_gene_ids[:, student_idx, :],
+                    student_attention_mask=student_attention_mask[:, student_idx, :],
+                    student_ibot_mask=student_ibot_mask[:, student_idx, :],
+                    use_masked_tokens=True,
+                )
+                if teacher_states.shape[0] > 0:
+                    teacher_logits = self.teacher_ibot_head(teacher_states)
+                    teacher_probs = self._teacher_token_probs(teacher_logits).detach()
+                    student_logits = self.student_ibot_head(student_states)
+                    pair_loss = self._cross_entropy_target(student_logits, teacher_probs)
+                    teacher_logits_for_center.append(teacher_logits.detach())
+
+                    if student_is_bin == teacher_is_bin:
+                        same_losses.append(pair_loss)
+                    else:
+                        cross_losses.append(pair_loss)
+
+                if unmasked_weight > 0.0:
+                    unmasked_teacher_states, unmasked_student_states = self._matched_token_states_for_pair(
+                        teacher_tokens=teacher_tokens,
+                        teacher_gene_ids=teacher_gene_ids[:, teacher_idx, :],
+                        teacher_attention_mask=teacher_attention_mask[:, teacher_idx, :],
+                        student_tokens=student_tokens,
+                        student_gene_ids=student_gene_ids[:, student_idx, :],
+                        student_attention_mask=student_attention_mask[:, student_idx, :],
+                        student_ibot_mask=student_ibot_mask[:, student_idx, :],
+                        use_masked_tokens=False,
+                    )
+                    if unmasked_teacher_states.shape[0] > 0:
+                        unmasked_teacher_logits = self.teacher_ibot_head(unmasked_teacher_states)
+                        unmasked_teacher_probs = self._teacher_token_probs(unmasked_teacher_logits).detach()
+                        unmasked_student_logits = self.student_ibot_head(unmasked_student_states)
+                        unmasked_pair_loss = self._cross_entropy_target(
+                            unmasked_student_logits,
+                            unmasked_teacher_probs,
+                        )
+                        if student_is_bin == teacher_is_bin:
+                            unmasked_same_losses.append(unmasked_pair_loss)
+                        else:
+                            unmasked_cross_losses.append(unmasked_pair_loss)
+
+        device = student_gene_ids.device
+        same_loss = (
+            torch.stack(same_losses).mean()
+            if same_losses
+            else torch.zeros([], device=device, dtype=torch.float32)
+        )
+        cross_loss = (
+            torch.stack(cross_losses).mean()
+            if cross_losses
+            else torch.zeros([], device=device, dtype=torch.float32)
+        )
+        unmasked_same_loss = (
+            torch.stack(unmasked_same_losses).mean()
+            if unmasked_same_losses
+            else torch.zeros([], device=device, dtype=torch.float32)
+        )
+        unmasked_cross_loss = (
+            torch.stack(unmasked_cross_losses).mean()
+            if unmasked_cross_losses
+            else torch.zeros([], device=device, dtype=torch.float32)
+        )
+        masked_total = same_loss + float(cross_weight) * cross_loss
+        unmasked_total = unmasked_same_loss + float(cross_weight) * unmasked_cross_loss
+        total = masked_total + unmasked_weight * unmasked_total
+        teacher_logits_center = torch.cat(teacher_logits_for_center, dim=0) if teacher_logits_for_center else None
+        self._last_token_unmasked_loss = unmasked_total.detach()
+        self._last_token_unmasked_same_mode_loss = unmasked_same_loss.detach()
+        self._last_token_unmasked_cross_mode_loss = unmasked_cross_loss.detach()
+        return total, same_loss, cross_loss, teacher_logits_center
+
+    def _teacher_momentum_value(self, global_step: int, max_steps: int) -> float:
+        start = float(getattr(self.cfg, "teacher_momentum", 0.996))
+        end = float(getattr(self.cfg, "teacher_momentum_end", 1.0))
+        if max_steps <= 1:
+            return start
+        progress = min(max(float(global_step), 0.0), float(max_steps - 1)) / float(max_steps - 1)
+        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+        return end - (end - start) * cosine
+
+    def _ema_update_module(self, student_module: nn.Module, teacher_module: nn.Module, momentum: float) -> None:
+        with ExitStack() as stack:
+            if FSDP is not None:
+                for wrapped in _leaf_fsdp_modules(student_module):
+                    stack.enter_context(FSDP.summon_full_params(wrapped, recurse=False, writeback=False))
+                for wrapped in _leaf_fsdp_modules(teacher_module):
+                    stack.enter_context(FSDP.summon_full_params(wrapped, recurse=False, writeback=True))
+
+            student_params = dict(student_module.named_parameters())
+            teacher_params = dict(teacher_module.named_parameters())
+            if student_params.keys() != teacher_params.keys():
+                missing_student = sorted(set(teacher_params.keys()) - set(student_params.keys()))
+                missing_teacher = sorted(set(student_params.keys()) - set(teacher_params.keys()))
+                raise RuntimeError(
+                    "Student/teacher parameter names do not match for EMA update: "
+                    f"missing_student={missing_student[:5]}, missing_teacher={missing_teacher[:5]}"
+                )
+
+            for name, teacher_param in teacher_params.items():
+                student_param = student_params[name]
+                if teacher_param.data.shape != student_param.data.shape:
+                    raise RuntimeError(
+                        f"EMA shape mismatch for {name}: teacher={tuple(teacher_param.shape)} "
+                        f"student={tuple(student_param.shape)}"
+                    )
+                teacher_param.data.mul_(momentum).add_(student_param.data, alpha=1.0 - momentum)
+
+    @torch.no_grad()
+    def update_teacher(self, *, global_step: int, max_steps: int) -> float:
+        momentum = self._teacher_momentum_value(global_step=global_step, max_steps=max_steps)
+        self.last_teacher_momentum.fill_(float(momentum))
+
+        self._ema_update_module(self.student_backbone, self.teacher_backbone, momentum)
+        self._ema_update_module(self.student_head, self.teacher_head, momentum)
+        self._ema_update_module(self.student_ibot_head, self.teacher_ibot_head, momentum)
+        if self.student_cls_expression_decoder is not None and self.teacher_cls_expression_decoder is not None:
+            self._ema_update_module(
+                self.student_cls_expression_decoder,
+                self.teacher_cls_expression_decoder,
+                momentum,
+            )
+        return momentum
+
+    def forward(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        student_gene_ids = batch["student_gene_ids"]
+        student_bin_ids = batch["student_bin_ids"]
+        student_attention_mask = batch["student_attention_mask"]
+        student_ibot_mask = batch["student_ibot_mask"]
+        student_view_use_bin = batch["student_view_use_bin"]
+        student_view_is_global = batch["student_view_is_global"]
+
+        teacher_gene_ids = batch["teacher_gene_ids"]
+        teacher_bin_ids = batch["teacher_bin_ids"]
+        teacher_attention_mask = batch["teacher_attention_mask"]
+        teacher_view_use_bin = batch["teacher_view_use_bin"]
+
+        student_capture_views = tuple(
+            idx for idx in range(int(student_view_is_global.numel())) if bool(student_view_is_global[idx].item())
+        )
+        teacher_capture_views = tuple(range(int(teacher_view_use_bin.numel())))
+
+        _, student_logits, student_token_views = self._encode_view_stack(
+            student_gene_ids,
+            student_bin_ids,
+            student_attention_mask,
+            student_view_use_bin,
+            backbone_name="student",
+            head=self.student_head,
+            capture_token_views=student_capture_views,
+        )
+
+        with torch.no_grad():
+            _, teacher_logits, teacher_token_views = self._encode_view_stack(
+                teacher_gene_ids,
+                teacher_bin_ids,
+                teacher_attention_mask,
+                teacher_view_use_bin,
+                backbone_name="teacher",
+                head=self.teacher_head,
+                capture_token_views=teacher_capture_views,
+            )
+
+        loss_dino, loss_same_mode, loss_cross_mode = self._pairwise_dino_loss(
+            student_logits,
+            student_view_use_bin,
+            teacher_logits,
+            teacher_view_use_bin,
+        )
+
+        token_loss_weight = float(getattr(self.cfg, "token_loss_weight", 1.0))
+        if token_loss_weight > 0.0:
+            loss_token, loss_token_same_mode, loss_token_cross_mode, teacher_token_center_logits = self._pairwise_token_loss(
+                student_token_views=student_token_views,
+                student_gene_ids=student_gene_ids,
+                student_attention_mask=student_attention_mask,
+                student_ibot_mask=student_ibot_mask,
+                student_view_use_bin=student_view_use_bin,
+                student_view_is_global=student_view_is_global,
+                teacher_token_views=teacher_token_views,
+                teacher_gene_ids=teacher_gene_ids,
+                teacher_attention_mask=teacher_attention_mask,
+                teacher_view_use_bin=teacher_view_use_bin,
+            )
+        else:
+            zero = torch.zeros([], device=student_logits.device, dtype=torch.float32)
+            loss_token = zero
+            loss_token_same_mode = zero
+            loss_token_cross_mode = zero
+            self._last_token_unmasked_loss = zero.detach()
+            self._last_token_unmasked_same_mode_loss = zero.detach()
+            self._last_token_unmasked_cross_mode_loss = zero.detach()
+            teacher_token_center_logits = None
+
+        decoder_loss_weight = float(getattr(self.cfg, "teacher_cls_decoder_loss_weight", 0.0))
+        if decoder_loss_weight > 0.0:
+            if float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0:
+                loss_decoder = self._teacher_student_cls_decoder_loss(
+                    teacher_gene_ids=teacher_gene_ids,
+                    teacher_bin_ids=teacher_bin_ids,
+                    teacher_attention_mask=teacher_attention_mask,
+                    teacher_view_use_bin=teacher_view_use_bin,
+                    student_gene_ids=student_gene_ids,
+                    student_bin_ids=student_bin_ids,
+                    student_attention_mask=student_attention_mask,
+                    student_view_use_bin=student_view_use_bin,
+                    target_gene_ids=batch["teacher_cls_decoder_target_gene_ids"],
+                    target_expr=batch["teacher_cls_decoder_target_expr"],
+                    target_mask=batch["teacher_cls_decoder_target_mask"],
+                    obs_features=batch.get("teacher_cls_decoder_obs_features"),
+                )
+            else:
+                loss_decoder = self._teacher_cls_decoder_loss(
+                    teacher_gene_ids=teacher_gene_ids,
+                    teacher_bin_ids=teacher_bin_ids,
+                    teacher_attention_mask=teacher_attention_mask,
+                    teacher_view_use_bin=teacher_view_use_bin,
+                    target_gene_ids=batch["teacher_cls_decoder_target_gene_ids"],
+                    target_expr=batch["teacher_cls_decoder_target_expr"],
+                    target_mask=batch["teacher_cls_decoder_target_mask"],
+                    obs_features=batch.get("teacher_cls_decoder_obs_features"),
+                )
+        else:
+            loss_decoder = torch.zeros([], device=student_logits.device, dtype=torch.float32)
+
+        total_loss = loss_dino + token_loss_weight * loss_token + decoder_loss_weight * loss_decoder
+
+        with torch.no_grad():
+            self._update_center(teacher_logits)
+            self._update_token_center(teacher_token_center_logits)
+
+        return {
+            "loss": total_loss,
+            "loss_dino": loss_dino.detach(),
+            "loss_dino_same_mode": loss_same_mode.detach(),
+            "loss_dino_cross_mode": loss_cross_mode.detach(),
+            "loss_token": loss_token.detach(),
+            "loss_token_same_mode": loss_token_same_mode.detach(),
+            "loss_token_cross_mode": loss_token_cross_mode.detach(),
+            "loss_token_unmasked": getattr(self, "_last_token_unmasked_loss", loss_token.new_zeros(())).detach(),
+            "loss_token_unmasked_same_mode": getattr(
+                self, "_last_token_unmasked_same_mode_loss", loss_token.new_zeros(())
+            ).detach(),
+            "loss_token_unmasked_cross_mode": getattr(
+                self, "_last_token_unmasked_cross_mode_loss", loss_token.new_zeros(())
+            ).detach(),
+            "loss_decoder": loss_decoder.detach(),
+            "loss_total": total_loss.detach(),
+            "loss_contrastive": total_loss.detach(),
+            "teacher_momentum": self.last_teacher_momentum.detach(),
+        }
+
+    def save_pretrained(
+        self,
+        save_directory: str,
+        *,
+        save_safetensors: bool = True,
+        state_dict: Optional[Dict[str, torch.Tensor]] = None,
+        gene_vocab: Optional[Dict[str, Any]] = None,
+        meta_vocab: Optional[Dict[str, Any]] = None,
+        manifest: Optional[Any] = None,
+        extra_files: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        os.makedirs(save_directory, exist_ok=True)
+
+        cfg_path = os.path.join(save_directory, "model_config.json")
+        _json_dump(self.cfg.__dict__, cfg_path)
+
+        if state_dict is None:
+            state_dict = self.state_dict()
+        cpu_state = {k: v.detach().cpu() for k, v in state_dict.items()}
+
+        if save_safetensors:
+            import safetensors.torch as st
+
+            st.save_file(cpu_state, os.path.join(save_directory, "model.safetensors"))
+        else:
+            torch.save(cpu_state, os.path.join(save_directory, "pytorch_model.bin"))
+
+        if gene_vocab is not None:
+            _json_dump(gene_vocab, os.path.join(save_directory, "gene_vocab.json"))
+        if meta_vocab is not None:
+            _json_dump(meta_vocab, os.path.join(save_directory, "meta_vocab.json"))
+        if manifest is not None:
+            _json_dump(manifest, os.path.join(save_directory, "manifest.json"))
+        if extra_files:
+            for name, obj in extra_files.items():
+                _json_dump(obj, os.path.join(save_directory, name))
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        load_directory: str,
+        *,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+        strict: bool = True,
+        load_assets: bool = True,
+    ) -> Tuple["CellContrastivePredictModel", Optional[Dict[str, Any]], Optional[Dict[str, Any]], Optional[Any]]:
+        cfg_path = os.path.join(load_directory, "model_config.json")
+        if not os.path.exists(cfg_path):
+            raise FileNotFoundError(f"Missing model_config.json in: {load_directory}")
+
+        with open(cfg_path, "r", encoding="utf-8") as f:
+            cfg_dict = json.load(f)
+
+        cfg_fields = {f.name for f in dataclass_fields(FoundationConfig)}
+        filtered_cfg = {k: v for k, v in cfg_dict.items() if k in cfg_fields}
+        cfg = FoundationConfig(**filtered_cfg)
+        model = cls(cfg)
+
+        st_path = os.path.join(load_directory, "model.safetensors")
+        pt_path = os.path.join(load_directory, "pytorch_model.bin")
+        if os.path.exists(st_path):
+            import safetensors.torch as st
+
+            state_dict = st.load_file(st_path, device="cpu")
+        elif os.path.exists(pt_path):
+            state_dict = torch.load(pt_path, map_location="cpu")
+        else:
+            raise FileNotFoundError(f"Missing model weights in: {load_directory}")
+
+        if "token_center" not in state_dict:
+            state_dict["token_center"] = torch.zeros_like(model.token_center)
+
+        for missing_prefix, source_prefix in (
+            ("student_ibot_head.", "student_head."),
+            ("teacher_ibot_head.", "teacher_head."),
+        ):
+            if any(key.startswith(missing_prefix) for key in state_dict):
+                continue
+            for model_key in model.state_dict().keys():
+                if not model_key.startswith(missing_prefix):
+                    continue
+                source_key = source_prefix + model_key[len(missing_prefix) :]
+                if source_key in state_dict:
+                    state_dict[model_key] = state_dict[source_key].clone()
+
+        model.load_state_dict(state_dict, strict=strict)
+        if dtype is not None:
+            model = model.to(dtype=dtype)
+        if device is not None:
+            model = model.to(device=device)
+
+        gene_vocab = None
+        meta_vocab = None
+        manifest = None
+        if load_assets:
+            gene_vocab_path = os.path.join(load_directory, "gene_vocab.json")
+            meta_vocab_path = os.path.join(load_directory, "meta_vocab.json")
+            manifest_path = os.path.join(load_directory, "manifest.json")
+            if os.path.exists(gene_vocab_path):
+                with open(gene_vocab_path, "r", encoding="utf-8") as f:
+                    gene_vocab = json.load(f)
+            if os.path.exists(meta_vocab_path):
+                with open(meta_vocab_path, "r", encoding="utf-8") as f:
+                    meta_vocab = json.load(f)
+            if os.path.exists(manifest_path):
+                with open(manifest_path, "r", encoding="utf-8") as f:
+                    manifest = json.load(f)
+        return model, gene_vocab, meta_vocab, manifest
+
+
+__all__ = ["CellContrastivePredictModel", "FoundationConfig"]

scRep_pretrain/train_utils.py

@@ -0,0 +1,299 @@
+import os
+import json
+import time
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+from transformers import Trainer
+
+from .manifest import build_manifest_from_dir
+from .vocab import build_gene_vocab_from_manifest, build_meta_vocab_from_manifest
+
+
+logger = logging.getLogger("scRep_train")
+
+
+# =============================================================================
+# Distributed helpers
+# =============================================================================
+def dist_is_initialized() -> bool:
+    return torch.distributed.is_available() and torch.distributed.is_initialized()
+
+
+def dist_rank_world() -> Tuple[int, int]:
+    if dist_is_initialized():
+        return torch.distributed.get_rank(), torch.distributed.get_world_size()
+    rank = int(os.environ.get("RANK", "0"))
+    world = int(os.environ.get("WORLD_SIZE", "1"))
+    return rank, world
+
+
+def dist_barrier() -> None:
+    if dist_is_initialized():
+        torch.distributed.barrier()
+
+
+def is_world_process_zero() -> bool:
+    rank, _ = dist_rank_world()
+    return rank == 0
+
+
+def wait_for_file(path: str, timeout: float = 300.0, poll: float = 0.5) -> None:
+    """Wait until a file exists and is readable."""
+    p = Path(path)
+    start = time.time()
+    last_err = None
+    while time.time() - start < timeout:
+        if p.exists():
+            try:
+                with open(p, "r") as f:
+                    _ = f.read(1)
+                return
+            except Exception as e:
+                last_err = e
+        time.sleep(poll)
+    raise TimeoutError(f"Timed out waiting for file: {path}. Last error: {last_err}")
+
+
+# =============================================================================
+# JSON helpers
+# =============================================================================
+def load_json(path: str) -> Any:
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def save_json(obj: Any, path: str) -> None:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    tmp_path = path + ".tmp"
+    with open(tmp_path, "w") as f:
+        json.dump(obj, f, ensure_ascii=False, indent=2)
+        f.flush()
+        os.fsync(f.fileno())
+    os.replace(tmp_path, path)
+
+
+# =============================================================================
+# FSDP / Trainer helpers
+# =============================================================================
+def parse_fsdp_list(fsdp_str: Optional[str]) -> Optional[List[str]]:
+    s = (fsdp_str or "").strip()
+    if not s:
+        return None
+    s = s.replace(",", " ")
+    parts = [p.strip() for p in s.split() if p.strip()]
+    return parts if parts else None
+
+
+def load_fsdp_config(path: Optional[str]) -> Optional[Dict[str, Any]]:
+    if not path:
+        return None
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"fsdp_config not found: {path}")
+    return load_json(path)
+
+
+def normalize_report_to(report_to: str):
+    return [] if report_to == "none" else report_to
+
+
+class TrainerBatchWrapperCollator:
+    """Wrap the collated batch as {"batch": ...} for Trainer."""
+    def __init__(self, base_collator):
+        self.base = base_collator
+
+    def __call__(self, examples):
+        return {"batch": self.base(examples)}
+
+
+# =============================================================================
+# Asset helpers
+# =============================================================================
+def gene_vocab_to_map(gene_vocab: Dict[str, Any]) -> Dict[str, int]:
+    gene_offset = int(gene_vocab.get("gene_offset", 2))
+    gene_names = list(gene_vocab["gene_names"])
+    return {str(name): i + gene_offset for i, name in enumerate(gene_names)}
+
+
+def _normalize_meta_keys(meta_keys: str) -> List[str]:
+    return [k.strip() for k in meta_keys.split(",") if k.strip()]
+
+
+def _asset_metadata(args) -> Dict[str, Any]:
+    return {
+        "data_dir": os.path.abspath(args.data_dir),
+        "manifest_json": os.path.abspath(args.manifest_json) if args.manifest_json else "",
+        "gene_vocab_json": os.path.abspath(args.gene_vocab_json) if getattr(args, "gene_vocab_json", "") else "",
+        "use_raw": bool(args.use_raw),
+        "backed": args.backed,
+        "meta_keys": _normalize_meta_keys(args.meta_keys),
+    }
+
+
+def prepare_scratch_assets(args, output_dir: str) -> Tuple[Any, Dict[str, Any], Dict[str, Any], Dict[str, int], Dict[str, Dict[str, int]], Dict[str, int]]:
+    is_main = is_world_process_zero()
+
+    manifest_path = os.path.join(output_dir, "manifest.json")
+    gene_vocab_path = os.path.join(output_dir, "gene_vocab.json")
+    meta_vocab_path = os.path.join(output_dir, "meta_vocab.json")
+    metadata_path = os.path.join(output_dir, "asset_metadata.json")
+
+    expected_metadata = _asset_metadata(args)
+    cache_complete = all(
+        os.path.exists(path)
+        for path in (manifest_path, gene_vocab_path, meta_vocab_path, metadata_path)
+    )
+    metadata_matches = False
+    cache_reason = "missing cached asset files"
+    if cache_complete:
+        cached_metadata = load_json(metadata_path)
+        metadata_matches = cached_metadata == expected_metadata
+        if not metadata_matches:
+            cache_reason = "cached asset metadata does not match current dataset/config"
+
+    if cache_complete and metadata_matches:
+        logger.info("loading cached scratch assets from %s", output_dir)
+        manifest = load_json(manifest_path)
+        gene_vocab = load_json(gene_vocab_path)
+        meta_vocab = load_json(meta_vocab_path)
+        gene_name_to_id = gene_vocab_to_map(gene_vocab)
+        meta_maps = meta_vocab["meta_maps"]
+        meta_vocab_sizes = meta_vocab["meta_vocab_sizes"]
+        logger.info(
+            "loaded cached scratch assets: manifest=%s datasets, genes=%s, meta_keys=%s",
+            len(manifest),
+            gene_vocab.get("gene_vocab_size", 0),
+            list(meta_vocab_sizes.keys()),
+        )
+        return manifest, gene_vocab, meta_vocab, gene_name_to_id, meta_maps, meta_vocab_sizes
+
+    if is_main:
+        logger.info("scratch assets cache miss in %s: %s", output_dir, cache_reason)
+        os.makedirs(output_dir, exist_ok=True)
+
+        if args.manifest_json and os.path.exists(args.manifest_json):
+            logger.info("building manifest from explicit manifest_json=%s", args.manifest_json)
+            manifest = load_json(args.manifest_json)
+            save_json(manifest, manifest_path)
+        else:
+            logger.info("building manifest from data_dir=%s", args.data_dir)
+            manifest_start = time.time()
+            manifest = build_manifest_from_dir(
+                args.data_dir,
+                backed=args.backed,
+                use_raw=bool(args.use_raw),
+            )
+            if len(manifest) == 0:
+                raise RuntimeError(
+                    "manifest is empty when building from `data_dir`. "
+                    "This usually means your h5ad files were filtered out.\n\n"
+                    "Common causes:\n"
+                    "1) `--use_raw` is enabled but your .h5ad has no `adata.raw`.\n"
+                    "2) Wrong `data_dir` / file pattern (no .h5ad found).\n"
+                    "3) Required obs/meta columns mismatch (meta_keys).\n\n"
+                    f"args.use_raw={args.use_raw}, data_dir={args.data_dir}"
+                )
+            save_json(manifest, manifest_path)
+            logger.info("manifest saved to %s in %.1fs", manifest_path, time.time() - manifest_start)
+
+        meta_keys = _normalize_meta_keys(args.meta_keys)
+        if getattr(args, "gene_vocab_json", ""):
+            logger.info("loading fixed gene vocab from %s", args.gene_vocab_json)
+            gene_vocab = load_json(args.gene_vocab_json)
+            gene_names = list(gene_vocab["gene_names"])
+            gene_name_to_id = gene_vocab_to_map(gene_vocab)
+            logger.info("fixed gene vocab loaded with %s genes", len(gene_names))
+        else:
+            logger.info("building gene vocab from %s manifest entries", len(manifest))
+            gene_start = time.time()
+            gene_name_to_id, gene_names = build_gene_vocab_from_manifest(
+                manifest,
+                use_raw=args.use_raw,
+                backed=args.backed,
+                gene_offset=2,
+                sort_paths=True,
+            )
+            logger.info("gene vocab built with %s genes in %.1fs", len(gene_names), time.time() - gene_start)
+
+        logger.info("building meta vocab for keys=%s", meta_keys)
+        meta_start = time.time()
+        meta_maps, meta_vocab_sizes = build_meta_vocab_from_manifest(
+            manifest,
+            keys=meta_keys,
+            backed=args.backed,
+            sort_paths=True,
+        )
+        logger.info("meta vocab built in %.1fs", time.time() - meta_start)
+
+        if not getattr(args, "gene_vocab_json", ""):
+            gene_vocab = {
+                "gene_offset": 2,
+                "gene_names": gene_names,
+                "gene_vocab_size": max(gene_name_to_id.values()) + 1,
+            }
+        meta_vocab = {
+            "meta_vocab_sizes": meta_vocab_sizes,
+            "meta_maps": meta_maps,
+        }
+        save_json(gene_vocab, gene_vocab_path)
+        save_json(meta_vocab, meta_vocab_path)
+        save_json(expected_metadata, metadata_path)
+        logger.info("scratch assets saved to %s", output_dir)
+    else:
+        logger.info("waiting for scratch assets to be prepared in %s", output_dir)
+        wait_for_file(metadata_path)
+        manifest = load_json(manifest_path)
+        gene_vocab = load_json(gene_vocab_path)
+        meta_vocab = load_json(meta_vocab_path)
+        gene_name_to_id = gene_vocab_to_map(gene_vocab)
+        meta_maps = meta_vocab["meta_maps"]
+        meta_vocab_sizes = meta_vocab["meta_vocab_sizes"]
+        logger.info("loaded scratch assets prepared by rank0 from %s", output_dir)
+
+    return manifest, gene_vocab, meta_vocab, gene_name_to_id, meta_maps, meta_vocab_sizes
+
+
+def save_bundle(
+    trainer: Trainer,
+    model,
+    save_dir: str,
+    *,
+    gene_vocab: Dict[str, Any],
+    meta_vocab: Dict[str, Any],
+    manifest: Any,
+) -> None:
+    """
+    Save a final bundle that is directly loadable by ``CellContrastivePredictModel.from_pretrained``.
+
+    Under FSDP we explicitly ask Accelerate for a gathered state dict, then save it
+    through the unwrapped model. All ranks participate in ``get_state_dict``; only
+    world-process-zero writes files.
+    """
+    os.makedirs(save_dir, exist_ok=True)
+
+    if hasattr(trainer, "accelerator") and trainer.accelerator is not None:
+        unwrapped_model = trainer.accelerator.unwrap_model(model)
+        state_dict = trainer.accelerator.get_state_dict(trainer.model)
+        if trainer.is_world_process_zero():
+            unwrapped_model.save_pretrained(
+                save_dir,
+                save_safetensors=bool(trainer.args.save_safetensors),
+                state_dict=state_dict,
+                gene_vocab=gene_vocab,
+                meta_vocab=meta_vocab,
+                manifest=manifest,
+                extra_files={"training_args.json": trainer.args.to_dict()},
+            )
+    else:
+        if not trainer.is_world_process_zero():
+            return
+        model.save_pretrained(
+            save_dir,
+            save_safetensors=bool(trainer.args.save_safetensors),
+            gene_vocab=gene_vocab,
+            meta_vocab=meta_vocab,
+            manifest=manifest,
+            extra_files={"training_args.json": trainer.args.to_dict()},
+        )

scRep_pretrain/vocab.py

@@ -0,0 +1,180 @@
+from dataclasses import dataclass
+import time
+import logging
+
+import anndata as ad
+from typing import Any, Dict, List, Optional, Tuple
+
+
+# =============================================================================
+# Debug / Logging
+# =============================================================================
+DEBUG_PRINT: bool = False
+logger = logging.getLogger("scRep_train")
+
+
+def _log(msg: str) -> None:
+    if DEBUG_PRINT:
+        print(f"[sc-token] {msg}")
+
+
+def _safe_close_anndata(adata_obj: Any) -> None:
+    f = getattr(adata_obj, "file", None)
+    if f is not None:
+        try:
+            f.close()
+        except Exception:
+            pass
+
+
+# =============================================================================
+# Global vocab builders
+# =============================================================================
+def build_gene_vocab_from_manifest(
+    manifest: List[Dict[str, Any]],
+    *,
+    use_raw: bool,
+    backed: Optional[str] = "r",
+    gene_offset: int = 2,
+    sort_paths: bool = True,
+    print_every: int = 50,
+) -> Tuple[Dict[str, int], List[str]]:
+    """
+    Build a global gene vocab:
+      gene_name -> token_id
+
+    Token ids start from `gene_offset` because:
+      0 = PAD
+      1 = CLS
+    """
+    items = list(manifest)
+    if sort_paths:
+        items.sort(key=lambda x: str(x.get("path", "")))
+
+    gene_name_to_id: Dict[str, int] = {}
+    _log(f"build gene vocab: files={len(items)}, use_raw={use_raw}")
+    logger.info("gene vocab scan start: files=%s use_raw=%s", len(items), use_raw)
+
+    start_time = time.time()
+    for idx, it in enumerate(items, start=1):
+        path = it["path"]
+        adata_obj = ad.read_h5ad(path, backed=backed)
+        try:
+            if use_raw:
+                if (adata_obj.raw is None) or (adata_obj.raw.var_names is None):
+                    raise ValueError(f"[{path}] use_raw=True but adata.raw.var_names missing.")
+                var_names = adata_obj.raw.var_names
+            else:
+                if adata_obj.var_names is None:
+                    raise ValueError(f"[{path}] adata.var_names missing.")
+                var_names = adata_obj.var_names
+
+            for g in map(str, list(var_names)):
+                if g not in gene_name_to_id:
+                    gene_name_to_id[g] = len(gene_name_to_id) + gene_offset
+        finally:
+            _safe_close_anndata(adata_obj)
+
+        if idx % max(1, int(print_every)) == 0:
+            logger.info(
+                "gene vocab scan progress: %s/%s unique_genes=%s elapsed=%.1fs",
+                idx, len(items), len(gene_name_to_id), time.time() - start_time,
+            )
+
+    gene_names: List[str] = [None] * len(gene_name_to_id)
+    for name, tid in gene_name_to_id.items():
+        gene_names[tid - gene_offset] = name
+
+    _log(f"gene vocab: n_genes={len(gene_names)}, vocab_size={len(gene_names) + gene_offset}")
+    logger.info(
+        "gene vocab scan done: n_genes=%s vocab_size=%s elapsed=%.1fs",
+        len(gene_names), len(gene_names) + gene_offset, time.time() - start_time,
+    )
+    return gene_name_to_id, gene_names
+
+
+def build_meta_vocab_from_manifest(
+    manifest: List[Dict[str, Any]],
+    *,
+    keys: List[str],
+    backed: Optional[str] = "r",
+    sort_paths: bool = True,
+    print_every: int = 50,
+) -> Tuple[Dict[str, Dict[str, int]], Dict[str, int]]:
+    """
+    Build global meta vocabs:
+      meta_value(str) -> id(int)
+
+    0 is reserved for unknown / missing.
+    """
+    items = list(manifest)
+    if sort_paths:
+        items.sort(key=lambda x: str(x.get("path", "")))
+
+    keys = list(dict.fromkeys(keys))
+    meta_maps: Dict[str, Dict[str, int]] = {k: {} for k in keys}
+    if len(keys) == 0:
+        logger.info("meta vocab scan skipped because keys=[]")
+        return {}, {}
+
+    _log(f"build meta vocabs: keys={keys}, files={len(items)}")
+    logger.info("meta vocab scan start: files=%s keys=%s", len(items), keys)
+
+    start_time = time.time()
+    for idx, it in enumerate(items, start=1):
+        path = it["path"]
+        adata_obj = ad.read_h5ad(path, backed=backed)
+        try:
+            for k in keys:
+                if k not in adata_obj.obs.columns:
+                    raise ValueError(f"[{path}] obs has no column '{k}'")
+                col = adata_obj.obs[k].astype("category")
+                for v in col.cat.categories.astype(str).tolist():
+                    if v not in meta_maps[k]:
+                        meta_maps[k][v] = len(meta_maps[k]) + 1
+        finally:
+            _safe_close_anndata(adata_obj)
+
+        if idx % max(1, int(print_every)) == 0:
+            logger.info(
+                "meta vocab scan progress: %s/%s elapsed=%.1fs",
+                idx, len(items), time.time() - start_time,
+            )
+
+    meta_vocab_sizes = {k: len(meta_maps[k]) + 1 for k in keys}
+    _log(f"meta vocab sizes: {meta_vocab_sizes}")
+    logger.info("meta vocab scan done: sizes=%s elapsed=%.1fs", meta_vocab_sizes, time.time() - start_time)
+    return meta_maps, meta_vocab_sizes
+
+# =============================================================================
+# Token vocab ids
+# =============================================================================
+@dataclass(frozen=True)
+class VocabSpec:
+    # gene vocab
+    gene_pad_id: int = 0
+    gene_cls_id: int = 1
+    gene_offset: int = 2  # valid gene ids: [gene_offset, gene_vocab_size-1]
+
+    # bin vocab
+    bin_pad_id: int = 0
+    bin_mask_id: int = 1
+
+    # Reserved exclusively for "unexpressed / dropout-style" states.
+    # This id is used for:
+    #   - weak zero-query targets
+    #   - zero-corruption of masked observed genes
+    bin_zero_id: int = 2
+
+    # Lowest strictly-positive expression bin id.
+    # If n_bins counts positive bins only, then:
+    #   bin vocab size = n_bins + bin_offset
+    # which now correctly becomes n_bins + 3.
+    bin_offset: int = 3
+
+
+def gene_vocab_to_map(gene_vocab: Dict[str, Any]) -> Dict[str, int]:
+    """Convert gene_vocab.json to a gene_name -> token_id mapping."""
+    gene_offset = int(gene_vocab.get("gene_offset", 2))
+    gene_names = list(gene_vocab["gene_names"])
+    return {str(name): i + gene_offset for i, name in enumerate(gene_names)}