diff --git a/utilities/ml.py b/utilities/ml.py index e5ab87a3..92b3e960 100644 --- a/utilities/ml.py +++ b/utilities/ml.py @@ -81,14 +81,6 @@ class Dataset(ABC): def generate_dataset(reactions, *args, **kwargs): pass - def at(self, position: int) -> Dataset: - """See __getitem__""" - return self[position] - - def limit(self, limit: int) -> Dataset: - """See __getitem__""" - return self[:limit] - def __iter__(self): """Use polars iter_rows for iterating over the dataset""" return self.df.iter_rows() @@ -180,7 +172,7 @@ class RuleBasedDataset(Dataset): def structure_id(self, index: int): """Get the UUID of a compound""" - return self.df.item(index, "structure_id") + return self.item(index, "structure_id") def X(self, exclude_id_col=True, na_replacement=0): """Get all the feature and trig columns"""