Merge branch 'dev' into feature/testing_framework

Author: aditya0by0

chebai/loss/semantic.py

@@ -2,10 +2,11 @@
 import math
 import os
 import pickle
-from typing import Literal
 
 import torch
 
+from typing import Literal, Union
+
 from chebai.loss.bce_weighted import BCEWeighted
 from chebai.preprocessing.datasets.chebi import ChEBIOver100, _ChEBIDataExtractor
 from chebai.preprocessing.datasets.pubchem import LabeledUnlabeledMixed
@@ -14,7 +15,7 @@
 class ImplicationLoss(torch.nn.Module):
     def __init__(
         self,
-        data_extractor: _ChEBIDataExtractor | LabeledUnlabeledMixed,
+        data_extractor: Union[_ChEBIDataExtractor, LabeledUnlabeledMixed],
         base_loss: torch.nn.Module = None,
         tnorm: Literal["product", "lukasiewicz", "xu19"] = "product",
         impl_loss_weight=0.1,  # weight of implication loss in relation to base_loss
@@ -114,7 +115,7 @@ class DisjointLoss(ImplicationLoss):
     def __init__(
         self,
         path_to_disjointness,
-        data_extractor: _ChEBIDataExtractor | LabeledUnlabeledMixed,
+        data_extractor: Union[_ChEBIDataExtractor, LabeledUnlabeledMixed],
         base_loss: torch.nn.Module = None,
         disjoint_loss_weight=100,
         **kwargs,

chebai/result/analyse_sem.py

@@ -7,7 +7,11 @@
 import pandas as pd
 import torch
 import wandb
+
 from torchmetrics.functional.classification import multilabel_auroc, multilabel_f1_score
+import gc
+from typing import List, Union
+
 from utils import *
 
 from chebai.loss.semantic import DisjointLoss
@@ -245,7 +249,7 @@ def analyse_run(
     labeled_data_cls=ChEBIOver100,  # use labels from this dataset for violations
     chebi_version=231,
     results_path=os.path.join("_semantic", "eval_results.csv"),
-    violation_metrics: [str | list[callable]] = "all",
+    violation_metrics: Union[str, List[callable]] = "all",
     verbose_violation_output=False,
 ):
     """Calculates all semantic metrics for given predictions (and supervised metrics if labels are provided),

chebai/trainer/CustomTrainer.py

@@ -64,7 +64,7 @@ def predict_from_file(
             smiles_strings = [inp.strip() for inp in input.readlines()]
         loaded_model.eval()
         predictions = self._predict_smiles(loaded_model, smiles_strings)
-        predictions_df = pd.DataFrame(predictions.detach().numpy())
+        predictions_df = pd.DataFrame(predictions.detach().cpu().numpy())
         if classes_path is not None:
             with open(classes_path, "r") as f:
                 predictions_df.columns = [cls.strip() for cls in f.readlines()]
@@ -74,7 +74,10 @@ def predict_from_file(
     def _predict_smiles(self, model: LightningModule, smiles: List[str]):
         reader = ChemDataReader()
         parsed_smiles = [reader._read_data(s) for s in smiles]
-        x = pad_sequence([torch.tensor(a) for a in parsed_smiles], batch_first=True)
+        x = pad_sequence(
+            [torch.tensor(a, device=model.device) for a in parsed_smiles],
+            batch_first=True,
+        )
         cls_tokens = (
             torch.ones(x.shape[0], dtype=torch.int, device=model.device).unsqueeze(-1)
             * CLS_TOKEN

setup.py

@@ -14,10 +14,7 @@
     author_email="martin.glauer@ovgu.de",
     description="",
     zip_safe=False,
-    # `|` operator used for type hint in chebai/loss/semantic.py is only supported for python_version >= 3.10.
-    # https://stackoverflow.com/questions/76712720/typeerror-unsupported-operand-types-for-type-and-nonetype
-    # python_requires="<=3.11.8",
-    python_requires=">=3.10.0, <3.11.8",
+    python_requires="<=3.11.8",
     install_requires=[
         "certifi",
         "idna",
@@ -51,8 +48,7 @@
         "iterative-stratification",
         "wandb",
         "chardet",
-        # --- commented below due to strange dependency error while setting up new env
-        # "yaml",`
+        "pyyaml",
         "torchmetrics",
     ],
     extras_require={"dev": ["black", "isort", "pre-commit"]},