diff --git a/docs/docs/usage/wandb_sweeps.md b/docs/docs/usage/wandb_sweeps.md new file mode 100644 index 0000000..20ea951 --- /dev/null +++ b/docs/docs/usage/wandb_sweeps.md @@ -0,0 +1,32 @@ +# How to use W&B Sweeps with ModelGenerator for hyperparameter tuning + +## Caveats +W&B agents cannot launch multi-node training jobs, which causes great difficulties integrating W&B Sweeps with ModelGenerator. This guide is based on a hacky workaround that introduces many limitations. + +### The workaround +An agent is configured to exit immediately after retrieving the next set of hyperparamenters and outputing the complete training command to stdout. This command is then executed on each node without being monitored by an active agent. + +### Limitations +1. All agent functionalities are lost. It is not possible to use agent to start/stop/resume/update training runs. Users must manually terminate training runs or implement early-stopping mechanisms. +2. Failed runs have to be re-run manually using your own sbatch scripts. The command for that run is availale in stdout of the failed run. +3. Parameter importance plots use wrong parameters by default, it can be manually fixed by selecting the right parameter names in your mgen config. + +>**NOTE**: Before proceeding, please make sure that your training job uses **WandbLogger**. +## SLURM +### Step 1: create a wandb sweep +The default `slurm_sweep.yaml` creates a wandb sweep with the training command `mgen fit --config .local/test.yaml` under the project `autotune-test`. Please modify it to suit your experiments. Key values to change are **project**, **command** and **parameters**. + +Run the following command to create a wandb sweep: +```bash +wandb sweep scripts/wandb_sweep/slurm_sweep.yaml +``` +Take a note of your sweep ID for step 2. It looks like `//` and is found in the output: `wandb: Run sweep agent with: wandb agent` +### Step 2: submit the next training job to SLURM +Similar to step 1, you need to edit `slurm_agent.sh` for your experiment. The most important changes are **WANDB_PROJECT** and **SWEEP_ID**. + +The following command creates one sweep agent that runs training with the next set of hyperparamenters. +```bash +sbatch scripts/wandb_sweep/slurm_agent.sh +``` + +>**TIPS**: To queue your other sweep runs, use `sbatch --dependency`. To launch your other sweep runs in parallel, use `sbatch --array=1-X` where `X` is the number of parallel runs. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index c47b873..9f26a7f 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -16,6 +16,7 @@ nav: - usage/exporting_models.md - usage/reproducing_experiments.md - usage/embedding_caching.md + - usage/wandb_sweeps.md - Tutorials: - tutorials/kfold_cross_validation.md - tutorials/finetuning_scheduler.md diff --git a/experiments/AIDO.Cell/extract_features.py b/experiments/AIDO.Cell/extract_features.py index 42ba571..a1f316f 100644 --- a/experiments/AIDO.Cell/extract_features.py +++ b/experiments/AIDO.Cell/extract_features.py @@ -14,6 +14,8 @@ model = model.to(device).to(torch.float16) adata = ad.read_h5ad('../../modelgenerator/cell-downstream-tasks/zheng/zheng_train.h5ad') +if not adata.obs_names.is_unique: + adata.obs_names_make_unique() batch_np = adata[:batch_size].X.toarray() batch_tensor = torch.from_numpy(batch_np).to(torch.float16).to(device) diff --git a/experiments/AIDO.Cell/tutorial_target_id.ipynb b/experiments/AIDO.Cell/tutorial_target_id.ipynb index 6159b3b..48b52f9 100644 --- a/experiments/AIDO.Cell/tutorial_target_id.ipynb +++ b/experiments/AIDO.Cell/tutorial_target_id.ipynb @@ -216,6 +216,7 @@ "sc.pp.log1p(adata)\n", "\n", "# Clustering + UMAP\n", + "sc.pp.pca(adata)\n", "sc.pp.neighbors(adata)\n", "sc.tl.leiden(adata, flavor='igraph', n_iterations=2, resolution=0.5)\n", "sc.tl.umap(adata)\n", diff --git a/huggingface/aido.cell/embed.py b/huggingface/aido.cell/embed.py index 2033da2..c6b4d52 100755 --- a/huggingface/aido.cell/embed.py +++ b/huggingface/aido.cell/embed.py @@ -49,6 +49,10 @@ print("Loading input data...") try: adata = ad.read_h5ad(INPUT_FILE) + # Check for duplicate obs_names + if not adata.obs_names.is_unique: + print("⚠ Duplicate AnnData obs_names detected. Automatically applying obs_names_make_unique().") + adata.obs_names_make_unique() print(f"✓ Loaded data with {adata.n_obs} cells and {adata.n_vars} genes\n") except Exception as e: print(f"Error loading data: {e}") diff --git a/huggingface/aido.cell/inference.py b/huggingface/aido.cell/inference.py index 6c9166c..c4bcbd7 100755 --- a/huggingface/aido.cell/inference.py +++ b/huggingface/aido.cell/inference.py @@ -55,6 +55,10 @@ print("Loading input data...") try: adata = ad.read_h5ad(INPUT_FILE) + # Check for duplicate obs_names + if not adata.obs_names.is_unique: + print("⚠ Duplicate AnnData obs_names detected. Automatically applying obs_names_make_unique().") + adata.obs_names_make_unique() print(f"✓ Loaded data with {adata.n_obs} cells and {adata.n_vars} genes\n") except Exception as e: print(f"Error loading data: {e}") diff --git a/modelgenerator/backbones/backbones.py b/modelgenerator/backbones/backbones.py index 0c91534..928923d 100644 --- a/modelgenerator/backbones/backbones.py +++ b/modelgenerator/backbones/backbones.py @@ -146,27 +146,6 @@ def setup(self): for name, param in self.encoder.named_parameters(): param.requires_grad = False - def process_batch(self, batch, device, add_special_tokens=True, **kwargs): - """Processes a batch of sequences to model input format. - - Args: - batch (List[str]): List of input sequences. - device (torch.device): Device to move the data to. - add_special_tokens (bool, optional): Whether to add special tokens. Defaults to True. - - Returns: - Dict: A dictionary containing required args for forward pass. - """ - seq_tokenized = self.tokenize( - batch["sequences"], padding=True, add_special_tokens=add_special_tokens, **kwargs - ) - for k, v in seq_tokenized.items(): - if v is not None: - if torch.is_tensor(v): - seq_tokenized[k] = v.to(dtype=torch.long, device=device) - else: - seq_tokenized[k] = torch.tensor(v, dtype=torch.long, device=device) - return seq_tokenized def forward( self, @@ -1555,28 +1534,6 @@ def get_decoder(self) -> nn.Module: """ return _Identity() - def process_batch( - self, batch: dict, device: torch.device, add_special_tokens: bool = True, **kwargs - ): - """Processes a batch of sequences to model input format. - - Args: - batch (dict): A dictionary containing input sequences. - device (torch.device): Device to move the data to. - - Returns: - Dict: A dictionary containing required args for forward pass. - """ - seq_tokenized = self.tokenize( - batch["sequences"], padding=True, add_special_tokens=add_special_tokens, **kwargs - ) - for k, v in seq_tokenized.items(): - if v is not None: - if torch.is_tensor(v): - seq_tokenized[k] = v.to(dtype=torch.long, device=device) - else: - seq_tokenized[k] = torch.tensor(v, dtype=torch.long, device=device) - return seq_tokenized def tokenize( self, @@ -1821,28 +1778,6 @@ def get_decoder(self) -> nn.Module: """ return _Identity() - def process_batch( - self, batch: dict, device: torch.device, add_special_tokens: bool = True, **kwargs - ): - """Processes a batch of sequences to model input format. - - Args: - batch (dict): A dictionary containing input sequences. - device (torch.device): Device to move the data to. - - Returns: - Dict: A dictionary containing required args for forward pass. - """ - seq_tokenized = self.tokenize( - batch["sequences"], padding=True, add_special_tokens=add_special_tokens, **kwargs - ) - for k, v in seq_tokenized.items(): - if v is not None: - if torch.is_tensor(v): - seq_tokenized[k] = v.to(dtype=torch.long, device=device) - else: - seq_tokenized[k] = torch.tensor(v, dtype=torch.long, device=device) - return seq_tokenized def tokenize(self, sequences: list[str], **kwargs) -> dict: """Tokenizes a list of sequences @@ -2033,17 +1968,6 @@ def get_decoder(self) -> nn.Module: """ return self.decoder - def process_batch(self, batch: dict, device: torch.device, **kwargs): - """Processes a batch of sequences to model input format. - - Args: - batch (dict): A dictionary containing input sequences. - - Returns: - Dict: A dictionary containing required args for forward pass. - """ - input_ids = self.tokenize(batch["sequences"])["input_ids"].to(device=device) - return {"input_ids": input_ids} def tokenize( self, @@ -2230,17 +2154,6 @@ def get_decoder(self) -> nn.Module: """ return self.decoder - def process_batch(self, batch: dict, device: torch.device, **kwargs): - """Processes a batch of sequences to model input format. - - Args: - batch (dict): A dictionary containing input sequences. - - Returns: - Dict: A dictionary containing required args for forward pass. - """ - input_ids = self.tokenize(batch["sequences"])["input_ids"].to(device=device) - return {"input_ids": input_ids} def tokenize( self, @@ -2494,24 +2407,6 @@ def get_decoder(self) -> nn.Module: """ return self.decoder - def process_batch(self, batch: dict, device: torch.device, **kwargs): - """Processes a batch of sequences to model input format. - - Args: - batch (dict): A dictionary containing input sequences. - device (torch.device): Device to move the data to. - - Returns: - Dict: A dictionary containing required args for forward pass. - """ - seq_tokenized = self.tokenize(batch["sequences"]) - for k, v in seq_tokenized.items(): - if v is not None: - if torch.is_tensor(v): - seq_tokenized[k] = v.to(dtype=torch.long, device=device) - else: - seq_tokenized[k] = torch.tensor(v, dtype=torch.long, device=device) - return seq_tokenized def tokenize(self, sequences: list[str], **kwargs) -> dict: """Tokenizes a list of sequences diff --git a/modelgenerator/backbones/base.py b/modelgenerator/backbones/base.py index a7e30e0..6628375 100644 --- a/modelgenerator/backbones/base.py +++ b/modelgenerator/backbones/base.py @@ -309,7 +309,16 @@ def process_batch( Returns: Dict: A dictionary containing required args for forward pass. """ - raise NotImplementedError + seq_tokenized = self.tokenize( + batch["sequences"], padding=True, add_special_tokens=add_special_tokens, **kwargs + ) + for k, v in seq_tokenized.items(): + if v is not None: + if torch.is_tensor(v): + seq_tokenized[k] = v.to(dtype=torch.long, device=device) + else: + seq_tokenized[k] = torch.tensor(v, dtype=torch.long, device=device) + return seq_tokenized def required_data_columns(self) -> List[str]: """List of required data columns for the model. diff --git a/modelgenerator/cell/utils.py b/modelgenerator/cell/utils.py index 8f1ccca..771fffa 100644 --- a/modelgenerator/cell/utils.py +++ b/modelgenerator/cell/utils.py @@ -4,8 +4,29 @@ import bionty as bt import numpy as np import pandas as pd +import logging from lightning.pytorch.utilities import rank_zero_info +logger = logging.getLogger(__name__) + + +def _ensure_unique_obs_names(adata: ad.AnnData) -> ad.AnnData: + """Ensures that the observation names in the AnnData object are unique. + + Args: + adata (ad.AnnData): The input AnnData object. + + Returns: + ad.AnnData: The AnnData object with unique observation names. + """ + if not adata.obs_names.is_unique: + logger.warning( + "Duplicate AnnData obs_names detected. " + "Automatically applying obs_names_make_unique()." + ) + adata.obs_names_make_unique() + return adata + def build_map(gene_symbols): # Get map of symbols to Ensembl IDs: diff --git a/modelgenerator/data/data.py b/modelgenerator/data/data.py index 2000831..eba47b9 100644 --- a/modelgenerator/data/data.py +++ b/modelgenerator/data/data.py @@ -1556,17 +1556,21 @@ def provided_columns(self) -> List[str]: if self.filter_columns is not None: return ["sequences"] + self.filter_columns adata = ad.read_h5ad(os.path.join(self.path, self.trainfile), backed="r") + adata = cell_utils._ensure_unique_obs_names(adata) return ["sequences"] + list(adata.obs.columns) def setup(self, stage: Optional[str] = None): """Set up the data module by loading the whole datasets and splitting them into training, validation, and test sets.""" adata_train = ad.read_h5ad(os.path.join(self.path, self.trainfile)) + adata_train = cell_utils._ensure_unique_obs_names(adata_train) adata_train = cell_utils.map_gene_symbols(adata_train, symbol_field="gene_symbols") adata_train = cell_utils.align_genes(adata_train, self.backbone_gene_list) adata_val = ad.read_h5ad(os.path.join(self.path, self.valfile)) + adata_val = cell_utils._ensure_unique_obs_names(adata_val) adata_val = cell_utils.map_gene_symbols(adata_val, symbol_field="gene_symbols") adata_val = cell_utils.align_genes(adata_val, self.backbone_gene_list) adata_test = ad.read_h5ad(os.path.join(self.path, self.testfile)) + adata_test = cell_utils._ensure_unique_obs_names(adata_test) adata_test = cell_utils.map_gene_symbols(adata_test, symbol_field="gene_symbols") adata_test = cell_utils.align_genes(adata_test, self.backbone_gene_list) @@ -1840,6 +1844,7 @@ def provided_columns(self) -> List[str]: def setup(self, stage: Optional[str] = None): """Set up the data module by loading the whole datasets and splitting them into training, validation, and test sets.""" adata = ad.read_h5ad(os.path.join(self.path, self.trainfile)) + adata = cell_utils._ensure_unique_obs_names(adata) adata = cell_utils.align_genes(adata, self.backbone_gene_list, ensembl_field="feature_id") adata_train = adata[adata.obs[self.split_column] == "train"] @@ -1913,6 +1918,7 @@ def __init__( # Load metadata in backed mode self.adata = ad.read_h5ad(self.file_path, backed="r") + self.adata = cell_utils._ensure_unique_obs_names(self.adata) self._process_metadata() self.neighbor_indices = self._precompute_neighbors() self.length = len(self.adata.obs) @@ -2176,13 +2182,17 @@ def provided_columns(self) -> List[str]: if self.filter_columns is not None: return ["sequences"] + self.filter_columns adata = ad.read_h5ad(os.path.join(self.path, self.file), backed="r") + adata = cell_utils._ensure_unique_obs_names(adata) return ["sequences"] + list(adata.obs.columns) - def setup(self, stage: Optional[str] = None): """Set up the data module by loading the whole datasets and splitting them into training, validation, and test sets.""" rank_zero_info("***") rank_zero_info(f"loading {self.file}") adata = ad.read_h5ad(os.path.join(self.path, self.file)) + # Check for duplicate obs_names + if not adata.obs_names.is_unique: + rank_zero_info("⚠ Duplicate AnnData obs_names detected. Automatically applying obs_names_make_unique().") + adata.obs_names_make_unique() adata = cell_utils.map_gene_symbols(adata, symbol_field="index") adata = cell_utils.align_genes(adata, self.backbone_gene_list) rank_zero_info(f"loaded {adata.shape[0]} cells") diff --git a/modelgenerator/main.py b/modelgenerator/main.py index 10fabc4..b73d6b9 100644 --- a/modelgenerator/main.py +++ b/modelgenerator/main.py @@ -9,18 +9,25 @@ class MyLightningCLI(LightningCLI): def add_arguments_to_parser(self, parser): - parser.link_arguments("data.init_args.batch_size", "model.init_args.batch_size") + parser.link_arguments( + "data.init_args.batch_size", + "model.init_args.batch_size", + apply_on="instantiate", + ) parser.link_arguments( "model.init_args.backbone.class_path", "trainer.strategy.init_args.auto_wrap_policy.init_args.backbone_classes", + apply_on="instantiate", ) parser.link_arguments( "model.init_args.backbone.class_path", "data.init_args.backbone_class_path", + apply_on="instantiate", ) parser.link_arguments( "model.init_args.backbone.init_args.use_peft", "trainer.strategy.init_args.auto_wrap_policy.init_args.use_peft", + apply_on="instantiate", ) parser.link_arguments( "data", diff --git a/modelgenerator/prot_inv_fold/pif_task.py b/modelgenerator/prot_inv_fold/pif_task.py index 0fccc08..33e1aff 100644 --- a/modelgenerator/prot_inv_fold/pif_task.py +++ b/modelgenerator/prot_inv_fold/pif_task.py @@ -38,6 +38,7 @@ def __init__( self.accuracies_str_enc = [] self.acc_metric = MyAccuracy() + @once_only def configure_model(self) -> None: self.lm = self.backbone_fn(None, None) self.tokenizer = self.lm.tokenizer diff --git a/modelgenerator/rna_inv_fold/rif_task.py b/modelgenerator/rna_inv_fold/rif_task.py index 4205610..9cefd4c 100644 --- a/modelgenerator/rna_inv_fold/rif_task.py +++ b/modelgenerator/rna_inv_fold/rif_task.py @@ -88,6 +88,7 @@ def __init__( print(self) + @once_only def configure_model(self) -> None: self.lm = self.backbone_fn(None, None) self.lm.setup() diff --git a/modelgenerator/rna_ss/rna_ss_task.py b/modelgenerator/rna_ss/rna_ss_task.py index 9655fc0..a9dc2be 100644 --- a/modelgenerator/rna_ss/rna_ss_task.py +++ b/modelgenerator/rna_ss/rna_ss_task.py @@ -48,6 +48,7 @@ def __init__( self.THRESHOLD_TUNE_METRIC = "f1" self.THRESHOLD_CANDIDATES = [i / 100 for i in range(1, 30, 1)] + @once_only def configure_model(self) -> None: self.backbone.setup() if self.use_legacy_adapter: diff --git a/tests/cell/test_utils.py b/tests/cell/test_utils.py new file mode 100644 index 0000000..d4ae6e8 --- /dev/null +++ b/tests/cell/test_utils.py @@ -0,0 +1,23 @@ +import pytest +import anndata as ad +import numpy as np +import pandas as pd +from modelgenerator.cell.utils import _ensure_unique_obs_names + +def test_ensure_unique_obs_names(): + """Test that _ensure_unique_obs_names automatically resolves duplicate observation names.""" + # Arrange + X = np.random.rand(3, 3) + obs = pd.DataFrame(index=["cell_1", "cell_1", "cell_2"]) + + # anndata 0.10.x requires explicit types for shape + adata = ad.AnnData(X=X, obs=obs) + + # Pre-condition: names are not unique + assert not adata.obs_names.is_unique + + # Act + adata = _ensure_unique_obs_names(adata) + + # Assert: names should now be unique + assert adata.obs_names.is_unique diff --git a/tests/data/test_obs_names_uniqueness.py b/tests/data/test_obs_names_uniqueness.py new file mode 100644 index 0000000..6acf465 --- /dev/null +++ b/tests/data/test_obs_names_uniqueness.py @@ -0,0 +1,40 @@ +import anndata as ad +import numpy as np +import pytest +import logging +from modelgenerator.cell.utils import _ensure_unique_obs_names + +def test_make_obs_names_unique(): + # Create AnnData with duplicate obs_names + data = np.random.rand(3, 2) + obs_names = ["cell1", "cell1", "cell2"] + adata = ad.AnnData(data) + adata.obs_names = obs_names + + assert not adata.obs_names.is_unique + + # Apply fix + _ensure_unique_obs_names(adata) + + assert adata.obs_names.is_unique + assert list(adata.obs_names) == ["cell1", "cell1-1", "cell2"] + +def test_ensure_unique_obs_names_no_duplicates(): + # Create AnnData with unique obs_names + data = np.random.rand(2, 2) + obs_names = ["cell1", "cell2"] + adata = ad.AnnData(data) + adata.obs_names = obs_names + + assert adata.obs_names.is_unique + + # Apply fix + _ensure_unique_obs_names(adata) + + assert adata.obs_names.is_unique + assert list(adata.obs_names) == ["cell1", "cell2"] + +if __name__ == "__main__": + test_make_obs_names_unique() + test_ensure_unique_obs_names_no_duplicates() + print("All tests passed!") diff --git a/tests/prot_inv_fold/test_pif_task.py b/tests/prot_inv_fold/test_pif_task.py new file mode 100644 index 0000000..e39f5ba --- /dev/null +++ b/tests/prot_inv_fold/test_pif_task.py @@ -0,0 +1,6 @@ +import pytest + +def test_pif_task_initialization(): + """Placeholder test for Protein Inverse Folding task initialization.""" + # TODO: Add deterministic tests for PIFTask using mocked backbones and adapters + pass diff --git a/tests/rna_inv_fold/test_rif_task.py b/tests/rna_inv_fold/test_rif_task.py new file mode 100644 index 0000000..a83bd7b --- /dev/null +++ b/tests/rna_inv_fold/test_rif_task.py @@ -0,0 +1,6 @@ +import pytest + +def test_rif_task_initialization(): + """Placeholder test for RNA Inverse Folding task initialization.""" + # TODO: Add deterministic tests for RIFTask using mocked backbones and adapters + pass diff --git a/tests/rna_ss/test_rna_ss_task.py b/tests/rna_ss/test_rna_ss_task.py new file mode 100644 index 0000000..23feb6e --- /dev/null +++ b/tests/rna_ss/test_rna_ss_task.py @@ -0,0 +1,6 @@ +import pytest + +def test_rna_ss_task_initialization(): + """Placeholder test for RNA Secondary Structure task initialization.""" + # TODO: Add deterministic tests for RNASecondaryStructureTask using mocked backbones and adapters + pass