updated README with logo and quickstart

README.md

@@ -10,13 +10,20 @@ tags:
 - tahoebio
 - pytorch
 ---
-# Tahoe-X1
+# Tahoe-x1
 
-Tahoe-X1 is a family of foundation models for single-cell RNA sequencing data developed by Tahoe Therapeutics. For more details, see our blog post and preprint [TODO: INSERT LINKS]. In this repository, we provide pretrained weights for three model sizes:
+Tahoe-x1 is a family of perturbation-trained single-cell foundation models  developed by Tahoe Therapeutics.
+For more details, see our blog post and [📄 preprint](http://www.tahoebio.ai/news/tahoe-x1).
+In this repository, we provide pretrained weights for three model sizes:
 
 - ~70M parameters (TahoeX1-70M)
 - ~1B parameters (TahoeX1-1B)
 - ~3B parameters (TahoeX1-3B)
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="./assets/abstract_logo_dark_mode.png">
+  <source media="(prefers-color-scheme: light)" srcset="./assets/abstract_logo_light_mode.png">
+  <img src="./assets/abstract_logo_light_mode.png" alt="Abstract Logo">
+</picture>
 
 ## Installation
 
@@ -30,57 +37,26 @@ To use the model, you must first install the `tahoe-x1` package from GitHub.
 
 2. Follow the installation steps for the docker or uv based insallation as described [here](https://github.com/tahoebio/tahoe-x1#installation) in the repository
 
-Additional files (including vocabulary files and data for the included benchmarks) are hosted on S3 in a publicly accesible bucket (s3://tahoe-hackathon-data). These files will be automatically downloaded as needed.
+Additional files (including vocabulary files and data for the included benchmarks) are hosted on S3 in a publicly accesible bucket (s3://tahoe-hackathon-data/MFM). These files will be automatically downloaded as needed.
 
-## Getting cell and gene embeddings
-
-1. Create a configuration file (see `scripts/inference/configs/predict.yaml`). 
-```bash
-# Configure:
-#   - paths.hf_repo_id: Hugging Face repository (tahoebio/TahoeX1)
-#   - paths.hf_model_size: model size (70M, 1B, or 3B)
-#   - paths.adata_output: where to save AnnData output including embeddings
-#   - predict.return_gene_embeddings: True (for extracting and saving gene embeddings)
-```
-
-You can also use a locally-downloaded model by specifying `paths.model_dir` instead of `paths.hf*`.
-
-2. Run the embedding script:
-```bash
-python scripts/inference/predict_embeddings.py path/to/config.yaml
-
-# Optional parameters can be added to overwrite config values:
-python scripts/inference/predict_embeddings.py path/to/config.yaml --paths.model_name=tx --batch_size=128
-```
-
-Cell embeddings are written to the output AnnData file. If `return_gene_embeddings` is `True`, gene embeddings are also saved. If you encounter an out-of-memory error, fall back to the lower-level `tasks.get_batch_embeddings` API, which loads gene embeddings in chunks:
+## Quickstart
+You can quickly load the model in this way:
 
 ```python
-model, vocab, model_cfg, collator_cfg = load_model(
-    model_dir,
-    device=device,
-    return_gene_embeddings=return_gene_embeddings,
-)
+from tahoex.model import ComposerTX
 
-cell_embs, gene_embs = get_batch_embeddings(
-    adata=adata,
-    model=model,
-    vocab=vocab,
-    model_cfg=model_cfg,
-    collator_cfg=collator_cfg,
-    return_gene_embeddings=True
+model, vocab, model_config, collator_config = ComposerTX.from_hf(
+    repo_id="tahoebio/Tahoe-x1",
+    model_size="70m",  # or "1b", "3b".
+    return_gene_embeddings=True,  # optional, default True
+    use_chem_inf=False  # optional, default False
 )
 ```
 
-## Training
-1. Define a YAML configuration file (see examples in `runai/` or `mcli/`) describing model, data and training parameters.
 
-2. Launch training with Composer:
-```bash
-composer scripts/train.py path/to/config.yaml
-```
 
 
-## License
 
+## License
 We release the model weights and associated code under the Apache 2.0 license.
+