mrna-bench 1.0.0__tar.gz

mrna_bench-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Ian Shi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mrna_bench-1.0.0/PKG-INFO

@@ -0,0 +1,160 @@
+Metadata-Version: 2.2
+Name: mrna-bench
+Version: 1.0.0
+Summary: Benchmarking suite for mRNA property prediction.
+Author-email: "Ruian (Ian) Shi" <ian.shi@mail.utoronto.ca>
+License: MIT License
+        
+        Copyright (c) 2024 Ian Shi
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Repository, https://github.com/IanShi1996/mRNABench
+Keywords: mrna,genomic foundation model,benchmark
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: Unix
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tqdm
+Requires-Dist: PyYAML
+Requires-Dist: scikit-learn
+Requires-Dist: numpy<2
+Requires-Dist: pandas<=2.2.2
+Requires-Dist: requests
+Provides-Extra: base-models
+Requires-Dist: torch==2.2.2; extra == "base-models"
+Requires-Dist: transformers==4.29.1; extra == "base-models"
+Requires-Dist: mamba-ssm==1.2.0.post1; extra == "base-models"
+Requires-Dist: rna-fm==0.2.2; extra == "base-models"
+Requires-Dist: accelerate; extra == "base-models"
+Provides-Extra: dev
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: flake8; extra == "dev"
+Requires-Dist: flake8-docstrings; extra == "dev"
+
+# mRNABench
+This repository contains a workflow to benchmark the embedding quality of genomic foundation models on (m)RNA specific tasks. The mRNABench contains a catalogue of datasets and training split logic which can be used to evaluate the embedding quality of several catalogued models.
+
+**Jump to:** [Model Catalog](#model-catalog) [Dataset Catalog](#dataset-catalog)
+
+## Setup
+Several configurations of the mRNABench are available.
+
+### Datasets Only
+If you are interested in the benchmark datasets **only**, you can run:
+
+```bash
+pip install mrna-bench
+```
+
+### Full Version
+The inference-capable version of mRNABench that can generate embeddings using
+Orthrus, DNA-BERT2, NucleotideTransformer, RNA-FM, and HyenaDNA can be 
+installed as shown below. Note that this requires PyTorch version 2.2.2 with 
+CUDA 12.1 and Triton uninstalled (due to a DNA-BERT2 issue).
+```bash
+conda create --name mrna_bench python=3.10
+conda activate mrna_bench
+
+pip install torch==2.2.2 --index-url https://download.pytorch.org/whl/cu121
+pip install mrna-bench[base_models] 
+pip uninstall triton
+```
+Inference with other models will require the installation of the model's
+dependencies first, which are usually listed on the model's GitHub page (see below).
+
+### Post-install
+After installation, please run the following in Python to set where data associated with the benchmarks will be stored.
+```python
+import mrna_bench as mb
+
+path_to_dir_to_store_data = "DESIRED_PATH"
+mb.update_data_path(path_to_dir_to_store_data)
+```
+
+## Usage
+Datasets can be retrieved using:
+
+```python
+import mrna_bench as mb
+
+dataset = mb.load_dataset("go-mf")
+data_df = dataset.data_df
+```
+
+The mRNABench can also be used to test out common genomic foundation models:
+```python
+import torch
+
+import mrna_bench as mb
+from mrna_bench.embedder import DatasetEmbedder
+from mrna_bench.linear_probe import LinearProbe
+
+device = torch.device("cuda")
+
+dataset = mb.load_dataset("go-mf")
+model = mb.load_model("Orthrus", "orthrus-large-6-track", device)
+
+embedder = DatasetEmbedder(model, dataset)
+embeddings = embedder.embed_dataset()
+embeddings = embeddings.detach().cpu().numpy()
+
+prober = LinearProbe(
+    dataset=dataset,
+    embeddings=embeddings,
+    task="multilabel",
+    target_col="target",
+    split_type="homology"
+)
+
+metrics = prober.run_linear_probe()
+print(metrics)
+```
+Also see the `scripts/` folder for example scripts that uses slurm to embed dataset chunks in parallel for reduce runtime, as well as an example of multi-seed linear probing.
+
+## Model Catalog
+The current models catalogued are:
+
+| Model Name |  Model Versions         | Description   | Citation | Supported<br>by<br>`base_models` |
+| :--------: |  ---------------------- | ------------- | :------: | :------------------------: |
+| `Orthrus` | `orthrus_large_6_track`<br> `orthrus_base_4_track` | Mamba-based RNA FM pre-trained using contrastive learning. | [code](https://github.com/bowang-lab/Orthrus) [paper](https://www.biorxiv.org/content/10.1101/2024.10.10.617658v2)| ✅ |
+| `AIDO.RNA` | `aido_rna_650m` <br> `aido_rna_650m_cds` <br> `aido_rna_1b600m` <br> `aido_rna_1b600m_cds` | Encoder Transformer-based RNA FM pre-trained using MLM on 42M ncRNA sequences. Version that is domain adapted to CDS is available. | [paper](https://www.biorxiv.org/content/10.1101/2024.11.28.625345v1) | |
+| `RNA-FM` | `rna-fm` <br> `mrna-fm` | Transformer-based RNA FM pre-trained using MLM on 23M ncRNA sequences. mRNA-FM trained on CDS using codon tokenizer. | [github](https://github.com/ml4bio/RNA-FM) | ✅ |
+| `DNABERT2` | `dnabert2` | Transformer-based DNA FM pre-trained using MLM on multispecies genomic dataset. Uses BPE and other modern architectural improvements for efficiency. | [github](https://github.com/MAGICS-LAB/DNABERT_2) | ✅ |
+| `NucleotideTransformer` | `2.5b-multi-species` <br> `2.5b-1000g` <br> `500m-human-ref` <br> `500m-1000g` <br> `v2-50m-multi-species` <br> `v2-100m-multi-species` <br> `v2-250m-multi-species` <br> `v2-500m-multi-species` | Transformer-based DNA FM pre-trained using MLM on a variety of possible datasets at various model sizes. Sequence is tokenized using 6-mers. | [github](https://github.com/instadeepai/nucleotide-transformer) | ✅ |
+| `HyenaDNA` | `hyenadna-large-1m-seqlen-hf` <br> `hyenadna-medium-450k-seqlen-hf` <br> `hyenadna-medium-160k-seqlen-hf` <br> `hyenadna-small-32k-seqlen-hf` <br> `hyenadna-tiny-16k-seqlen-d128-hf` | Hyena-based DNA FM pre-trained using NTP on the human reference genome. Available at various model sizes and pretraining sequence contexts. | [github](https://github.com/HazyResearch/hyena-dna) | ✅ |
+
+### Adding a new model
+All models should inherit from the template `EmbeddingModel`. Each model file should lazily load dependencies within its `__init__` methods so each model can be used individually without install all other models. Models must implement `get_model_short_name(model_version)` which fetches the internal name for the model. This must be unique for every model version and must not contain underscores. Models should implement either `embed_sequence` or `embed_sequence_sixtrack` (see code for method signature). New models should be added to `MODEL_CATALOG`.
+
+## Dataset Catalog
+The current datasets catalogued are:
+| Dataset Name | Catalogue Identifier | Description | Tasks | Citation |
+|---|---|---|---|---|
+| GO Molecular Function | <code>go-mf</code> | Classification of the molecular function of a transcript's  product as defined by the GO Resource. | `multilabel` | [website](https://geneontology.org/) |
+| Mean Ribosome Load (Sugimoto) | <code>mrl&#8209;sugimoto</code> | Mean Ribosome Load per transcript isoform as measured in Sugimoto et al. 2022. | `regression` | [paper](https://www.nature.com/articles/s41594-022-00819-2) |
+| RNA Half-life (Human) | <code>rnahl&#8209;human</code> | RNA half-life of human transcripts collected by Agarwal et al. 2022. | `regression` | [paper](https://genomebiology.biomedcentral.com/articles/10.1186/s13059-022-02811-x) |
+| RNA Half-life (Mouse) | <code>rnahl&#8209;mouse</code> | RNA half-life of mouse transcripts collected by Agarwal et al. 2022. | `regression` | [paper](https://genomebiology.biomedcentral.com/articles/10.1186/s13059-022-02811-x) |
+| Protein Subcellular Localization | <code>prot&#8209;loc</code> | Subcellular localization of transcript protein product defined in Protein Atlas. | `multilabel` | [website](https://www.proteinatlas.org/) |
+| Protein Coding Gene Essentiality | <code>pcg&#8209;ess</code> | Essentiality of PCGs as measured by CRISPR knockdown. Log-fold expression and binary essentiality available on several cell lines. | `regression` `classification`| [paper](https://www.cell.com/cell/fulltext/S0092-8674(24)01203-0)|
+
+### Adding a new dataset
+New datasets should inherit from `BenchmarkDataset`. Dataset names cannot contain underscores. Each new dataset should download raw data and process it into a dataframe by overriding `process_raw_data`. This dataframe should store transcript as rows, using string encoding in the `sequence` column. If homology splitting is required, a column `gene` containing gene names is required. Six track embedding also requires columns `cds` and `splice`. The target column can have any name, as it is specified at time of probing. New datasets should be added to `DATASET_CATALOG`.

mrna_bench-1.0.0/README.md

@@ -0,0 +1,108 @@
+# mRNABench
+This repository contains a workflow to benchmark the embedding quality of genomic foundation models on (m)RNA specific tasks. The mRNABench contains a catalogue of datasets and training split logic which can be used to evaluate the embedding quality of several catalogued models.
+
+**Jump to:** [Model Catalog](#model-catalog) [Dataset Catalog](#dataset-catalog)
+
+## Setup
+Several configurations of the mRNABench are available.
+
+### Datasets Only
+If you are interested in the benchmark datasets **only**, you can run:
+
+```bash
+pip install mrna-bench
+```
+
+### Full Version
+The inference-capable version of mRNABench that can generate embeddings using
+Orthrus, DNA-BERT2, NucleotideTransformer, RNA-FM, and HyenaDNA can be 
+installed as shown below. Note that this requires PyTorch version 2.2.2 with 
+CUDA 12.1 and Triton uninstalled (due to a DNA-BERT2 issue).
+```bash
+conda create --name mrna_bench python=3.10
+conda activate mrna_bench
+
+pip install torch==2.2.2 --index-url https://download.pytorch.org/whl/cu121
+pip install mrna-bench[base_models] 
+pip uninstall triton
+```
+Inference with other models will require the installation of the model's
+dependencies first, which are usually listed on the model's GitHub page (see below).
+
+### Post-install
+After installation, please run the following in Python to set where data associated with the benchmarks will be stored.
+```python
+import mrna_bench as mb
+
+path_to_dir_to_store_data = "DESIRED_PATH"
+mb.update_data_path(path_to_dir_to_store_data)
+```
+
+## Usage
+Datasets can be retrieved using:
+
+```python
+import mrna_bench as mb
+
+dataset = mb.load_dataset("go-mf")
+data_df = dataset.data_df
+```
+
+The mRNABench can also be used to test out common genomic foundation models:
+```python
+import torch
+
+import mrna_bench as mb
+from mrna_bench.embedder import DatasetEmbedder
+from mrna_bench.linear_probe import LinearProbe
+
+device = torch.device("cuda")
+
+dataset = mb.load_dataset("go-mf")
+model = mb.load_model("Orthrus", "orthrus-large-6-track", device)
+
+embedder = DatasetEmbedder(model, dataset)
+embeddings = embedder.embed_dataset()
+embeddings = embeddings.detach().cpu().numpy()
+
+prober = LinearProbe(
+    dataset=dataset,
+    embeddings=embeddings,
+    task="multilabel",
+    target_col="target",
+    split_type="homology"
+)
+
+metrics = prober.run_linear_probe()
+print(metrics)
+```
+Also see the `scripts/` folder for example scripts that uses slurm to embed dataset chunks in parallel for reduce runtime, as well as an example of multi-seed linear probing.
+
+## Model Catalog
+The current models catalogued are:
+
+| Model Name |  Model Versions         | Description   | Citation | Supported<br>by<br>`base_models` |
+| :--------: |  ---------------------- | ------------- | :------: | :------------------------: |
+| `Orthrus` | `orthrus_large_6_track`<br> `orthrus_base_4_track` | Mamba-based RNA FM pre-trained using contrastive learning. | [code](https://github.com/bowang-lab/Orthrus) [paper](https://www.biorxiv.org/content/10.1101/2024.10.10.617658v2)| ✅ |
+| `AIDO.RNA` | `aido_rna_650m` <br> `aido_rna_650m_cds` <br> `aido_rna_1b600m` <br> `aido_rna_1b600m_cds` | Encoder Transformer-based RNA FM pre-trained using MLM on 42M ncRNA sequences. Version that is domain adapted to CDS is available. | [paper](https://www.biorxiv.org/content/10.1101/2024.11.28.625345v1) | |
+| `RNA-FM` | `rna-fm` <br> `mrna-fm` | Transformer-based RNA FM pre-trained using MLM on 23M ncRNA sequences. mRNA-FM trained on CDS using codon tokenizer. | [github](https://github.com/ml4bio/RNA-FM) | ✅ |
+| `DNABERT2` | `dnabert2` | Transformer-based DNA FM pre-trained using MLM on multispecies genomic dataset. Uses BPE and other modern architectural improvements for efficiency. | [github](https://github.com/MAGICS-LAB/DNABERT_2) | ✅ |
+| `NucleotideTransformer` | `2.5b-multi-species` <br> `2.5b-1000g` <br> `500m-human-ref` <br> `500m-1000g` <br> `v2-50m-multi-species` <br> `v2-100m-multi-species` <br> `v2-250m-multi-species` <br> `v2-500m-multi-species` | Transformer-based DNA FM pre-trained using MLM on a variety of possible datasets at various model sizes. Sequence is tokenized using 6-mers. | [github](https://github.com/instadeepai/nucleotide-transformer) | ✅ |
+| `HyenaDNA` | `hyenadna-large-1m-seqlen-hf` <br> `hyenadna-medium-450k-seqlen-hf` <br> `hyenadna-medium-160k-seqlen-hf` <br> `hyenadna-small-32k-seqlen-hf` <br> `hyenadna-tiny-16k-seqlen-d128-hf` | Hyena-based DNA FM pre-trained using NTP on the human reference genome. Available at various model sizes and pretraining sequence contexts. | [github](https://github.com/HazyResearch/hyena-dna) | ✅ |
+
+### Adding a new model
+All models should inherit from the template `EmbeddingModel`. Each model file should lazily load dependencies within its `__init__` methods so each model can be used individually without install all other models. Models must implement `get_model_short_name(model_version)` which fetches the internal name for the model. This must be unique for every model version and must not contain underscores. Models should implement either `embed_sequence` or `embed_sequence_sixtrack` (see code for method signature). New models should be added to `MODEL_CATALOG`.
+
+## Dataset Catalog
+The current datasets catalogued are:
+| Dataset Name | Catalogue Identifier | Description | Tasks | Citation |
+|---|---|---|---|---|
+| GO Molecular Function | <code>go-mf</code> | Classification of the molecular function of a transcript's  product as defined by the GO Resource. | `multilabel` | [website](https://geneontology.org/) |
+| Mean Ribosome Load (Sugimoto) | <code>mrl&#8209;sugimoto</code> | Mean Ribosome Load per transcript isoform as measured in Sugimoto et al. 2022. | `regression` | [paper](https://www.nature.com/articles/s41594-022-00819-2) |
+| RNA Half-life (Human) | <code>rnahl&#8209;human</code> | RNA half-life of human transcripts collected by Agarwal et al. 2022. | `regression` | [paper](https://genomebiology.biomedcentral.com/articles/10.1186/s13059-022-02811-x) |
+| RNA Half-life (Mouse) | <code>rnahl&#8209;mouse</code> | RNA half-life of mouse transcripts collected by Agarwal et al. 2022. | `regression` | [paper](https://genomebiology.biomedcentral.com/articles/10.1186/s13059-022-02811-x) |
+| Protein Subcellular Localization | <code>prot&#8209;loc</code> | Subcellular localization of transcript protein product defined in Protein Atlas. | `multilabel` | [website](https://www.proteinatlas.org/) |
+| Protein Coding Gene Essentiality | <code>pcg&#8209;ess</code> | Essentiality of PCGs as measured by CRISPR knockdown. Log-fold expression and binary essentiality available on several cell lines. | `regression` `classification`| [paper](https://www.cell.com/cell/fulltext/S0092-8674(24)01203-0)|
+
+### Adding a new dataset
+New datasets should inherit from `BenchmarkDataset`. Dataset names cannot contain underscores. Each new dataset should download raw data and process it into a dataframe by overriding `process_raw_data`. This dataframe should store transcript as rows, using string encoding in the `sequence` column. If homology splitting is required, a column `gene` containing gene names is required. Six track embedding also requires columns `cds` and `splice`. The target column can have any name, as it is specified at time of probing. New datasets should be added to `DATASET_CATALOG`.

mrna_bench-1.0.0/mrna_bench/__init__.py

@@ -0,0 +1,9 @@
+from mrna_bench.loader.loader import load_model, load_dataset
+from mrna_bench.utils import update_data_path, get_data_path
+
+__all__ = [
+    "load_model",
+    "load_dataset",
+    "update_data_path",
+    "get_data_path"
+]

mrna_bench-1.0.0/mrna_bench/data_splitter/data_splitter.py

@@ -0,0 +1,53 @@
+from abc import ABC, abstractmethod
+
+import pandas as pd
+
+
+class DataSplitter(ABC):
+    """Generates reproducible train test splits."""
+
+    @abstractmethod
+    def split_df(
+        self,
+        df: pd.DataFrame,
+        test_size: float,
+        random_seed: int
+    ) -> tuple[pd.DataFrame, pd.DataFrame]:
+        """Split dataframe into train and test dataframes.
+
+        Args:
+            df: Dataframe to split.
+            test_size: Fraction of dataset to assign to test split.
+            random_seed: Random seed used during split sampling.
+
+        Returns:
+            Train and test dataframes.
+        """
+        pass
+
+    def get_all_splits_df(
+        self,
+        df: pd.DataFrame,
+        split_ratios: tuple[float, float, float],
+        random_seed: int,
+    ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+        """Get train, validation, test splits.
+
+        Args:
+            df: Dataframe to split.
+            split_ratios: Ratio of training, validation, test split size.
+            random_seed: Random seed used to generate splits.
+
+        Returns:
+            Dataframe containing training, validation, test splits.
+        """
+        if sum(split_ratios) != 1:
+            raise ValueError("Split ratios must sum to 1.")
+
+        tv_split_size = split_ratios[1] + split_ratios[2]
+        test_split_size = split_ratios[2] / tv_split_size
+
+        train_df, tv_df = self.split_df(df, tv_split_size, random_seed)
+        val_df, test_df = self.split_df(tv_df, test_split_size, random_seed)
+
+        return train_df, val_df, test_df

mrna_bench-1.0.0/mrna_bench/data_splitter/homology_split.py

@@ -0,0 +1,200 @@
+from pathlib import Path
+import shutil
+import zipfile
+
+import pandas as pd
+import numpy as np
+
+from mrna_bench import get_data_path
+from mrna_bench.data_splitter.data_splitter import DataSplitter
+from mrna_bench.utils import download_file
+
+
+def train_test_split_homologous(
+    genes: list[str],
+    homology_df: pd.DataFrame,
+    test_size: float = 0.3,
+    random_state: int | None = None
+) -> tuple[list[int], list[int]]:
+    """Split genes into two sets with homologous genes in the same set.
+
+    Args:
+        genes: List of gene names.
+        homology_df: DataFrame with columns "gene_name" and "gene_group".
+        test_size: Fraction of data allocated to test split.
+        random_state: Defaults to None.
+
+    Returns:
+        Dictionary with keys "train_indices" and "test_indices" containing the
+        indices of the genes in the train and test sets respectively.
+    """
+    # Map genes to their respective homology groups
+    homo_group_map = homology_df.set_index("gene_name")["gene_group"].to_dict()
+
+    gene_groups = np.array([homo_group_map.get(gene, None) for gene in genes])
+
+    group_to_index: dict[int, list[int]] = {}
+
+    # Populate the group_to_index dictionary
+    for i, group in enumerate(gene_groups):
+        if group is not None:
+            group_to_index.setdefault(group, []).append(i)
+
+    gene_index = np.arange(len(genes))
+
+    np.random.seed(random_state)
+    np.random.shuffle(gene_index)
+
+    len_of_train = int(len(gene_index) * (1 - test_size))
+
+    train_indices: list[int] = []
+    test_indices: list[int] = []
+
+    seen_groups = set()
+
+    for index in gene_index:
+        group = gene_groups[index]
+        if group is not None and group in seen_groups:
+            continue
+
+        seen_groups.add(group)
+
+        if len(train_indices) < len_of_train:
+            train_indices.extend(group_to_index.get(group, [index]))
+        else:
+            test_indices.extend(group_to_index.get(group, [index]))
+
+    train_indices = [int(ind) for ind in train_indices]
+    test_indices = [int(ind) for ind in test_indices]
+
+    return train_indices, test_indices
+
+
+class HomologySplitter(DataSplitter):
+    """Homology-based data splitter.
+
+    Uses an external homology mapping file to construct train / test splits.
+    Genes which are homologous are kept within the same 'side' of the data
+    split to reduce data leakage.
+    """
+
+    HOMO_URL = (
+        "https://zenodo.org/records/13910050/files/"
+        "homology_maps_homologene.zip"
+    )
+
+    def __init__(
+        self,
+        species: list[str],
+        homology_map_path: str | None = None,
+        force_redownload: bool = False
+    ):
+        """Initialize HomologySplitter.
+
+        Homology splitting requires a dataframe mapping genes to a group
+        of homologous genes. This dataframe must have columns "gene_name" and
+        "gene_group". If homology_map_path is not specified, relationships
+        processed from homologene will be downloaded.
+
+        Homology map files should be named as: {species}_homology_map.csv
+
+        Args:
+            species: List of species that genes are from.
+            homology_map_path: Path to homology maps.
+            force_redownload: Forces redownload of homology maps.
+        """
+        super().__init__()
+
+        if homology_map_path is None:
+            data_storage_path = get_data_path()
+            self.homology_map_path = data_storage_path + "/homology_maps"
+
+            if not Path(self.homology_map_path).exists():
+                Path(self.homology_map_path).mkdir(exist_ok=True)
+                force_redownload = True
+        else:
+            self.homology_map_path = homology_map_path
+
+        if force_redownload:
+            out = download_file(self.HOMO_URL, self.homology_map_path, True)[0]
+
+            with zipfile.ZipFile(out, "r") as zip_ref:
+                for file in zip_ref.namelist():
+                    if not file.endswith("/"):
+                        source = zip_ref.open(file)
+                        path = Path(self.homology_map_path) / Path(file).name
+
+                        with open(path, "wb") as target:
+                            shutil.copyfileobj(source, target)
+
+            Path(out).unlink()
+
+        self.species = species
+        self.homology_df = self.get_homology_df()
+
+    def get_homology_df(self) -> pd.DataFrame:
+        """Get homology dataframe for specified species.
+
+        If multiple species exists, naively merges dataframes, merging homology
+        groups if gene name overlap exists.
+
+        Returns:
+            Dataframe with gene names mapped to homology groups.
+        """
+        species_maps = []
+        for s in self.species:
+            homology_dir = Path(self.homology_map_path)
+            homology_map = homology_dir / "{}_homology_map.csv".format(s)
+            if not homology_map.exists():
+                print("Homology map for species {} is missing.".format(s))
+            else:
+                species_maps.append(pd.read_csv(homology_map))
+
+        max_groups = [max(s_df["gene_group"]) for s_df in species_maps]
+        cusum_max_groups = np.cumsum(max_groups)
+
+        for i, s_df in enumerate(species_maps[1:]):
+            s_df["gene_group"] += cusum_max_groups[i]
+
+        out_df = pd.concat(species_maps, axis=0, ignore_index=True)
+
+        # Merge identical gene names
+        out_df.groupby("gene_name", as_index=False).min()
+        out_df.reset_index(drop=True, inplace=True)
+
+        return out_df
+
+    def split_df(
+        self,
+        df: pd.DataFrame,
+        test_size: float,
+        random_seed: int
+    ) -> tuple[pd.DataFrame, pd.DataFrame]:
+        """Split dataframe rows into train and test df using homology split.
+
+        Args:
+            df: Dataframe to split.
+            test_size: Fraction of dataset to assign to test split.
+            random_seed: Random seed used for sampling rows during splitting.
+
+        Returns:
+            Dataframe containing train data and test data.
+        """
+        if "gene" not in df.columns:
+            raise ValueError("Gene must be column for homology split.")
+
+        genes = df["gene"]
+
+        train_indices, test_indices = train_test_split_homologous(
+            genes,
+            self.homology_df,
+            test_size,
+            random_seed
+        )
+
+        assert len(set(train_indices).intersection(set(test_indices))) == 0
+
+        train_df = df.iloc[train_indices]
+        test_df = df.iloc[test_indices]
+
+        return train_df, test_df

mrna_bench-1.0.0/mrna_bench/data_splitter/sklearn_split.py

@@ -0,0 +1,31 @@
+from sklearn.model_selection import train_test_split
+
+import pandas as pd
+
+from mrna_bench.data_splitter.data_splitter import DataSplitter
+
+
+class SklearnSplitter(DataSplitter):
+    """Wrapper for sklearn train_test_split."""
+
+    def split_df(
+        self,
+        df: pd.DataFrame,
+        test_size: float,
+        random_seed: int
+    ) -> tuple[pd.DataFrame, pd.DataFrame]:
+        """Split dataframe into train and test split using sklearn.
+
+        Args:
+            df: Dataframe to split.
+            test_size: Fraction of dataset to place in test set.
+            random_seed: Random seed used during sampling.
+
+        Return:
+            Train and test dataframes.
+        """
+        return train_test_split(
+            df,
+            test_size=test_size,
+            random_state=random_seed
+        )

mrna_bench-1.0.0/mrna_bench/data_splitter/split_catalog.py

@@ -0,0 +1,11 @@
+from collections.abc import Callable
+
+from mrna_bench.data_splitter.data_splitter import DataSplitter
+from mrna_bench.data_splitter.homology_split import HomologySplitter
+from mrna_bench.data_splitter.sklearn_split import SklearnSplitter
+
+
+SPLIT_CATALOG: dict[str, Callable[..., DataSplitter]] = {
+    "default": SklearnSplitter,
+    "homology": HomologySplitter
+}

mrna_bench-1.0.0/mrna_bench/datasets/__init__.py

@@ -0,0 +1,7 @@
+from mrna_bench.datasets.benchmark_dataset import BenchmarkDataset
+from mrna_bench.datasets.dataset_catalog import DATASET_CATALOG
+
+__all__ = [
+    "BenchmarkDataset",
+    "DATASET_CATALOG"
+]