chebai - docstring + typehints

Author: aditya0by0

chebai/__init__.py

@@ -1,10 +1,29 @@
 import os
-
 import torch
+from typing import Any
 
+# Get the absolute path of the current file's directory
 MODULE_PATH = os.path.abspath(os.path.dirname(__file__))
 
 
 class CustomTensor(torch.Tensor):
-    def __new__(cls, data):
+    """
+    A custom tensor class inheriting from `torch.Tensor`.
+
+    This class allows for the creation of tensors using the provided data.
+
+    Attributes:
+        data (Any): The data to be converted into a tensor.
+    """
+
+    def __new__(cls, data: Any) -> "CustomTensor":
+        """
+        Creates a new instance of CustomTensor.
+
+        Args:
+            data (Any): The data to be converted into a tensor.
+
+        Returns:
+            CustomTensor: A tensor containing the provided data.
+        """
         return torch.tensor(data)

chebai/__main__.py

@@ -1,4 +1,10 @@
 from chebai.cli import cli
 
 if __name__ == "__main__":
+    """
+    Entry point for the CLI application.
+
+    This script calls the `cli` function from the `chebai.cli` module
+    when executed as the main program.
+    """
     cli()

chebai/callbacks.py

@@ -3,29 +3,74 @@
 
 from lightning.pytorch.callbacks import BasePredictionWriter
 import torch
+from typing import Any, Dict, List, Union
 
 
 class ChebaiPredictionWriter(BasePredictionWriter):
-    def __init__(self, output_dir, write_interval, target_file="predictions.json"):
+    """
+    A custom prediction writer for saving batch and epoch predictions during model training.
+
+    This class inherits from `BasePredictionWriter` and is designed to save predictions
+    in a specified output directory at specified intervals.
+
+    Args:
+        output_dir (str): The directory where predictions will be saved.
+        write_interval (str): The interval at which predictions will be written.
+        target_file (str): The name of the file where epoch predictions will be saved (default: "predictions.json").
+    """
+
+    def __init__(
+        self,
+        output_dir: str,
+        write_interval: str,
+        target_file: str = "predictions.json",
+    ) -> None:
         super().__init__(write_interval)
         self.output_dir = output_dir
         self.target_file = target_file
 
     def write_on_batch_end(
         self,
-        trainer,
-        pl_module,
-        prediction,
-        batch_indices,
-        batch,
-        batch_idx,
-        dataloader_idx,
-    ):
-        outpath = os.path.join(self.output_dir, dataloader_idx, f"{batch_idx}.pt")
-        os.makedirs(outpath, exist_ok=True)
+        trainer: Any,
+        pl_module: Any,
+        prediction: Union[torch.Tensor, List[torch.Tensor]],
+        batch_indices: List[int],
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        """
+        Saves batch predictions at the end of each batch.
+
+        Args:
+            trainer (Any): The trainer instance.
+            pl_module (Any): The LightningModule instance.
+            prediction (Union[torch.Tensor, List[torch.Tensor]]): The prediction output from the model.
+            batch_indices (List[int]): The indices of the batch.
+            batch (Any): The current batch.
+            batch_idx (int): The index of the batch.
+            dataloader_idx (int): The index of the dataloader.
+        """
+        outpath = os.path.join(self.output_dir, str(dataloader_idx), f"{batch_idx}.pt")
+        os.makedirs(os.path.dirname(outpath), exist_ok=True)
         torch.save(prediction, outpath)
 
-    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
+    def write_on_epoch_end(
+        self,
+        trainer: Any,
+        pl_module: Any,
+        predictions: List[Dict[str, Any]],
+        batch_indices: List[int],
+    ) -> None:
+        """
+        Saves all predictions at the end of each epoch in a JSON file.
+
+        Args:
+            trainer (Any): The trainer instance.
+            pl_module (Any): The LightningModule instance.
+            predictions (List[Dict[str, Any]]): The list of prediction outputs from the model.
+            batch_indices (List[int]): The indices of the batches.
+        """
         pred_list = []
         for p in predictions:
             idents = p["data"]["idents"]
@@ -35,7 +80,7 @@ def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
             else:
                 labels = [None for _ in idents]
             output = torch.sigmoid(p["output"]["logits"]).tolist()
-            for i, l, p in zip(idents, labels, output):
-                pred_list.append(dict(ident=i, labels=l, predictions=p))
+            for i, l, o in zip(idents, labels, output):
+                pred_list.append(dict(ident=i, labels=l, predictions=o))
         with open(os.path.join(self.output_dir, self.target_file), "wt") as fout:
             json.dump(pred_list, fout)

chebai/cli.py

@@ -1,15 +1,38 @@
 from typing import Dict, Set
-
 from lightning.pytorch.cli import LightningArgumentParser, LightningCLI
-
 from chebai.trainer.CustomTrainer import CustomTrainer
 
 
 class ChebaiCLI(LightningCLI):
-    def __init__(self, *args, **kwargs):
+    """
+    Custom CLI subclass for Chebai project based on PyTorch Lightning's LightningCLI.
+
+    Args:
+        save_config_kwargs (dict): Keyword arguments for saving configuration.
+        parser_kwargs (dict): Keyword arguments for parser configuration.
+
+    Attributes:
+        save_config_kwargs (dict): Configuration options for saving.
+        parser_kwargs (dict): Configuration options for the argument parser.
+    """
+
+    def __init__(self, save_config_kwargs: dict, parser_kwargs: dict):
+        """
+        Initialize ChebaiCLI with custom trainer and configure parser settings.
+
+        Args:
+            save_config_kwargs (dict): Keyword arguments for saving configuration.
+            parser_kwargs (dict): Keyword arguments for parser configuration.
+        """
         super().__init__(trainer_class=CustomTrainer, *args, **kwargs)
 
     def add_arguments_to_parser(self, parser: LightningArgumentParser):
+        """
+        Add custom arguments to the argument parser.
+
+        Args:
+            parser (LightningArgumentParser): Argument parser instance.
+        """
         for kind in ("train", "val", "test"):
             for average in ("micro-f1", "macro-f1", "balanced-accuracy"):
                 parser.link_arguments(
@@ -25,7 +48,12 @@ def add_arguments_to_parser(self, parser: LightningArgumentParser):
 
     @staticmethod
     def subcommands() -> Dict[str, Set[str]]:
-        """Defines the list of available subcommands and the arguments to skip."""
+        """
+        Defines the list of available subcommands and the arguments to skip.
+
+        Returns:
+            Dict[str, Set[str]]: Dictionary where keys are subcommands and values are sets of arguments to skip.
+        """
         return {
             "fit": {"model", "train_dataloaders", "val_dataloaders", "datamodule"},
             "validate": {"model", "dataloaders", "datamodule"},
@@ -36,6 +64,9 @@ def subcommands() -> Dict[str, Set[str]]:
 
 
 def cli():
+    """
+    Main function to instantiate and run the ChebaiCLI.
+    """
     r = ChebaiCLI(
         save_config_kwargs={"config_filename": "lightning_config.yaml"},
         parser_kwargs={"parser_mode": "omegaconf"},