remove unused imports + minor type hints update

Author: aditya0by0

chebai/callbacks.py

@@ -3,7 +3,7 @@
 
 from lightning.pytorch.callbacks import BasePredictionWriter
 import torch
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Literal
 
 
 class ChebaiPredictionWriter(BasePredictionWriter):
@@ -22,7 +22,7 @@ class ChebaiPredictionWriter(BasePredictionWriter):
     def __init__(
         self,
         output_dir: str,
-        write_interval: str,
+        write_interval: Literal["batch", "epoch", "batch_and_epoch"],
         target_file: str = "predictions.json",
     ) -> None:
         super().__init__(write_interval)

chebai/callbacks/prediction_callback.py

@@ -3,7 +3,7 @@
 import torch
 import os
 import pickle
-from typing import Sequence, Any
+from typing import Sequence, Any, Literal
 
 
 class PredictionWriter(BasePredictionWriter):
@@ -15,7 +15,11 @@ class PredictionWriter(BasePredictionWriter):
         write_interval (str): When to write predictions. Options are "batch" or "epoch".
     """
 
-    def __init__(self, output_dir: str, write_interval: str):
+    def __init__(
+        self,
+        output_dir: str,
+        write_interval: Literal["batch", "epoch", "batch_and_epoch"],
+    ):
         super().__init__(write_interval)
         self.output_dir = output_dir
         self.prediction_file_name = "predictions.pkl"

chebai/loss/bce_weighted.py

@@ -1,9 +1,10 @@
+from typing import Optional
+
 import torch
 from chebai.preprocessing.datasets.base import XYBaseDataModule
 from chebai.preprocessing.datasets.pubchem import LabeledUnlabeledMixed
 import pandas as pd
 import os
-import pickle
 
 
 class BCEWeighted(torch.nn.BCEWithLogitsLoss):

chebai/molecule.py

@@ -10,7 +10,6 @@
 
 import networkx as nx
 import numpy as np
-import six
 import torch
 
 logger = logging.getLogger(__name__)
@@ -131,6 +130,9 @@ def create_directed_graphs(self):
     def create_feature_vectors(self):
         """
         Creates feature vectors based on local environments of atoms.
+
+        Note:
+            create a three-dimensional matrix I, such that I_{i,j} is the local input vector for jth vertex in ith DAG
         """
         length_of_bond_features = Molecule.num_bond_features()
         length_of_atom_features = Molecule.num_atom_features()

chebai/preprocessing/datasets/tox21.py

@@ -3,21 +3,17 @@
 import csv
 import gzip
 import os
-import random
 import shutil
 import zipfile
-from typing import List, Dict, Generator
+from typing import List, Dict, Generator, Optional
 
 from rdkit import Chem
 from sklearn.model_selection import GroupShuffleSplit, train_test_split
 import numpy as np
 import torch
-import pysmiles
 
 from chebai.preprocessing import reader as dr
-from chebai.preprocessing.datasets.base import MergedDataset, XYBaseDataModule
-from chebai.preprocessing.datasets.chebi import JCIExtendedTokenData
-from chebai.preprocessing.datasets.pubchem import Hazardous
+from chebai.preprocessing.datasets.base import XYBaseDataModule
 
 
 class Tox21MolNet(XYBaseDataModule):
@@ -218,7 +214,7 @@ def download(self) -> None:
         )
 
     def _retrieve_file(
-        self, url: str, target_file: str, compression: str = None
+        self, url: str, target_file: str, compression: Optional[str] = None
     ) -> None:
         """Retrieves a file from a URL and saves it locally.
 

chebai/preprocessing/structures.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Tuple, Union
 from torch.utils.data.dataset import T_co
 import networkx as nx
 import torch

chebai/result/classification.py

@@ -1,22 +1,16 @@
-import os
-from typing import List, Optional, Tuple
+from typing import List
 
 import matplotlib.pyplot as plt
 import pandas as pd
 import seaborn as sns
-import torch
 from torch import Tensor
 from torchmetrics.classification import (
     MultilabelF1Score,
     MultilabelPrecision,
     MultilabelRecall,
 )
-import tqdm
 
 from chebai.callbacks.epoch_metrics import MacroF1
-from chebai.models import ChebaiBaseNet
-from chebai.models.electra import Electra
-from chebai.preprocessing.datasets import XYBaseDataModule
 from chebai.result.utils import *