add fingerprint dataset, logistic regression model

Author: sfluegel05

chebai/models/base.py

@@ -106,7 +106,8 @@ def _get_prediction_and_labels(
         Returns:
             Tuple[torch.Tensor, torch.Tensor]: Predictions and labels.
         """
-        return output, labels
+        # cast labels to int
+        return output, labels.to(torch.int) if labels is not None else labels
 
     def _process_labels_in_batch(self, batch: XYData) -> torch.Tensor:
         """

chebai/models/classic_ml.py

@@ -0,0 +1,62 @@
+from typing import Any, Dict
+
+import torch
+import tqdm
+from sklearn.exceptions import NotFittedError
+from sklearn.linear_model import LogisticRegression as SklearnLogisticRegression
+
+from chebai.models.base import ChebaiBaseNet
+
+
+class LogisticRegression(ChebaiBaseNet):
+    """
+    Logistic Regression model using scikit-learn, wrapped to fit the ChebaiBaseNet interface.
+    """
+
+    def __init__(self, out_dim: int, input_dim: int, **kwargs):
+        super().__init__(out_dim=out_dim, input_dim=input_dim, **kwargs)
+        self.models = [SklearnLogisticRegression(solver="liblinear") for _ in range(5)]
+
+    def forward(self, x: Dict[str, Any], **kwargs) -> torch.Tensor:
+        print(
+            f"forward called with x[features].shape {x['features'].shape}, self.training {self.training}"
+        )
+        if self.training:
+            self.fit_sklearn(x["features"], x["labels"])
+        try:
+            preds = [
+                torch.from_numpy(model.predict(x["features"]))
+                .to(
+                    x["features"].device
+                    if isinstance(x["features"], torch.Tensor)
+                    else "cpu"
+                )
+                .float()
+                for model in self.models
+            ]
+        except NotFittedError:
+            # Not fitted yet, return zeros
+            print(
+                f"returning default 0s with shape {(x['features'].shape[0], self.out_dim)}"
+            )
+            return torch.zeros(
+                (x["features"].shape[0], self.out_dim),
+                device=(
+                    x["features"].device
+                    if isinstance(x["features"], torch.Tensor)
+                    else "cpu"
+                ),
+            )
+        preds = torch.stack(preds, dim=1)
+        print(f"preds shape {preds.shape}")
+        return preds
+
+    def fit_sklearn(self, X, y):
+        """
+        Fit the underlying sklearn model. X and y should be numpy arrays.
+        """
+        for i, model in tqdm.tqdm(enumerate(self.models), desc="Fitting models"):
+            model.fit(X, y[:, i])
+
+    def configure_optimizers(self, **kwargs):
+        pass

chebai/preprocessing/datasets/chebi.py

@@ -809,6 +809,22 @@ class ChEBIOver50Partial(ChEBIOverXPartial, ChEBIOver50):
     pass
 
 
+class ChEBIOverXFingerprints(ChEBIOverX):
+    """A class that uses Fingerprints for the processed data (used for fixed-length ML models)."""
+
+    READER = dr.FingerprintReader
+
+
+class ChEBIOver100Fingerprints(ChEBIOverXFingerprints, ChEBIOver100):
+    """
+    A class for extracting data from the ChEBI dataset with Fingerprints reader and a threshold of 100.
+
+    Inherits from ChEBIOverXFingerprints and ChEBIOver100.
+    """
+
+    pass
+
+
 class JCIExtendedBPEData(JCIExtendedBase):
     READER = dr.ChemBPEReader
 

chebai/preprocessing/reader.py

@@ -372,3 +372,33 @@ def name(cls) -> str:
     def _read_data(self, raw_data: str) -> List[int]:
         """Convert characters in raw data to their ordinal values."""
         return [ord(s) for s in raw_data]
+
+
+class FingerprintReader(DataReader):
+    """
+    Data reader for chemical data using RDKit fingerprints.
+
+    Args:
+        collator_kwargs: Optional dictionary of keyword arguments for the collator.
+        kwargs: Additional keyword arguments.
+    """
+
+    COLLATOR = DefaultCollator
+
+    def __init__(self, fingerprint_size=1024, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.fingerprint_size = fingerprint_size
+
+    @classmethod
+    def name(cls) -> str:
+        """Returns the name of the data reader."""
+        return "rdkit_fingerprint"
+
+    def _read_data(self, raw_data: str) -> List[int]:
+        """Generate RDKit fingerprint from raw SMILES data."""
+        mol = Chem.MolFromSmiles(raw_data.strip())
+        if mol is None:
+            raise ValueError(f"Invalid SMILES: {raw_data}")
+        return list(
+            Chem.RDKFingerprint(mol, fpSize=self.fingerprint_size).ToBitString()
+        )