Correct README: 512-d shared space, fix package import + GitHub URL, fill in verified eval numbers

Fixes from end-to-end reproducer validation (zinagoodlab/spatialwhisperer PR #1):
- shared embedding dim: 2048 -> 512 (2048 was the pooled BioBERT feature dim, not the projected shared space)
- code example: from spatialwhisperer import load_spatialwhisperer_model (was cellwhisperer.utils.model_io.load_cellwhisperer_model, pre-rename)
- GitHub URL: zinagoodlab/spatialwhisperer (was Good-Lab, which is the HF org only)
- removed reference to non-existent examples/zero_shot_celltype.py
- dropped "private repo access" framing + pass api_keys/huggingface_write shell example
- evaluation table: filled in Lizard (0.764) and PanNuke (0.689) AUROCs with verified-reproduced numbers; dropped the suspect +XX% column (baseline comparisons live in the paper)
- ckpt size: ~530 MB (was ~505 MB)

README.md

@@ -1,94 +1,119 @@
 ---
 license: cc-by-nc-4.0
 tags:
-  - histopathology
-  - spatial-transcriptomics
-  - multimodal
-  - vision-language
-  - clip
-  - cell-type-annotation
+- biology
+- histopathology
+- spatial-transcriptomics
+- multimodal
+- pathology
+- gene-expression
+- biobert
+- vision-language
 library_name: pytorch
-pipeline_tag: zero-shot-image-classification
 ---
 
 # SpatialWhisperer
 
-SpatialWhisperer is a trimodal embedding model that aligns hematoxylin & eosin (H&E) image patches, gene-expression profiles, and free-text descriptions into a shared 2048-dimensional space. It enables zero-shot cell-type annotation of H&E patches and natural-language querying over histopathology and spatial transcriptomics data.
+A trimodal embedding model that aligns hematoxylin & eosin (H&E) image patches, gene-expression profiles, and free-text descriptions into a shared 512-dimensional space. Enables zero-shot cell-type annotation of H&E patches and natural-language querying over histopathology and spatial transcriptomics data.
 
-This repository hosts the main checkpoint (seed=0) from the ICML 2026 paper *[Transitive Representation Learning Enhances Histopathology Annotation](https://openreview.net/forum?id=Ze7U293Zw4)* (Schaefer et al., PMLR vol. 306).
+This checkpoint (seed 0) is from the ICML 2026 paper *[Transitive Representation Learning Enhances Histopathology Annotation](https://openreview.net/forum?id=Ze7U293Zw4)* (Schaefer et al., PMLR vol. 306). The paper refers to this configuration as the **Trimodal** model: three encoders trained on two paired-modality datasets (transcriptome↔text and image↔transcriptome) that together span three modalities.
 
-## Model architecture
+- **Code & reproduction pipeline:** <https://github.com/zinagoodlab/spatialwhisperer>
+- **Paper:** <https://openreview.net/forum?id=Ze7U293Zw4>
 
-Three encoders project into a shared embedding space:
+## Architecture
 
-| Modality       | Encoder          | Freezing |
-|----------------|------------------|----------|
-| Image (H&E)    | UNI2             | locked   |
-| Transcriptome  | Geneformer (12L) | locked   |
-| Text           | BioBERT v1.1     | unfrozen |
+| Modality | Encoder | Status |
+|----------|---------|--------|
+| Image (H&E) | UNI2 (`MahmoodLab/UNI2-h`) | locked |
+| Transcriptome | Geneformer-12L-30M | locked |
+| Text | BioBERT v1.1 | trained |
 
-Following LiT convention, the freezing pattern is **LUL** (image locked, text unlocked, transcriptome locked). Only the text tower and the three projection heads are trained. Projection dimension is 2048.
+Three projection heads map each encoder's pooled features into a shared 512-dimensional space. Only the text tower and the three projection heads are trained.
 
 ## Training data
 
-Three paired datasets cover the three modality pairs:
+Three paired-modality datasets:
 
 - **HEST-1K** — H&E ↔ spatial gene expression (Visium-style spots)
-- **cellxgene_census** — gene expression ↔ free-text cell/sample metadata
+- **CellxGene Census** — gene expression ↔ free-text cell/sample metadata
 - **ARCHS4/GEO** — gene expression ↔ free-text sample descriptions
 
-Training was 4 epochs with AdamW at learning rate 1e-5 and cosine schedule (warmup 3%), batch size 512, on a single H100 GPU. This checkpoint reflects epoch 3, global step 14624.
+Training: 4 epochs, AdamW (lr 1e-5), cosine schedule (3% warmup), batch size 512, single H100. This checkpoint is from epoch 3, global step 14624.
 
-## Evaluation
+## What this checkpoint contains
+
+- `spatialwhisperer.ckpt` — Lightning state-dict (~530 MB, 236 tensors: trained BioBERT text tower + three projection heads) plus the `hyper_parameters` block. Optimizer/scheduler state is stripped.
 
-Reported AUROC on cell-type benchmarks (mean across cell types):
+**The locked foundation-model weights are NOT included.** UNI2 and Geneformer are re-instantiated at load time from their original providers. The `load_spatialwhisperer_model()` convenience wrapper fetches both on first call.
 
-| Benchmark | SpatialWhisperer | Best published baseline | Δ rel.  |
-|-----------|------------------|-------------------------|---------|
-| PathoCell | **0.630**        | 0.554                   | +13.7%  |
-| Lizard    | (see paper)      | —                       | +15.9%  |
-| PanNuke   | (see paper)      | —                       | +13.7%  |
+## Usage
 
-Modality-pair benchmarks (Tabula Sapiens, HEST-1K, Skin Conditions) confirm the trimodal model retains per-pair performance under low-n subsampling. See the paper for full numbers.
+Install the [code repository](https://github.com/zinagoodlab/spatialwhisperer) (pixi env), then:
 
-## How to use
+```python
+from spatialwhisperer import load_spatialwhisperer_model
+
+model, tokenizer, transcriptome_processor, image_processor = load_spatialwhisperer_model()
+# model: TranscriptomeTextDualEncoderLightning (frozen, eval mode, on CUDA if available)
+```
 
-The checkpoint is a stripped Lightning state-dict (~505 MB, 236 tensors covering the trained BioBERT text tower and the three 2048-d projection heads) plus its `hyper_parameters` block. **Foundation model weights are NOT included** — the locked UNI2 image encoder and locked Geneformer transcriptome encoder are re-instantiated at load time from their original providers (and remain under their respective licenses). The ckpt's `hyper_parameters.model_config.use_cache = True` flag triggers the `FrozenCachedModel` wrapping that excludes the locked towers from `state_dict` during load.
+First call downloads the SpatialWhisperer checkpoint plus UNI2 and Geneformer weights; subsequent calls load from cache.
 
-Loading requires the cellwhisperer code at <https://github.com/Good-Lab/spatialwhisperer> (model code) and the foundation models (UNI2, Geneformer, BioBERT v1.1), which are downloaded by the cellwhisperer setup scripts.
+Each `get_<modality>_features` call returns `(pooled_features, projected_embed_in_shared_space)`. The second element is the 512-D shared-space embedding to compare across modalities.
 
 ```python
-from cellwhisperer.utils.model_io import load_cellwhisperer_model
+import torch
 
-model, tokenizer, transcriptome_proc, image_proc = load_cellwhisperer_model(
-    model_path="hf://Good-Lab/spatialwhisperer"
-)
-# model is a TranscriptomeTextDualEncoderLightning in eval mode
+prompts = ["cytotoxic T cells", "plasma cells"]
+text_inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
+
+with torch.no_grad():
+    _, text_emb = model.model.get_text_features(normalize_embeds=True, **text_inputs)
+print(text_emb.shape)  # (2, 512)
 ```
 
-While the repo is private, export a token first:
+Image and transcriptome embeddings follow the same pattern; see the [GitHub README](https://github.com/zinagoodlab/spatialwhisperer#use-the-model) for complete examples covering all three modalities.
 
-```bash
-export HUGGINGFACE_TOKEN=$(pass api_keys/huggingface_write)  # or any read token with access
-```
+## Foundation-model setup
+
+UNI2 (`MahmoodLab/UNI2-h`) is a gated HuggingFace model. Before first use:
+
+1. Accept the terms at <https://huggingface.co/MahmoodLab/UNI2-h>.
+2. Make a read token visible to your environment — the loader checks `HF_TOKEN` / `HUGGINGFACE_TOKEN`, otherwise falls back to `huggingface_hub`'s on-disk cache. The simplest setup is:
+   ```bash
+   huggingface-cli login    # paste your read token
+   ```
 
-To compute image–text similarities for zero-shot cell-type annotation, encode patches and class-name strings, then take cosine similarity in the shared 2048-d space. See `examples/zero_shot_celltype.py` in the model code repository.
+Geneformer downloads without gating.
+
+## Evaluation
+
+The accompanying [code repository](https://github.com/zinagoodlab/spatialwhisperer) reproduces every paper benchmark with one command: `pixi run snakemake -j N paper_all`. Verified end-to-end on a fresh ILC ampere4 clone (2026-05-31):
+
+| Benchmark | Macro AUROC (this ckpt, seed 0) |
+|-----------|-----------------|
+| PathoCell CRC (13-class) | 0.630 |
+| Lizard (3-class reduced) | 0.764 |
+| PanNuke (4-class reduced) | 0.689 |
+| Kriegsmann Skin Conditions (16-class, clinical labels) | 0.698 |
+
+These match the paper's reported Trimodal seed-0 numbers exactly. Comparisons to CONCH, PLIP, OmiCLIP, and other baselines are in the paper.
 
 ## Intended use & limitations
 
-**Intended.** Research on multimodal histopathology, cell-type annotation, spatial transcriptomics analysis, and natural-language querying over H&E and gene-expression data.
+**Intended:** research on multimodal histopathology, zero-shot cell-type annotation, spatial transcriptomics analysis, and natural-language querying over H&E and gene-expression data.
 
-**Not intended.** Clinical diagnosis or treatment decisions. The model was trained on academic datasets and is not validated for clinical use.
+**Not intended:** clinical diagnosis or treatment decisions. The model was trained on academic datasets and is not validated for clinical use.
 
-**Known limitations.**
+**Limitations:**
 - Trained on Visium-scale spots (~55 μm); finer-grained image–expression alignment is not guaranteed.
 - BioBERT vocabulary constrains the text tower; rare technical terms may be out-of-distribution.
 - The image tower (UNI2) is locked; performance on tissue types poorly represented in UNI2's pretraining will be lower.
 
-## File contents
+## License
 
-- `spatialwhisperer.ckpt` — Lightning checkpoint (state_dict + hyper_parameters; optimizer/scheduler state stripped).
-- `README.md` — this card.
+**CC BY-NC 4.0** (research use). Foundation-model weights (UNI2, Geneformer, BioBERT) carry their own licenses; consult the upstream repositories.
 
 ## Citation
 
@@ -106,7 +131,3 @@ To compute image–text similarities for zero-shot cell-type annotation, encode
   url       = {https://openreview.net/forum?id=Ze7U293Zw4}
 }
 ```
-
-## License
-
-CC BY-NC 4.0 (research use). Foundation model weights (UNI2, Geneformer, BioBERT) carry their own licenses; please consult upstream repositories.