make ruff rules similar between models and recipes (#1369)

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

bionemo-recipes/recipes/.ruff.toml

@@ -2,7 +2,7 @@ line-length = 119
 target-version = "py312"
 
 [lint]
-ignore = ["D100", "E501", "N811", "N814"]
+ignore = ["C901", "D100", "E501", "N811", "N814"]
 select = [
     "C",    # Pylint conventions
     "D",    # Documentation formatting

bionemo-recipes/recipes/codonfm_ptl_te/src/config.py

@@ -41,7 +41,7 @@
 
 
 # Datasets
-def get_dataset_config(args: Any, process_item_cfg: fdl.Partial) -> fdl.Config:  # noqa: C901
+def get_dataset_config(args: Any, process_item_cfg: fdl.Partial) -> fdl.Config:
     """Builds the dataset configuration."""
     class_name = args.dataset_name
     if class_name == "CodonMemmapDataset":

bionemo-recipes/recipes/codonfm_ptl_te/src/data/codon_memmap_dataset.py

@@ -62,7 +62,7 @@ class CodonMemmapDataset(torch.utils.data.Dataset):
     **IMPORTANT** this works because we aren't mutating the underlying data between these splits, and only the indices are modified.
     """
 
-    def __init__(  # noqa: C901, D107
+    def __init__(  # noqa: D107
         self,
         data_path: str,
         tokenizer: Callable,

bionemo-recipes/recipes/codonfm_ptl_te/src/models/components/encodon_te.py

@@ -106,7 +106,7 @@ def __init__(self, config):
         self.te_rope_emb = None
         self._init_weights()
 
-    def reset_cls_parameters(self):  # noqa: C901
+    def reset_cls_parameters(self):
         """Resets the parameters of the classification head."""
         for module in self.cls.modules():
             if isinstance(module, (nn.Linear, transformer_engine.pytorch.Linear)):
@@ -136,7 +136,7 @@ def reset_cls_parameters(self):  # noqa: C901
                 if module.layer_norm_bias is not None:
                     module.layer_norm_bias.data.zero_()
 
-    def _init_weights(self):  # noqa: C901
+    def _init_weights(self):
         """Initializes the weights of the model using the MAGNETO initialization scheme."""
         for name, module in self.named_modules():
             if isinstance(module, (nn.Linear, transformer_engine.pytorch.Linear)):

bionemo-recipes/recipes/codonfm_ptl_te/src/models/components/encodon_te_layer.py

@@ -266,7 +266,7 @@ class EncodonTELayer(torch.nn.Module):
                     QK normalization at different points.
     """  # noqa: D205
 
-    def __init__(  # noqa: D107, C901
+    def __init__(  # noqa: D107
         self,
         hidden_size: int,
         ffn_hidden_size: int,
@@ -573,7 +573,7 @@ def set_context_parallel_group(
             if hasattr(child, "set_context_parallel_group"):
                 child.set_context_parallel_group(cp_group, cp_global_ranks, cp_stream, cp_comm_type)
 
-    def forward(  # noqa: C901
+    def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,

bionemo-recipes/recipes/codonfm_ptl_te/src/models/components/encodon_te_mha.py

@@ -591,7 +591,7 @@ def set_context_parallel_group(
             if hasattr(child, "set_context_parallel_group"):
                 child.set_context_parallel_group(cp_group, cp_global_ranks, cp_stream, cp_comm_type)
 
-    def forward(  # noqa: C901
+    def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,

bionemo-recipes/recipes/codonfm_ptl_te/src/models/encodon_pl.py

@@ -139,7 +139,7 @@ def __init__(  # noqa: D107
             # When downstream head is disabled, use cross-entropy for language modeling
             self.loss = torch.nn.CrossEntropyLoss(ignore_index=ignore_index)
 
-    def configure_model(self, state_dict: Optional[Dict[str, Any]] = None) -> None:  # noqa: C901
+    def configure_model(self, state_dict: Optional[Dict[str, Any]] = None) -> None:
         """Configure the underlying model and optionally load weights.
 
         Sets up the base EnCodon model, attaches optional downstream heads, and

bionemo-recipes/recipes/codonfm_ptl_te/src/models/encodon_te_pl.py

@@ -143,7 +143,7 @@ def __init__(  # noqa: D107
 
         self.attn_input_format = attn_input_format
 
-    def configure_model(self, state_dict: Optional[Dict[str, Any]] = None) -> None:  # noqa: C901
+    def configure_model(self, state_dict: Optional[Dict[str, Any]] = None) -> None:
         """Configure the model, optionally loading a state dict.
 
         Args:

bionemo-recipes/recipes/codonfm_ptl_te/src/runner.py

@@ -228,7 +228,7 @@ def get_parser():  # noqa: D103
     return parser
 
 
-def main():  # noqa: C901, D103
+def main():  # noqa: D103
     parser = get_parser()
     args = parser.parse_args()
     if args.mode in ["eval"] and not args.checkpoint_path:

bionemo-recipes/recipes/codonfm_ptl_te/src/utils/pred_writer.py

@@ -104,7 +104,7 @@ def _save_predictions(self, trainer, batch_idx):
             np.save(file_path, value)
         self.predictions_buffer.clear()
 
-    def _merge_predictions(self):  # noqa: C901
+    def _merge_predictions(self):
         """Merge per-rank shard files into a single array per key.
 
         The method reconstructs the original sample ordering based on the DDP