Replace bare excepts with specific exception types and add explicit weights_only to torch.load

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."

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."

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(

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:

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

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(

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(

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"])

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

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