ruff

Author: ncfrey

.DS_Store

@@ -13,11 +13,15 @@ def instantiate_callbacks(callbacks_cfg: DictConfig) -> list[Callback]:
         return callbacks
 
     if not isinstance(callbacks_cfg, DictConfig):
-        raise TypeError("[instantiate_callbacks] Callbacks config must be a DictConfig!")
+        raise TypeError(
+            "[instantiate_callbacks] Callbacks config must be a DictConfig!"
+        )
 
     for _, cb_conf in callbacks_cfg.items():
         if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf:
-            print(f"[instantiate_callbacks] Instantiating callback <{cb_conf._target_}>")
+            print(
+                f"[instantiate_callbacks] Instantiating callback <{cb_conf._target_}>"
+            )
             callbacks.append(hydra.utils.instantiate(cb_conf))
 
     return callbacks

src/lobster/cmdline/_utils.py

@@ -99,7 +99,9 @@ def __call__(self, raw_batch: Sequence[Tuple[str, str]]):
         if self._contact_maps:
             # Flatten the output of Atom3D transforms
             flattened_batch = [
-                (a, b, c) for ((a, b), c) in raw_batch if (a is not None) and (b is not None)
+                (a, b, c)
+                for ((a, b), c) in raw_batch
+                if (a is not None) and (b is not None)
             ]
             batch_size = len(flattened_batch)
             if batch_size == 0:
@@ -114,8 +116,12 @@ def __call__(self, raw_batch: Sequence[Tuple[str, str]]):
 
         if self.truncation_seq_length:
             # NOTE - This removes eos token for long sequences. Should we re-add eos or keep as is?
-            seq1_tokenized = [seq[: self.truncation_seq_length] for seq in seq1_tokenized]
-            seq2_tokenized = [seq[: self.truncation_seq_length] for seq in seq2_tokenized]
+            seq1_tokenized = [
+                seq[: self.truncation_seq_length] for seq in seq1_tokenized
+            ]
+            seq2_tokenized = [
+                seq[: self.truncation_seq_length] for seq in seq2_tokenized
+            ]
 
         tokens1 = pad_sequence(
             seq1_tokenized, batch_first=True, padding_value=self.tokenizer.pad_token_id

src/lobster/data/_collate.py

@@ -3,7 +3,11 @@
 
 class FarthestFirstTraversal:
     def __init__(
-        self, num_samples: int, k: int = 10, random_seed: int = 0xDEADBEEF, p_norm: int = 2
+        self,
+        num_samples: int,
+        k: int = 10,
+        random_seed: int = 0xDEADBEEF,
+        p_norm: int = 2,
     ):
         """
         Parameters
@@ -45,7 +49,10 @@ def str_fft(self, inputs: list[str]):
         inputs = [inputs[i] for i in perm]
         centroids = [inputs[i] for i in range(self._k)]
         while len(centroids) < self._num_samples:
-            dist = [min(self._levenshtein(str1, str2) for str2 in centroids) for str1 in inputs]
+            dist = [
+                min(self._levenshtein(str1, str2) for str2 in centroids)
+                for str1 in inputs
+            ]
             farthest = dist.index(max(dist))
             if inputs[farthest] in centroids:
                 break

src/lobster/data/_farthest_first_traversal.py

@@ -22,7 +22,9 @@ def summarize_binary_classification_metrics(preds, labels):
     """
     # Initialize metric objects
     accuracy = Accuracy(task="binary")
-    precision = Precision(task="binary", num_classes=2, average="micro")  # binary classification
+    precision = Precision(
+        task="binary", num_classes=2, average="micro"
+    )  # binary classification
     recall = Recall(task="binary", num_classes=2, average="micro")
     f1_score = F1Score(task="binary", num_classes=2, average="micro")
     auroc = AUROC(task="binary", num_classes=1)  # binary classification

src/lobster/metrics/_binary_classification.py

@@ -134,7 +134,7 @@ def __init__(
         key_bias=True,
         value_bias=True,
         intermediate_bias=True,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(pad_token_id=pad_token_id, **kwargs)
 

src/lobster/model/_seq2seq_configuration.py

@@ -139,10 +139,14 @@ def __init__(self, config, **kwargs):
             act,
         )
         for _i in range(config.num_inner_mlps):
-            self.implicit_filter.append(nn.Linear(config.filter_order, config.filter_order))
+            self.implicit_filter.append(
+                nn.Linear(config.filter_order, config.filter_order)
+            )
             self.implicit_filter.append(act)
 
-        self.implicit_filter.append(nn.Linear(config.filter_order, config.d_model, bias=False))
+        self.implicit_filter.append(
+            nn.Linear(config.filter_order, config.d_model, bias=False)
+        )
 
         self.modulation = HyenaExponentialModulation(config.d_model)
 
@@ -191,7 +195,11 @@ def __init__(
         self.out_proj = nn.Linear(self.d_model, self.d_model)
 
         self.short_filter = nn.Conv1d(
-            inner_width, inner_width, config.short_filter_order, padding=2, groups=inner_width
+            inner_width,
+            inner_width,
+            config.short_filter_order,
+            padding=2,
+            groups=inner_width,
         )
         self.filter_fn = HyenaFilter(config)
 
@@ -297,7 +305,9 @@ def __init__(self, config, padding_idx=None):
             vocab_size += config.pad_vocab_size_multiple - (
                 vocab_size % config.pad_vocab_size_multiple
             )
-        self.word_embeddings = nn.Embedding(vocab_size, config.d_model, padding_idx=padding_idx)
+        self.word_embeddings = nn.Embedding(
+            vocab_size, config.d_model, padding_idx=padding_idx
+        )
 
     def forward(self, input_ids):
         """
@@ -330,7 +340,9 @@ def forward(self, input_ids, inputs_embeds=None, output_hidden_states=False):
 
         for layer in self.layers:
             if self.gradient_checkpointing and self.training:
-                hidden_states = self._gradient_checkpointing_func(layer.__call__, hidden_states)
+                hidden_states = self._gradient_checkpointing_func(
+                    layer.__call__, hidden_states
+                )
             else:
                 hidden_states = layer(hidden_states)
             if output_hidden_states:
@@ -349,7 +361,9 @@ class HyenaDNAPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["HyenaBlock"]
     _skip_keys_device_placement = "past_key_values"
-    _keys_to_ignore_on_load_missing = [r"freq"]  # Shared tensors that safetensors merges
+    _keys_to_ignore_on_load_missing = [
+        r"freq"
+    ]  # Shared tensors that safetensors merges
 
     def _init_weights(self, module, initializer_range=0.02):
         if isinstance(module, nn.Linear):
@@ -368,13 +382,17 @@ def _init_weights(self, module, initializer_range=0.02):
             if name in ["out_proj.weight", "fc2.weight"]:
                 # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
                 nn.init.normal_(
-                    p, mean=0.0, std=initializer_range / math.sqrt(2 * self.config.num_layers)
+                    p,
+                    mean=0.0,
+                    std=initializer_range / math.sqrt(2 * self.config.num_layers),
                 )
             # If using GLU activation for now, we scale the std by 2
             elif name in ["output_linear.0.weight"]:
                 # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
                 nn.init.normal_(
-                    p, mean=0.0, std=initializer_range / math.sqrt(2 * self.config.num_layers)
+                    p,
+                    mean=0.0,
+                    std=initializer_range / math.sqrt(2 * self.config.num_layers),
                 )
 
 
@@ -388,16 +406,22 @@ def __init__(self, config, **kwargs) -> None:
         # Initialize weights and apply final processing
         self.post_init()
 
-    def forward(self, input_ids, inputs_embeds=None, output_hidden_states=None, return_dict=None):
+    def forward(
+        self, input_ids, inputs_embeds=None, output_hidden_states=None, return_dict=None
+    ):
         output_hidden_states = (
             output_hidden_states
             if output_hidden_states is not None
             else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
 
         hidden_states, all_hidden_states = self.backbone(
-            input_ids, inputs_embeds=inputs_embeds, output_hidden_states=output_hidden_states
+            input_ids,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
         )
         if return_dict:
             return BaseModelOutputWithNoAttention(
@@ -451,13 +475,14 @@ def forward(
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[Tuple, CausalLMOutput]:
-
         output_hidden_states = (
             output_hidden_states
             if output_hidden_states is not None
             else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.hyena(
@@ -525,7 +550,9 @@ def forward(
             config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
             `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
         """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
 
         transformer_outputs = self.hyena(
             input_ids,
@@ -542,7 +569,9 @@ def forward(
             batch_size = inputs_embeds.shape[0]
 
         if self.config.pad_token_id is None and batch_size != 1:
-            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+            raise ValueError(
+                "Cannot handle batch sizes > 1 if no padding token is defined."
+            )
         if self.config.pad_token_id is None:
             sequence_lengths = -1
         else:
@@ -553,7 +582,9 @@ def forward(
             else:
                 sequence_lengths = -1
 
-        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+        pooled_logits = logits[
+            torch.arange(batch_size, device=logits.device), sequence_lengths
+        ]
 
         loss = None
         if labels is not None:
@@ -576,7 +607,9 @@ def forward(
                     loss = loss_fct(pooled_logits, labels)
             elif self.config.problem_type == "single_label_classification":
                 loss_fct = nn.CrossEntropyLoss()
-                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+                loss = loss_fct(
+                    pooled_logits.view(-1, self.num_labels), labels.view(-1)
+                )
             elif self.config.problem_type == "multi_label_classification":
                 loss_fct = nn.BCEWithLogitsLoss()
                 loss = loss_fct(pooled_logits, labels)

src/lobster/model/hyena/_hyena_base.py

@@ -22,7 +22,7 @@ def __init__(
         pad_token: str = "[PAD]",
         cls_token: str = "[CLS]",
         mask_token: str = "[MASK]",
-        **kwargs
+        **kwargs,
     ):
         super().__init__(
             vocab_file=vocab_file,
@@ -33,7 +33,7 @@ def __init__(
             pad_token=pad_token,
             cls_token=cls_token,
             mask_token=mask_token,
-            **kwargs
+            **kwargs,
         )
         self.padding_idx = self.vocab[pad_token]
         self.masking_idx = self.vocab[mask_token]