Metadata-Version: 2.4
Name: alphagenome_pt
Version: 0.3.0
Summary: An implementation of AlphaGenome in PyTorch.
Keywords: alphagenome,pytorch,deep-learning,genomics,dna,transformer
Author: Rylie Weaver, Gomathi Lakshmanan, John Lagergren
Author-email: Rylie Weaver <rylieweaver9@gmail.com>, Gomathi Lakshmanan <gomathi.l@hotmail.com>, John Lagergren <lagergrenjh@ornl.gov>
License-Expression: Apache-2.0
License-File: LICENSE
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Requires-Dist: torch>=2.0
Requires-Dist: numpy>=1.22
Requires-Dist: einops>=0.6
Requires-Dist: einx>=0.3
Requires-Dist: huggingface-hub>=0.20
Requires-Dist: tomli>=2.0 ; python_full_version < '3.11'
Requires-Dist: pytest>=7 ; extra == 'dev'
Requires-Dist: alphagenome ; extra == 'jax2pt'
Requires-Dist: jax ; extra == 'jax2pt'
Requires-Dist: dm-haiku ; extra == 'jax2pt'
Requires-Python: >=3.10
Project-URL: Homepage, https://github.com/RylieWeaver/AlphaGenome_PyTorch
Project-URL: Issues, https://github.com/RylieWeaver/AlphaGenome_PyTorch/issues
Provides-Extra: dev
Provides-Extra: jax2pt
Description-Content-Type: text/markdown

# AlphaGenome_PyTorch
<p align="left">
  <img src="images/AG_PyTorch.png" width="800"/>
</p>
<!-- <p align="left">
  <img src="images/AG_Total.png" width="800"/>
</p> -->
<!-- <p align="left">
  <img src="images/AG_Architecture.png" width="400"/>
</p> -->
*Image generated by Nano Banana 2*


## Quick Start

### From-Scratch
```python
from alphagenome_pt import deepmind_model

model = deepmind_model()
```

### Finetuned
Load the converted DeepMind state into a model. By default this downloads from
Hugging Face cache if the state is not already local:

```python
from alphagenome_pt import deepmind_model

model = deepmind_model(load_state=True)
```

### Get Predictions/Embeddings
```python
from alphagenome_pt import synthetic_batch
batch = synthetic_batch(model.metadata, batch_size=1, seq_len=model.max_seq_len)
predictions, embeddings = model(batch)

print(
  embeddings.embeddings_1bp.shape,
  embeddings.embeddings_128bp.shape,
  embeddings.embeddings_pair.shape,
)
print(predictions.keys())
```

For finetuning with custom downstream heads, build custom metadata:

```python
from alphagenome_pt import (
    HeadName,
    deepmind_model,
    synthetic_batch,
    synthetic_metadata,
)

metadata = synthetic_metadata(
    heads=(HeadName.RNA_SEQ, HeadName.ATAC, HeadName.SPLICE_SITES_CLASSIFICATION),
    num_organisms=2,
    num_tracks=8,
)
# NOTE: can set heads=False to skip loading heads and/or organisms=False to skip loading organisms.
# They are prefix-loaded by default.
model = deepmind_model(
    metadata=metadata.metadata,
    load_state=True,
)
batch = synthetic_batch(metadata, batch_size=1, seq_len=8192)
loss, scalars, predictions = model.loss(batch)
loss.backward()
```


### Metadata
Metadata defines organisms, output heads, track counts, masks, and track means.
For examples and smoke tests, generate synthetic metadata:

```python
from alphagenome_pt import HeadName, synthetic_metadata

metadata = synthetic_metadata(
    heads=(HeadName.RNA_SEQ, HeadName.CONTACT_MAPS),
    num_organisms=2,
    num_tracks=8,
)
```

To use the converted official DeepMind metadata:

```python
from alphagenome_pt import deepmind_metadata

metadata_dict = deepmind_metadata()
```

To download all converted metadata files and all checkpoint folds ahead of time:

```bash
alphagenome-pt-download --local-dir checkpoints
```




## Docs
The `docs` directory contains information on environment setup, explanations of the data structure and model architecture, and example trainings. To train on your own data, it's strongly recommended to read the `model.md` and `data.md` markdown files in the `*/AlphaGenome_PyTorch/docs/guides` directory so that you understand the structure of data batches and metadata.


### Environment
See `*/AlphaGenome_PyTorch/docs/environment` for instructions on how to set up a UV environment to run AlphaGenome_PyTorch.

### Official JAX Checkpoint Conversion
Most users should download the converted PyTorch checkpoint files from Hugging
Face:

```bash
alphagenome-pt-download --local-dir checkpoints
```

The package also includes optional JAX-to-PyTorch conversion tooling for
recreating the PyTorch checkpoint from the official public JAX weights:

```bash
pip install "alphagenome-pt[jax2pt]"
python -m alphagenome_pt.jax2pt.convert_state
```

By default this converts every supported official checkpoint. To convert only one checkpoint, pass `--fold all_folds`, `--fold fold_0`, `--fold fold_1`, `--fold fold_2`, or `--fold fold_3`; outputs default to `alphagenome_{fold_name}.pt`.

The converter also requires the public `alphagenome_research` checkout used by the official loader:

```bash
git clone https://github.com/google-deepmind/alphagenome_research.git
pip install -e alphagenome_research
```

### Guides
See `*/AlphaGenome_PyTorch/docs/guides` for explanations on the AlphaGenome model and its data structure (very helpful for understanding examples).

### Examples
See `*/AlphaGenome_PyTorch/docs/examples` for examples of:
- Masked Language Modeling (MLM) training (`train_mlm.py`)
- Training on Downstream Tasks (RNA-Seq, CAGE, ATAC, Splice Sites Classification/Usage/Junction) (`train_downstream.py`)
- MLM Pretraining --> Training on Downstream Tasks (`train_downstream_from_pretrained.py`)


## Acknowledgements
This repository is a reimplementation of the AlphaGenome model in PyTorch, with an added option for Masked Language Modeling (MLM).

Within the `alphagenome_pt` directory, some components are direct ports of the released AlphaGenome code [Link1](https://github.com/google-deepmind/alphagenome) [Link2](https://github.com/google-deepmind/alphagenome_research) (licensed under [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0)), some are reimplementations based on pseudocode from the [BioArXiV paper](https://www.biorxiv.org/content/10.1101/2025.06.25.661532v1), and others are original additions (e.g., the MLM head). Attribution is made clear at the top of each file in the `alphagenome_pt` directory. The `docs` and `tests` directories are original work (with LLM coding assistance).

Developing this project used resources of the Oak Ridge Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC05-00OR22725.


## Intended Audience
This intended audience of this repo is for those who might want to train the AlphaGenome architecture in PyTorch with maximal flexibility to hyperparameters and new organisms/contexts. You should be able to prepare a batch of tensor data and set up a train/val/test loop, but don't want the hassle of replicating the architecture.


## Other Implementations
There is another AlphaGenome PyTorch implementation out [here](https://pypi.org/project/alphagenome-pytorch/0.2.8/#description) by Phillip Wang (a.k.a. LucidRains) which is quite good. The GitHub page is down as of March 2nd, 2026, but the PyPi package remains. The main advantages of that implementation (as of version 0.2.8) are in evaluation (loading the published weights and running variant scoring). The main advantage of this implementation is research training (an MLM head and track masks that can vary by batch in training). This implementation also has a `.loss()` function in the model to compute multi-resolution losses for you, and one head per task with dense weight tensor of shape [O, D, T] rather than separate weights tensors of shape [D, T] for each [organism x task], which is mathematically equivalent but more in-line with the original AlphaGenome implementation.


## Reaching Out
Want a new feature or find a bug? Feel free to leave an issue on the GitHub repository.
