reformat using black

Author: sfluegel

chebai/callbacks/prediction_callback.py

@@ -3,14 +3,23 @@
 import os
 import pickle
 
+
 class PredictionWriter(BasePredictionWriter):
     def __init__(self, output_dir, write_interval):
         super().__init__(write_interval)
         self.output_dir = output_dir
         self.prediction_file_name = "predictions.pkl"
 
     def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
-        results = [dict(ident=row["data"]["idents"][0], predictions=torch.sigmoid(row["output"]["logits"]).numpy(),
-              labels=row["labels"][0].numpy() if row["labels"] is not None else None) for row in predictions]
-        with open(os.path.join(self.output_dir, self.prediction_file_name), "wb") as fout:
+        results = [
+            dict(
+                ident=row["data"]["idents"][0],
+                predictions=torch.sigmoid(row["output"]["logits"]).numpy(),
+                labels=row["labels"][0].numpy() if row["labels"] is not None else None,
+            )
+            for row in predictions
+        ]
+        with open(
+            os.path.join(self.output_dir, self.prediction_file_name), "wb"
+        ) as fout:
             pickle.dump(results, fout)

chebai/loss/bce_weighted.py

@@ -15,15 +15,15 @@ def __init__(self, beta: float = None, data_extractor: _ChEBIDataExtractor = Non
         self.data_extractor = data_extractor
         super().__init__()
 
-
     def set_pos_weight(self, input):
         if (
             self.beta is not None
             and self.data_extractor is not None
             and all(
                 os.path.exists(os.path.join(self.data_extractor.raw_dir, raw_file))
                 for raw_file in self.data_extractor.raw_file_names
-            ) and self.pos_weight is None
+            )
+            and self.pos_weight is None
         ):
             complete_data = pd.concat(
                 [
@@ -42,9 +42,13 @@ def set_pos_weight(self, input):
             value_counts = []
             for c in complete_data.columns[3:]:
                 value_counts.append(len([v for v in complete_data[c] if v]))
-            weights = [(1 - self.beta) / (1 - pow(self.beta, value)) for value in value_counts]
+            weights = [
+                (1 - self.beta) / (1 - pow(self.beta, value)) for value in value_counts
+            ]
             mean = sum(weights) / len(weights)
-            self.pos_weight = torch.tensor([w / mean for w in weights], device=input.device)
+            self.pos_weight = torch.tensor(
+                [w / mean for w in weights], device=input.device
+            )
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         self.set_pos_weight(input)

chebai/models/electra.py

@@ -18,7 +18,8 @@
 
 logging.getLogger("pysmiles").setLevel(logging.CRITICAL)
 
-from chebai.loss.semantic import DisjointLoss as ElectraChEBIDisjointLoss # noqa
+from chebai.loss.semantic import DisjointLoss as ElectraChEBIDisjointLoss  # noqa
+
 
 class ElectraPre(ChebaiBaseNet):
     NAME = "ElectraPre"