torchcell.datasets.Esm2Dataset

class Esm2Dataset(root: str, genome: SCerevisiaeGenome, model_name: str | None = None, transform: Callable[[...], Any] | None = None, pre_transform: Callable[[...], Any] | None = None)

Bases: BaseEmbeddingDataset

Embedding dataset that computes ESM2 protein embeddings per gene.

static parse_genome(genome: SCerevisiaeGenome | None) → ParsedGenome | None

Return a ParsedGenome holding the genome’s gene set, or None.

initialize_model() → Esm2

Return the ESM2 model corresponding to the selected model name.

process() → None

Compute ESM2 embeddings for each gene and save the processed dataset.

Each gene stores a [1, hidden] float32 row (zeros for an excluded ORF classification), so the collate is [n_genes, hidden].