diff --git a/helical/models/caduceus/__init__.py b/helical/models/caduceus/__init__.py index 9f6d65f3..58b339d2 100644 --- a/helical/models/caduceus/__init__.py +++ b/helical/models/caduceus/__init__.py @@ -4,7 +4,7 @@ from .model import Caduceus from .caduceus_config import CaduceusConfig from .fine_tuning_model import CaduceusFineTuningModel -except: +except ImportError: LOGGER = logging.getLogger(__name__) LOGGER.error( "Caduceus not available: If you want to use this model, ensure you have a CUDA GPU and have installed the optional helical[mamba-ssm] dependencies." diff --git a/helical/models/evo_2/__init__.py b/helical/models/evo_2/__init__.py index d9c8e891..8243b3ea 100644 --- a/helical/models/evo_2/__init__.py +++ b/helical/models/evo_2/__init__.py @@ -3,7 +3,7 @@ try: from .evo_2_config import Evo2Config from .model import Evo2 -except: +except ImportError: LOGGER = logging.getLogger(__name__) LOGGER.info( "Evo 2 not available: If you want to use this model, please follow the installation instructions in the evo_2/README." diff --git a/helical/models/evo_2/model.py b/helical/models/evo_2/model.py index 4593601c..5b894277 100644 --- a/helical/models/evo_2/model.py +++ b/helical/models/evo_2/model.py @@ -25,7 +25,7 @@ from vortex.model.generation import generate as vortex_generate from vortex.model.model import StripedHyena from vortex.model.utils import dotdict, load_checkpoint -except: +except ImportError: message = "Vortex is required for Evo 2 model. Please follow the instructions in the README to install it." LOGGER.error(message) raise ImportError( @@ -380,7 +380,7 @@ def _load_evo2_model(self, model_name: str = None): filename=filename, ) # If file is split, download and join parts - except: + except huggingface_hub.utils.EntryNotFoundError: print(f"Loading checkpoint shards for {filename}") # If file is split, get the first part's directory to use the same cache location weights_path = os.path.join( diff --git a/helical/models/geneformer/geneformer_utils.py b/helical/models/geneformer/geneformer_utils.py index eba6ff04..660c0d81 100644 --- a/helical/models/geneformer/geneformer_utils.py +++ b/helical/models/geneformer/geneformer_utils.py @@ -295,7 +295,7 @@ def pad_tensor(tensor, pad_token_id, max_len): def pad_3d_tensor(tensor, pad_token_id, max_len, dim): if dim == 0: - raise Exception("dim 0 usually does not need to be padded.") + raise ValueError("dim 0 usually does not need to be padded.") if dim == 1: pad = (0, 0, 0, max_len - tensor.size()[dim]) elif dim == 2: diff --git a/helical/models/hyena_dna/pretrained_model.py b/helical/models/hyena_dna/pretrained_model.py index f126f64a..60b327e0 100644 --- a/helical/models/hyena_dna/pretrained_model.py +++ b/helical/models/hyena_dna/pretrained_model.py @@ -47,8 +47,8 @@ def load_weights(scratch_dict, pretrained_dict, checkpointing=False): key_loaded = inject_substring(key_loaded) try: scratch_dict[key] = pretrained_dict[key_loaded] - except: - raise Exception("key mismatch in the state dicts!") + except KeyError: + raise KeyError("key mismatch in the state dicts!") # scratch_dict has been updated return scratch_dict diff --git a/helical/models/scgpt/fine_tuning_model.py b/helical/models/scgpt/fine_tuning_model.py index abd8131e..cfe1ff13 100644 --- a/helical/models/scgpt/fine_tuning_model.py +++ b/helical/models/scgpt/fine_tuning_model.py @@ -168,7 +168,7 @@ def train( try: use_batch_labels = train_input_data.batch_ids is not None - except: + except AttributeError: use_batch_labels = False collator = DataCollator( @@ -305,7 +305,7 @@ def get_outputs( try: use_batch_labels = dataset.batch_ids is not None - except: + except AttributeError: use_batch_labels = False collator = DataCollator( diff --git a/helical/models/scgpt/model.py b/helical/models/scgpt/model.py index 9019e4be..101a237b 100644 --- a/helical/models/scgpt/model.py +++ b/helical/models/scgpt/model.py @@ -131,7 +131,7 @@ def get_embeddings( try: use_batch_labels = dataset.batch_ids is not None - except: + except AttributeError: use_batch_labels = False collator = DataCollator( diff --git a/helical/models/scgpt/scgpt_utils.py b/helical/models/scgpt/scgpt_utils.py index 34e5ad52..4a24be96 100644 --- a/helical/models/scgpt/scgpt_utils.py +++ b/helical/models/scgpt/scgpt_utils.py @@ -100,7 +100,7 @@ def load_model(model_configs: scGPTConfig): load_pretrained( model, - torch.load(model_configs["model_path"], map_location=model_configs["device"]), + torch.load(model_configs["model_path"], map_location=model_configs["device"], weights_only=True), verbose=False, ) model.to(model_configs["device"]) diff --git a/helical/models/uce/gene_embeddings.py b/helical/models/uce/gene_embeddings.py index 79332d44..7c50d1ca 100644 --- a/helical/models/uce/gene_embeddings.py +++ b/helical/models/uce/gene_embeddings.py @@ -66,7 +66,8 @@ def load_gene_embeddings_adata( species: { gene_symbol.lower(): gene_embedding for gene_symbol, gene_embedding in torch.load( - species_to_gene_embedding_path[species] + species_to_gene_embedding_path[species], + weights_only=False, # gene embedding dicts contain non-tensor metadata ).items() } for species in species_names @@ -104,7 +105,8 @@ def load_gene_embeddings_adata( species: [ gene_symbol.lower() for gene_symbol, _ in torch.load( - species_to_gene_embedding_path[species] + species_to_gene_embedding_path[species], + weights_only=False, # gene embedding dicts contain non-tensor metadata ).items() ] for species in species_names diff --git a/helical/models/uce/uce_utils.py b/helical/models/uce/uce_utils.py index 96463545..dcabcd4b 100644 --- a/helical/models/uce/uce_utils.py +++ b/helical/models/uce/uce_utils.py @@ -53,7 +53,7 @@ def get_ESM2_embeddings(token_file: Union[Path, str], token_dim: int) -> torch.T The token file loaded as a torch.Tensor. """ - all_pe = torch.load(token_file) + all_pe = torch.load(token_file, weights_only=True) # TODO: Why this if clause and why this magic number 143574? if all_pe.shape[0] == 143574: @@ -120,9 +120,9 @@ def prepare_expression_counts_file( LOGGER.info( f"Passed the gene expressions (with shape={shape} and max gene count data {gene_expression.max()}) to {filename}" ) - except: + except Exception: LOGGER.error(f"Error during preparation of npz file {filename}.") - raise Exception + raise ## writing a funciton to load the model @@ -155,7 +155,7 @@ def load_model( empty_pe.requires_grad = False model.pe_embedding = torch.nn.Embedding.from_pretrained(empty_pe) model.load_state_dict( - torch.load(model_path, map_location=model_config["device"]), strict=True + torch.load(model_path, map_location=model_config["device"], weights_only=True), strict=True ) # This will make sure that you don't overwrite the tokens in case you're embedding species from the training data