Update README.md

README.md

@@ -12,22 +12,55 @@ pipeline_tag: other
 
 # AssayFormer
 
+
+[Website](https://genentech.github.io/AssayLoop) | [Models and Data](https://huggingface.co/collections/Genentech/assaybench) | [Paper](README_DETAILED.md) | [![Benchmark Package](https://img.shields.io/pypi/v/assaybench)](https://pypi.org/project/assaybench/)
+
+![The task: a screen is a library of genes, a phenotype and a hit set. Each round, a method
+sees the phenotype and everything it has already assayed, and chooses the next hundred
+genes.](docs/assets/figures/figure1.png)
+
+
 The amortized gene ranker from **AssayLoop**: *Biology-in-the-loop: Amortized
 Adaptive Hit Discovery in CRISPR Screens*. Given a screen description and the
-labels revealed so far, it scores every gene in a 24,217-gene vocabulary and
-proposes the next batch to assay.
+hit labels revealed so far, it scores every gene in a 24,217-gene vocabulary and
+proposes the next batch to assay. In the paper, a 21,147 gene subset was used. 
+
+
+## Usage
+
+```python
+from huggingface_hub import snapshot_download
+from assayloop.models.amortized_ranker import AmortizedRankerModel
+
+ckpt = snapshot_download("Genentech/assayformer")
+model = AmortizedRankerModel(checkpoint=ckpt, ckpt_file="model_last.pt")
+```
+
+Note the explicit `ckpt_file`: the loader defaults to `model.pt`, which is not
+in this repository.
+
+Then run it on a held-out screen and score the trajectory:
+
+```python
+from assaybench import SequentialLoop, enrichment_factor
+from assayloop.acquisitions.greedy_from_model import GreedyFromModel
+from assayloop.tasks import AssayBenchGeneBatchTask, gene_universe, load_screens
+
+screens = load_screens(target_set="public")          # the paper's 20-screen test set
+universe = gene_universe(screens)                    # the f2 pool, 21,147 genes
 
-This is the `+ GRPO` row of the paper's main table — the transformer with BPMF
-gene embeddings, supervised on 1,349 screens and then fine-tuned with GRPO.
+screen = screens[0]
+task = AssayBenchGeneBatchTask(screen, universe_genes=universe)
+run = SequentialLoop(task, model, GreedyFromModel(), metrics=[], batch_size=100).run(n_steps=10)
+
+picked = [g for step in run.history for g in step.acquired_batch]
+hits = [g for g, h in zip(screen.genes, screen.hits) if h]
+print(enrichment_factor(picked, screen.genes, hits, universe=universe, budget=1000)) # 7.61 on U_1733_merged
+```
 
-- **Code:** https://github.com/genentech/assayloop
-- **Benchmark data:** [Genentech/assaybench](https://huggingface.co/datasets/Genentech/assaybench)
-- **Benchmark package:** [`assaybench`](https://pypi.org/project/assaybench/) on PyPI
 
 ## Results
 
-Twenty held-out genome-wide screens, ten rounds of one hundred genes, against a
-21,147-gene universe (`--min-screen-freq 2`).
 
 | metric | value |
 |---|---|
@@ -38,8 +71,6 @@ Twenty held-out genome-wide screens, ten rounds of one hundred genes, against a
 | DepMap common-essential fraction | 0.407 |
 | Effective pathways EP-B / EP-S / EP-D | 19.3 / 46.7 / 65.0 |
 
-EP is the effective number of Reactome level-2 groups, rarefied to 30 / 200 /
-6000 annotated genes; a uniform draw scores 21.8 / 56.4 / 82.1.
 
 ## Files
 
@@ -49,10 +80,7 @@ EP is the effective number of Reactome level-2 groups, rarefied to 30 / 200 /
 | `config.json` | architecture and RL hyperparameters |
 | `vocab.json` | the 24,217-entry gene vocabulary (index 0 is `<unk>`) |
 
-`model_last.pt` is the checkpoint the paper reports, not the best-on-validation
-one. Held-out EF dips and recovers during GRPO, so the run also kept a
-best-on-val snapshot; every number above, every figure, and the handoff
-experiments use the final epoch. If you want the gene embedding table on its
+`model_last.pt` is the checkpoint the paper reports. If you want the gene embedding table on its
 own it is `state_dict["gene_emb.weight"]`, shape (24217, 10).
 
 ## Architecture
@@ -69,85 +97,13 @@ gene embedding.
 | context | up to 1024 observations |
 
 **The description input is a precomputed embedding, not text.** The model
-expects a 1536-d vector from OpenAI `text-embedding-3-small`; there is no text
-encoder in the checkpoint. The `bert_model` field in `config.json` is inert for
-this `encoder_type` — it applies only to the `modernbert_text` variant. The
+expects a 1536-d vector from OpenAI `text-embedding-3-small`. The
 repository ships the screen-description embeddings for the benchmark screens, so
 you only need an embedding backend for screens of your own.
 
-## Usage
-
-```python
-from huggingface_hub import snapshot_download
-from assayloop.models.amortized_ranker import AmortizedRankerModel
-
-ckpt = snapshot_download("Genentech/assayformer")
-model = AmortizedRankerModel(checkpoint=ckpt, ckpt_file="model_last.pt")
-```
-
-Note the explicit `ckpt_file`: the loader defaults to `model.pt`, which is not
-in this repository.
-
-Then run it on a held-out screen and score the trajectory:
-
-```python
-from assaybench import SequentialLoop, enrichment_factor
-from assayloop.acquisitions.greedy_from_model import GreedyFromModel
-from assayloop.tasks.gene_batch import AssayBenchGeneBatchTask
-from assayloop.tasks.screen_sets import load_screens
-
-screen = load_screens(target_set="public")[0]        # the paper's 20-screen test set
-task = AssayBenchGeneBatchTask(screen)
-run = SequentialLoop(task, model, GreedyFromModel(), metrics=[], batch_size=100).run(n_steps=10)
-
-picked = [g for step in run.history for g in step.acquired_batch]
-hits = [g for g, h in zip(screen.genes, screen.hits) if h]
-print(enrichment_factor(picked, screen.genes, hits, budget=1000))   # 6.96 on U_1733_merged
-```
-
-No API key is needed: the screen-description embeddings for the benchmark
-screens ship with the repository, so `task_context` resolves offline.
-
-That prints one number for one screen, which is not the 4.83 in the table
-above. Two differences: the table is the mean over all twenty screens, and it
-scores against the full-genome candidate pool — the union of the twenty
-libraries, restricted to genes measured in at least two of them — which the
-task takes as `universe_genes=`. Picking from a 21,147-gene pool rather than
-one screen's own library is harder, and it is the number the paper reports.
-`python -m assayloop.scripts.full_genome_table --min-screen-freq 2` reproduces
-the table row for row.
-
-## Training
-
-Supervised on 1,349 screens (hits objective, leave-one-out, 40 epochs), then
-GRPO for 100 epochs on the `public_train` split with `public_validation` held
-out: group size 8, 10 steps of 100 genes, lr 1e-5, entropy coefficient 0.01, no
-KL penalty, context reset every 25 epochs, seed 19. The reward is the
-context-delta form; `config.json` has the rest.
-
-## Gene embeddings and BPMF
-
-The gene embedding table was initialised from the posterior mean of a K=10
-Bayesian PMF fit and then updated by supervised training and GRPO. Every row
-moved, but not far: median cosine to the BPMF initialisation is 0.9999 and the
-median relative displacement is 1.6%. Treat the table as approximately the BPMF
-gene factors rather than as independently learned representations. The BPMF fit
-itself — the screen loadings, the posterior samples, the predicted
-probabilities — is not part of this release.
-
-## Limitations
-
-- Trained on BioGRID ORCS screens; the split is by year, and the vocabulary is
-  fixed at training time. Genes outside it map to `<unk>`.
-- EF is computed against each screen's own library. A gene the screen never
-  assayed can never be a hit, and 10% of this model's picks fall there.
-- 40.7% of its picks are DepMap common-essential — a real fraction of the hit
-  rate comes from genes that are essential in most cell lines.
 
 ## Citation
 
-The arXiv identifier is not assigned yet; this entry will be updated when it is.
-
 ```bibtex
 @article{assayloop2026,
   title   = {Biology-in-the-loop: Amortized Adaptive Hit Discovery in CRISPR Screens},