torchcell.datasets.ProtT5Dataset

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

Bases: BaseEmbeddingDataset

Embedding dataset of ProtT5 protein representations per gene.

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

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

initialize_model() → ProtT5

Return a ProtT5 model instance for the uniref50 checkpoint.

process() → None

Compute ProtT5 embeddings for all genes and save the processed data.

Returns before building a backbone when no model is named. Each gene stores a [1, hidden] float32 CPU tensor, hidden read from the backbone’s config (zeros for an excluded ORF classification), so the collate is one [n_genes, hidden] tensor.