Reduce logging duplication in AptaTransLightning (#241)

fix #224 
Added _log_metric helper to avoid repeating the same self.log() params
in both _step methods.

---------

Co-authored-by: tarun111111 <tarunpuri2544@gmail.com>

Author: Tarun-goswamii

pyaptamer/aptatrans/_model_lightning.py

@@ -76,7 +76,13 @@ def __init__(
         self.weight_decay = weight_decay
         self.betas = betas
 
-    def _step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
+    def _log_metric(self, name: str, value: Tensor) -> None:
+        """Log metric at epoch level with progress bar display."""
+        self.log(name, value, on_epoch=True, on_step=False, prog_bar=True)
+
+    def _step(
+        self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int, stage: str
+    ) -> Tensor:
         """Defines a single (mini-batch) step in the training/test loop.
 
         Parameters
@@ -85,6 +91,8 @@ def _step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
             A batch of data containing aptamer sequences, protein sequences, and labels.
         batch_idx: int
             Index of the batch.
+        stage: str
+            The stage of the step, either "train" or "test".
 
         Returns
         -------
@@ -100,7 +108,10 @@ def _step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
         y_pred = (y_hat > 0.5).float()
         accuracy = (y_pred == y.float()).float().mean()
 
-        return loss, accuracy
+        self._log_metric(f"{stage}_loss", loss)
+        self._log_metric(f"{stage}_accuracy", accuracy)
+
+        return loss
 
     def training_step(
         self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int
@@ -119,14 +130,7 @@ def training_step(
         Tensor
             The computed loss for the batch.
         """
-        loss, accuracy = self._step(batch, batch_idx)
-
-        self.log("train_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
-        self.log(
-            "train_accuracy", accuracy, on_epoch=True, on_step=False, prog_bar=True
-        )
-
-        return loss
+        return self._step(batch, batch_idx, "train")
 
     def test_step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
         """Defines a single (mini-batch) step in the test loop.
@@ -143,12 +147,7 @@ def test_step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Ten
         Tensor
             The computed loss for the batch.
         """
-        loss, accuracy = self._step(batch, batch_idx)
-
-        self.log("test_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
-        self.log("test_accuracy", accuracy, on_epoch=True, on_step=False, prog_bar=True)
-
-        return loss
+        return self._step(batch, batch_idx, "test")
 
     def configure_optimizers(self) -> torch.optim.Optimizer:
         """Defines the optimizer to be used during training."""
@@ -256,7 +255,9 @@ def __init__(
         self.weight_mlm = weight_mlm
         self.weight_ssp = weight_ssp
 
-    def _step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
+    def _step(
+        self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int, stage: str
+    ) -> Tensor:
         """Defines a single (mini-batch) step in the training/test loop.
 
         The loss function is a weighted sum of the masked language modeling (MLM)
@@ -268,6 +269,8 @@ def _step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
             A batch of data containing aptamer sequences, protein sequences, and labels.
         batch_idx: int
             Index of the batch.
+        stage: str
+            The stage of the step, either "train" or "test".
 
         Returns
         -------
@@ -284,46 +287,8 @@ def _step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
         loss_ssp = F.cross_entropy(y_ssp_hat.transpose(1, 2), y_ssp.long())
         loss = self.weight_mlm * loss_mlm + self.weight_ssp * loss_ssp
 
-        return loss
-
-    def training_step(
-        self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int
-    ) -> Tensor:
-        """Defines a single (mini-batch) step in the training loop.
+        self._log_metric(f"{stage}_loss", loss)
 
-        Parameters
-        ----------
-        batch: tuple[Tensor, Tensor, Tensor]
-            A batch of data containing aptamer sequences, protein sequences, and labels.
-        batch_idx: int
-            Index of the batch.
-
-        Returns
-        -------
-        Tensor
-            The computed loss for the batch.
-        """
-        loss = self._step(batch, batch_idx)
-        self.log("train_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
-        return loss
-
-    def test_step(self, batch: tuple[Tensor, Tensor, Tensor], batch_idx: int) -> Tensor:
-        """Defines a single (mini-batch) step in the test loop.
-
-        Parameters
-        ----------
-        batch: tuple[Tensor, Tensor, Tensor]
-            A batch of data containing aptamer sequences, protein sequences, and labels.
-        batch_idx: int
-            Index of the batch.
-
-        Returns
-        -------
-        Tensor
-            The computed loss for the batch.
-        """
-        loss = self._step(batch, batch_idx)
-        self.log("test_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
         return loss
 
     def configure_optimizers(self) -> torch.optim.Optimizer: