""" Geneformer multi-task cell classifier. **Input data:** | Single-cell transcriptomes as Geneformer rank value encodings with cell state labels for each task in Geneformer .dataset format (generated from single-cell RNAseq data by tokenizer.py). Must contain "unique_cell_id" column for logging. **Usage:** .. code-block :: python >>> from geneformer import MTLClassifier >>> mc = MTLClassifier(task_columns = ["task1", "task2"], ... study_name = "mtl", ... pretrained_path = "/path/pretrained/model", ... train_path = "/path/train/set", ... val_path = "/path/eval/set", ... test_path = "/path/test/set", ... model_save_path = "/results/directory/save_path", ... trials_result_path = "/results/directory/results.txt", ... results_dir = "/results/directory", ... tensorboard_log_dir = "/results/tblogdir", ... hyperparameters = hyperparameters) >>> mc.run_optuna_study() >>> mc.load_and_evaluate_test_model() >>> mc.save_model_without_heads() """ import logging import os from .mtl import eval_utils, train_utils, utils logger = logging.getLogger(__name__) class MTLClassifier: valid_option_dict = { "task_columns": {list}, "train_path": {None, str}, "val_path": {None, str}, "test_path": {None, str}, "pretrained_path": {None, str}, "model_save_path": {None, str}, "results_dir": {None, str}, "batch_size": {None, int}, "n_trials": {None, int}, "study_name": {None, str}, "max_layers_to_freeze": {None, dict}, "epochs": {None, int}, "tensorboard_log_dir": {None, str}, "use_data_parallel": {None, bool}, "use_attention_pooling": {None, bool}, "use_task_weights": {None, bool}, "hyperparameters": {None, dict}, "manual_hyperparameters": {None, dict}, "use_manual_hyperparameters": {None, bool}, "use_wandb": {None, bool}, "wandb_project": {None, str}, "gradient_clipping": {None, bool}, "max_grad_norm": {None, int, float}, "seed": {None, int}, "trials_result_path": {None, str}, } def __init__( self, task_columns=None, train_path=None, val_path=None, test_path=None, pretrained_path=None, model_save_path=None, results_dir=None, trials_result_path=None, batch_size=4, n_trials=15, study_name="mtl", max_layers_to_freeze=None, epochs=1, tensorboard_log_dir="/results/tblogdir", use_data_parallel=False, use_attention_pooling=True, use_task_weights=True, hyperparameters=None, # Default is None manual_hyperparameters=None, # Default is None use_manual_hyperparameters=False, # Default is False use_wandb=False, wandb_project=None, gradient_clipping=False, max_grad_norm=None, seed=42, # Default seed value ): """ Initialize Geneformer multi-task classifier. **Parameters:** task_columns : list | List of tasks for cell state classification | Input data columns are labeled with corresponding task names study_name : None, str | Study name for labeling output files pretrained_path : None, str | Path to pretrained model train_path : None, str | Path to training dataset with task columns and "unique_cell_id" column val_path : None, str | Path to validation dataset with task columns and "unique_cell_id" column test_path : None, str | Path to test dataset with task columns and "unique_cell_id" column model_save_path : None, str | Path to directory to save output model (either full model or model without heads) trials_result_path : None, str | Path to directory to save hyperparameter tuning trial results results_dir : None, str | Path to directory to save results tensorboard_log_dir : None, str | Path to directory for Tensorboard logging results use_data_parallel : None, bool | Whether to use data parallelization use_attention_pooling : None, bool | Whether to use attention pooling use_task_weights : None, bool | Whether to use task weights batch_size : None, int | Batch size to use n_trials : None, int | Number of trials for hyperparameter tuning epochs : None, int | Number of epochs for training max_layers_to_freeze : None, dict | Dictionary with keys "min" and "max" indicating the min and max layers to freeze from fine-tuning (int) | 0: no layers will be frozen; 2: first two layers will be frozen; etc. hyperparameters : None, dict | Dictionary of categorical max and min for each hyperparameter for tuning | For example: | {"learning_rate": {"type":"float", "low":"1e-5", "high":"1e-3", "log":True}, "task_weights": {...}, ...} manual_hyperparameters : None, dict | Dictionary of manually set value for each hyperparameter | For example: | {"learning_rate": 0.001, "task_weights": [1, 1], ...} use_manual_hyperparameters : None, bool | Whether to use manually set hyperparameters use_wandb : None, bool | Whether to use Weights & Biases for logging wandb_project : None, str | Weights & Biases project name gradient_clipping : None, bool | Whether to use gradient clipping max_grad_norm : None, int, float | Maximum norm for gradient clipping seed : None, int | Random seed """ self.task_columns = task_columns self.train_path = train_path self.val_path = val_path self.test_path = test_path self.pretrained_path = pretrained_path self.model_save_path = model_save_path self.results_dir = results_dir self.trials_result_path = trials_result_path self.batch_size = batch_size self.n_trials = n_trials self.study_name = study_name if max_layers_to_freeze is None: # Dynamically determine the range of layers to freeze layer_freeze_range = utils.get_layer_freeze_range(pretrained_path) self.max_layers_to_freeze = {"min": 1, "max": layer_freeze_range["max"]} else: self.max_layers_to_freeze = max_layers_to_freeze self.epochs = epochs self.tensorboard_log_dir = tensorboard_log_dir self.use_data_parallel = use_data_parallel self.use_attention_pooling = use_attention_pooling self.use_task_weights = use_task_weights self.hyperparameters = ( hyperparameters if hyperparameters is not None else { "learning_rate": { "type": "float", "low": 1e-5, "high": 1e-3, "log": True, }, "warmup_ratio": {"type": "float", "low": 0.005, "high": 0.01}, "weight_decay": {"type": "float", "low": 0.01, "high": 0.1}, "dropout_rate": {"type": "float", "low": 0.0, "high": 0.7}, "lr_scheduler_type": {"type": "categorical", "choices": ["cosine"]}, "task_weights": {"type": "float", "low": 0.1, "high": 2.0}, } ) self.manual_hyperparameters = ( manual_hyperparameters if manual_hyperparameters is not None else { "learning_rate": 0.001, "warmup_ratio": 0.01, "weight_decay": 0.1, "dropout_rate": 0.1, "lr_scheduler_type": "cosine", "use_attention_pooling": False, "task_weights": [1, 1], "max_layers_to_freeze": 2, } ) self.use_manual_hyperparameters = use_manual_hyperparameters self.use_wandb = use_wandb self.wandb_project = wandb_project self.gradient_clipping = gradient_clipping self.max_grad_norm = max_grad_norm self.seed = seed if self.use_manual_hyperparameters: logger.warning( "Hyperparameter tuning is highly recommended for optimal results." ) self.validate_options() # set up output directories if self.results_dir is not None: self.trials_results_path = f"{self.results_dir}/results.txt".replace( "//", "/" ) for output_dir in [self.model_save_path, self.results_dir]: if not os.path.exists(output_dir): os.makedirs(output_dir) self.config = { key: value for key, value in self.__dict__.items() if key in self.valid_option_dict } def validate_options(self): # confirm arguments are within valid options and compatible with each other for attr_name, valid_options in self.valid_option_dict.items(): attr_value = self.__dict__[attr_name] if not isinstance(attr_value, (list, dict)): if attr_value in valid_options: continue valid_type = False for option in valid_options: if (option in [int, float, list, dict, bool, str]) and isinstance( attr_value, option ): valid_type = True break if valid_type: continue logger.error( f"Invalid option for {attr_name}. " f"Valid options for {attr_name}: {valid_options}" ) raise ValueError( f"Invalid option for {attr_name}. Valid options for {attr_name}: {valid_options}" ) def run_manual_tuning(self): """ Manual hyperparameter tuning and multi-task fine-tuning of pretrained model. """ required_variable_names = [ "train_path", "val_path", "pretrained_path", "model_save_path", "results_dir", ] required_variables = [ self.train_path, self.val_path, self.pretrained_path, self.model_save_path, self.results_dir, ] req_var_dict = dict(zip(required_variable_names, required_variables)) self.validate_additional_options(req_var_dict) if not self.use_manual_hyperparameters: raise ValueError( "Manual hyperparameters are not enabled. Set use_manual_hyperparameters to True." ) # Ensure manual_hyperparameters are set in the config self.config["manual_hyperparameters"] = self.manual_hyperparameters self.config["use_manual_hyperparameters"] = True train_utils.run_manual_tuning(self.config) def validate_additional_options(self, req_var_dict): missing_variable = False for variable_name, variable in req_var_dict.items(): if variable is None: logger.warning( f"Please provide value to MTLClassifier for required variable {variable_name}" ) missing_variable = True if missing_variable is True: raise ValueError("Missing required variables for MTLClassifier") def run_optuna_study( self, ): """ Hyperparameter optimization and/or multi-task fine-tuning of pretrained model. """ required_variable_names = [ "train_path", "val_path", "pretrained_path", "model_save_path", "results_dir", ] required_variables = [ self.train_path, self.val_path, self.pretrained_path, self.model_save_path, self.results_dir, ] req_var_dict = dict(zip(required_variable_names, required_variables)) self.validate_additional_options(req_var_dict) train_utils.run_optuna_study(self.config) def load_and_evaluate_test_model( self, ): """ Loads previously fine-tuned multi-task model and evaluates on test data. """ required_variable_names = ["test_path", "model_save_path", "results_dir"] required_variables = [self.test_path, self.model_save_path, self.results_dir] req_var_dict = dict(zip(required_variable_names, required_variables)) self.validate_additional_options(req_var_dict) eval_utils.load_and_evaluate_test_model(self.config) # def save_model_without_heads( # self, # ): # """ # Save previously fine-tuned multi-task model without classification heads. # """ # required_variable_names = ["model_save_path"] # required_variables = [self.model_save_path] # req_var_dict = dict(zip(required_variable_names, required_variables)) # self.validate_additional_options(req_var_dict) # utils.save_model_without_heads( # os.path.join(self.model_save_path, "GeneformerMultiTask") # )