D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/src/Analyzer/interpretation.py

@@ -0,0 +1,232 @@
+"""Row- and index-aligned ISA interpretation results."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Mapping, Sequence
+
+import numpy as np
+import torch
+from rdkit import Chem
+
+
+@dataclass(frozen=True)
+class ISAAnalysisRow:
+    """Scores and features aligned to one molecule's atoms or fragments."""
+
+    row_index: int
+    inputs: Mapping[str, Any]
+    atom_count: int
+    fragment_atom_indices: tuple[tuple[int, ...], ...]
+    scores: Mapping[str, np.ndarray]
+    score_mode: str
+    features: Mapping[str, np.ndarray]
+    feature_mode: str | None
+    prediction: np.ndarray | None = None
+
+    def atom_scores(self, key="positive") -> np.ndarray:
+        """Return a score for every atom, expanding fragment scores if needed."""
+
+        if key not in self.scores:
+            raise KeyError(f"Score {key!r} is unavailable. Available scores: {list(self.scores)}.")
+        values = np.asarray(self.scores[key])
+        if self.score_mode == "atom":
+            if values.shape[0] != self.atom_count:
+                raise ValueError(
+                    f"Atom score {key!r} has {values.shape[0]} rows, expected {self.atom_count}."
+                )
+            return values
+        output_shape = (self.atom_count,) + values.shape[1:]
+        expanded = np.zeros(output_shape, dtype=values.dtype)
+        for fragment_index, atoms in enumerate(self.fragment_atom_indices):
+            expanded[list(atoms)] = values[fragment_index]
+        return expanded
+
+
+@dataclass(frozen=True)
+class ISAAnalysisResult(Sequence[ISAAnalysisRow]):
+    """Valid ISA analysis rows in original input order."""
+
+    rows: tuple[ISAAnalysisRow, ...]
+    invalid_rows: tuple[Any, ...] = ()
+
+    def __len__(self):
+        return len(self.rows)
+
+    def __getitem__(self, index):
+        return self.rows[index]
+
+
+class ISAInterpreter:
+    """Execute score/feature extraction and validate all index boundaries."""
+
+    def __init__(
+        self,
+        core,
+        *,
+        score_keys=("positive",),
+        feature_keys=(),
+        score_mode=None,
+    ):
+        self.core = core
+        self.score_keys = tuple(score_keys)
+        self.feature_keys = tuple(feature_keys)
+        self.score_mode = score_mode
+        sculptor = getattr(getattr(core.dm, "gg", None), "sculptor", None)
+        if sculptor is None:
+            raise ValueError(
+                f"Model {str(core.artifacts.root)!r} does not expose the ISA sculptor "
+                "required for fragment-aligned interpretation."
+            )
+        self.sculptor = sculptor
+
+    def _fragments(self, smiles):
+        molecule = Chem.MolFromSmiles(smiles)
+        if molecule is None:
+            raise ValueError(f"Could not parse SMILES {smiles!r} for ISA fragmentation.")
+        fragments = self.sculptor.fragmentation_with_condition(
+            molecule, draw=False, get_index=False
+        )
+        indices = tuple(tuple(int(index) for index in fragment.atoms) for fragment in fragments)
+        flattened = sorted(index for fragment in indices for index in fragment)
+        expected = list(range(molecule.GetNumAtoms()))
+        if flattened != expected:
+            raise ValueError(
+                f"ISA fragments for {smiles!r} do not cover each atom exactly once. "
+                f"Observed atom indices: {flattened}; expected: {expected}."
+            )
+        return molecule.GetNumAtoms(), indices
+
+    @staticmethod
+    def _numpy_mapping(value, kind):
+        if not isinstance(value, dict):
+            raise ValueError(
+                f"ISA {kind} extraction must return a dictionary, got {type(value).__name__}. "
+                "This saved network does not implement the requested interpretation contract."
+            )
+        return {
+            key: item.detach().cpu().numpy() if isinstance(item, torch.Tensor) else np.asarray(item)
+            for key, item in value.items()
+        }
+
+    @staticmethod
+    def _split(mapping, counts, *, skip=()):
+        total = sum(counts)
+        result = [dict() for _ in counts]
+        for key, values in mapping.items():
+            if key in skip:
+                continue
+            array = np.asarray(values)
+            if array.shape[0] != total:
+                raise ValueError(
+                    f"ISA output {key!r} has leading dimension {array.shape[0]}, "
+                    f"but aligned rows require {total} ({counts})."
+                )
+            start = 0
+            for row, count in zip(result, counts):
+                row[key] = array[start:start + count]
+                start += count
+        return result
+
+    def analyze(self, *args, include_features=True, **kwargs):
+        normalized = self.core.normalize_inputs(args, kwargs)
+        loader, valid_indices = self.core._prepare(normalized)
+        if not valid_indices:
+            prediction = self.core._result_from_scores(
+                normalized, valid_indices, np.empty((0, len(self.core.targets)))
+            )
+            return ISAAnalysisResult((), prediction.invalid_rows)
+
+        draw_column = self.core.molecule_columns[0]
+        fragment_rows = []
+        for row_index in valid_indices:
+            fragment_rows.append(self._fragments(normalized[draw_column][row_index]))
+        atom_counts = [item[0] for item in fragment_rows]
+        fragment_counts = [len(item[1]) for item in fragment_rows]
+
+        score_output = self._numpy_mapping(
+            self.core.tm.get_score(self.core.nm, loader), "score"
+        )
+        prediction_values = score_output.pop("prediction", None)
+        network_name = str(self.core.config.get("network", ""))
+        missing_scores = [key for key in self.score_keys if key not in score_output]
+        if missing_scores:
+            raise ValueError(
+                f"Saved network {network_name!r} did not return required ISA scores "
+                f"{missing_scores!r}. Returned keys: {list(score_output)!r}."
+            )
+        score_output = {key: score_output[key] for key in self.score_keys}
+        score_mode = self.score_mode or (
+            "fragment"
+            if network_name in {"GC_model", "ISATPN_model", "ISATPM_model"}
+            else "atom"
+        )
+        score_counts = fragment_counts if score_mode == "fragment" else atom_counts
+        score_rows = self._split(score_output, score_counts)
+
+        feature_rows = [dict() for _ in valid_indices]
+        feature_mode = None
+        if include_features:
+            if not self.feature_keys:
+                raise ValueError(
+                    f"Saved network {network_name!r} does not provide ISATPN hidden features. "
+                    "Call analyze_rows(..., include_features=False) for score-only analysis."
+                )
+            feature_output = self._numpy_mapping(
+                self.core.tm.get_feature(self.core.nm, loader), "feature"
+            )
+            # ISATPN's get_feature branch returns dot/group-node hidden states.
+            missing_features = [
+                key for key in self.feature_keys if key not in feature_output
+            ]
+            if missing_features:
+                raise ValueError(
+                    f"Saved network {network_name!r} did not return required ISATPN "
+                    f"features {missing_features!r}. Returned keys: "
+                    f"{list(feature_output)!r}."
+                )
+            feature_values = {
+                key: feature_output[key] for key in self.feature_keys
+            }
+            feature_rows = self._split(feature_values, fragment_counts)
+            feature_mode = "fragment"
+
+        if prediction_values is not None:
+            prediction_values = self.core._inverse_transform(prediction_values)
+            if prediction_values.shape[0] != len(valid_indices):
+                raise ValueError(
+                    f"ISA prediction output has {prediction_values.shape[0]} rows, "
+                    f"expected {len(valid_indices)}."
+                )
+
+        rows = []
+        for position, row_index in enumerate(valid_indices):
+            rows.append(
+                ISAAnalysisRow(
+                    row_index=row_index,
+                    inputs={
+                        column: normalized[column][row_index]
+                        for column in self.core.input_columns
+                    },
+                    atom_count=atom_counts[position],
+                    fragment_atom_indices=fragment_rows[position][1],
+                    scores=score_rows[position],
+                    score_mode=score_mode,
+                    features=feature_rows[position],
+                    feature_mode=feature_mode,
+                    prediction=(
+                        None if prediction_values is None
+                        else np.asarray(prediction_values[position])
+                    ),
+                )
+            )
+        prediction_result = self.core._result_from_scores(
+            normalized,
+            valid_indices,
+            (
+                prediction_values
+                if prediction_values is not None
+                else np.zeros((len(valid_indices), len(self.core.targets)))
+            ),
+        )
+        return ISAAnalysisResult(tuple(rows), prediction_result.invalid_rows)

D4CMPP2/src/Analyzer/results.py

@@ -0,0 +1,101 @@
+"""Structured, row-preserving results for model inference."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Iterable, Iterator, Mapping, Sequence
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class PredictionRow:
+    """One input row and its prediction status."""
+
+    row_index: Any
+    inputs: Mapping[str, Any]
+    prediction: np.ndarray | None = None
+    status: str = "ok"
+    error: str | None = None
+
+
+@dataclass(frozen=True)
+class PredictionResult(Sequence[PredictionRow]):
+    """Prediction rows in the same order as the caller's input."""
+
+    rows: tuple[PredictionRow, ...]
+    targets: tuple[str, ...]
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, index):
+        return self.rows[index]
+
+    def __iter__(self) -> Iterator[PredictionRow]:
+        return iter(self.rows)
+
+    @property
+    def valid_rows(self) -> tuple[PredictionRow, ...]:
+        return tuple(row for row in self.rows if row.status == "ok")
+
+    @property
+    def invalid_rows(self) -> tuple[PredictionRow, ...]:
+        return tuple(row for row in self.rows if row.status != "ok")
+
+    def legacy_dict(self, input_columns: Iterable[str]) -> dict[tuple[Any, ...], np.ndarray]:
+        """Return the historical tuple-keyed mapping.
+
+        Duplicate input tuples cannot be represented by the historical mapping. Callers
+        that need every row should use this structured result or ``to_dataframe()``.
+        """
+
+        columns = tuple(input_columns)
+        result = {}
+        for row in self.valid_rows:
+            result[tuple(row.inputs[column] for column in columns)] = row.prediction
+        return result
+
+    def to_dataframe(self):
+        """Convert all rows, including invalid rows, to a pandas DataFrame."""
+
+        import pandas as pd
+
+        records = []
+        for row in self.rows:
+            record = {"row_index": row.row_index, **dict(row.inputs)}
+            if row.prediction is None:
+                for target in self.targets:
+                    record[f"{target}_pred"] = np.nan
+            else:
+                values = np.asarray(row.prediction).reshape(-1)
+                for index, target in enumerate(self.targets):
+                    record[f"{target}_pred"] = values[index]
+            record["prediction_status"] = row.status
+            record["prediction_error"] = row.error
+            records.append(record)
+        return pd.DataFrame.from_records(records)
+
+    def to_csv(self, path, **kwargs) -> str:
+        """Write all rows atomically and return the resolved output path."""
+
+        destination = Path(path).expanduser().resolve()
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        temporary = destination.with_name(destination.name + ".tmp")
+        self.to_dataframe().to_csv(temporary, index=False, **kwargs)
+        temporary.replace(destination)
+        return str(destination)
+
+
+@dataclass(frozen=True)
+class UncertaintyResult:
+    """Repeated stochastic predictions for one row-preserving input."""
+
+    mean: PredictionResult
+    std: PredictionResult
+    samples: tuple[PredictionResult, ...]
+    method: str
+    seed: int | None
+

D4CMPP2/src/DataManager/Dataset/GraphDataset.py

@@ -0,0 +1,314 @@
+import torch
+from torch.utils.data import Dataset
+from torch_geometric.data import Batch
+import numpy as np
+from ..contracts import (
+    GENERAL_BATCH_CONTRACT,
+    LEGACY_GENERAL_BATCH_CONTRACT,
+    LEGACY_SOLVENT_BATCH_CONTRACT,
+)
+
+class GraphDataset_legacy(Dataset):
+    batch_contract = LEGACY_GENERAL_BATCH_CONTRACT
+    def __init__(self, graphs=None, target=None, smiles=None):
+        if graphs is None: return
+        self.graphs = graphs
+        self.node_feature = [g.x for g in graphs]
+        self.edge_feature = [g.edge_attr for g in graphs]
+        self.target = torch.tensor(target).float()
+        self.smiles = smiles
+
+    def __len__(self):
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        return self.graphs[idx], self.node_feature[idx], self.edge_feature[idx], self.target[idx], self.smiles[idx]
+    
+    # reload the dataset with new data
+    def reload(self, data):
+        self.graphs, self.node_feature, self.edge_feature, self.target, self.smiles = data
+
+        self.target= torch.stack(self.target,dim=0)
+        if self.target.dim() == 1:
+            self.target = self.target.unsqueeze(-1)
+
+    # get a subset of the dataset with given indices
+    def get_subDataset(self, idx):
+        graphs = [self.graphs[i] for i in idx]
+        node_feature = [self.node_feature[i] for i in idx]
+        edge_feature = [self.edge_feature[i] for i in idx]
+        target = [self.target[i] for i in idx]
+        smiles = [self.smiles[i] for i in idx]
+        
+        dataset = GraphDataset_legacy()
+        dataset.reload((graphs, node_feature, edge_feature, target, smiles))
+        return dataset
+        
+
+    @staticmethod
+    def collate(samples):
+        graphs, node_feature, edge_feature, target, smiles = map(list, zip(*samples))
+        batched_graph = Batch.from_data_list(graphs)
+        return batched_graph, torch.concat(node_feature,dim=0), torch.concat(edge_feature,dim=0), torch.stack(target,dim=0), smiles
+    
+    @staticmethod
+    def unwrapper(batch_graph, node_feature, edge_feature, target, smiles,device='cpu'):
+        batch_graph = batch_graph.to(device=device)
+        node_feature = node_feature.float().to(device=device)
+        edge_feature = edge_feature.float().to(device=device)
+        target = target.float().to(device=device)
+        return {"graph":batch_graph, "node_feats":node_feature, "edge_feats":edge_feature, "target":target, "smiles":smiles}
+
+class GraphDataset_withSolv(GraphDataset_legacy):
+    batch_contract = LEGACY_SOLVENT_BATCH_CONTRACT
+    def __init__(self, graphs=None, solv_graphs=None, target=None, smiles=None, solv_smiles=None):
+        super().__init__(graphs, target, smiles)
+        if graphs is None: return
+        self.solv_graphs = solv_graphs
+        self.solv_node_feature = [g.x for g in solv_graphs]
+        self.solv_edge_feature = [g.edge_attr for g in solv_graphs]
+        self.solv_smiles = solv_smiles
+
+    def __getitem__(self, idx):
+        args = super().__getitem__(idx)
+        return args + (self.solv_graphs[idx], self.solv_node_feature[idx], self.solv_edge_feature[idx], self.solv_smiles[idx])
+    
+    def reload(self, data):
+        self.graphs, self.node_feature, self.edge_feature, self.target, self.smiles, self.solv_graphs, self.solv_node_feature, self.solv_edge_feature, self.solv_smiles = data
+        
+        self.target= torch.stack(self.target,dim=0)
+        if self.target.dim() == 1:
+            self.target = self.target.unsqueeze(-1)
+
+
+    # def subDataset(self, idx):
+    #     super().subDataset(idx)
+    #     self.solv_graphs = [self.solv_graphs[i] for i in idx]
+    #     self.solv_node_feature = [self.solv_node_feature[i] for i in idx]
+    #     self.solv_edge_feature = [self.solv_edge_feature[i] for i in idx]
+    #     self.solv_smiles = [self.solv_smiles[i] for i in idx]
+
+    def get_subDataset(self, idx):
+        graphs = [self.graphs[i] for i in idx]
+        node_feature = [self.node_feature[i] for i in idx]
+        edge_feature = [self.edge_feature[i] for i in idx]
+        target = [self.target[i] for i in idx]
+        smiles = [self.smiles[i] for i in idx]        
+        
+        solv_graphs = [self.solv_graphs[i] for i in idx]
+        solv_node_feature = [self.solv_node_feature[i] for i in idx]
+        solv_edge_feature = [self.solv_edge_feature[i] for i in idx]
+        solv_smiles = [self.solv_smiles[i] for i in idx]
+
+        dataset = GraphDataset_withSolv()
+        dataset.reload((graphs, node_feature, edge_feature, target, smiles, solv_graphs, solv_node_feature, solv_edge_feature, solv_smiles))
+        return dataset
+
+
+    @staticmethod
+    def collate(samples):
+        graphs, node_feature, edge_feature, target, smiles, solv_graphs, solv_node_feature, solv_edge_feature, solv_smiles = map(list, zip(*samples))
+
+        batched_graph = Batch.from_data_list(graphs)
+        batched_solv_graph = Batch.from_data_list(solv_graphs)
+        return batched_graph, torch.concat(node_feature,dim=0), torch.concat(edge_feature,dim=0), batched_solv_graph, torch.concat(solv_node_feature,dim=0), torch.concat(solv_edge_feature,dim=0), torch.stack(target,dim=0), smiles, solv_smiles
+    
+    @staticmethod
+    def unwrapper(batch_graph, node_feature, edge_feature, batch_solv_graph, solv_node_feature, solv_edge_feature, target, smiles, solv_smiles, device='cpu'):
+        batch_graph = batch_graph.to(device=device)
+        node_feature = node_feature.float().to(device=device)
+        edge_feature = edge_feature.float().to(device=device)
+        batch_solv_graph = batch_solv_graph.to(device=device)
+        solv_node_feature = solv_node_feature.float().to(device=device)
+        solv_edge_feature = solv_edge_feature.float().to(device=device)
+        target = target.float().to(device=device)
+
+        return {"graph":batch_graph, "node_feats":node_feature, "edge_feats":edge_feature, "solv_graph":batch_solv_graph, "solv_node_feats":solv_node_feature, "solv_edge_feats":solv_edge_feature, "target":target, "smiles":smiles, "solv_smiles":solv_smiles}
+    
+
+class GraphDataset(GraphDataset_legacy):
+    batch_contract = GENERAL_BATCH_CONTRACT
+    def __init__(self, graphs : dict = None, numeric_inputs : dict = None, target : list = None, smiles : dict = None, row_indices=None):
+        if graphs is None: 
+            self.graphs = {}
+            self.target = None
+            self.smiles = {}
+            self.data_keys = []
+            return
+
+        for key in graphs:
+            setattr(self, key + '_graphs', graphs[key])
+            setattr(self, key + '_node_feature', [g.x for g in graphs[key]])
+            setattr(self, key + '_edge_feature', [g.edge_attr for g in graphs[key]])
+
+        if numeric_inputs is not None:
+            for key in numeric_inputs:
+                setattr(self, key + '_var', numeric_inputs[key])
+        else:
+            numeric_inputs = {}
+
+        for key in smiles:
+            if key not in graphs:
+                raise ValueError(f"Key '{key}' in smiles is not found in graphs.")
+            setattr(self, key + '_smiles', smiles[key])
+
+        self.target = torch.as_tensor(target).float() if target is not None else None
+        self.original_row_index = (
+            torch.as_tensor(row_indices, dtype=torch.long) if row_indices is not None else None
+        )
+        self.data_keys = list(graphs.keys()) + list(numeric_inputs.keys()) + ['target']
+
+
+    def __len__(self):
+        if getattr(self, 'target', None) is not None:
+            return len(self.target)
+        else:
+            if len(self.data_keys) == 0:
+                return 0
+            key = self.data_keys[0]
+            if hasattr(self, key + '_graphs'):
+                return len(getattr(self, key + '_graphs'))
+            elif hasattr(self, key + '_var'):
+                return len(getattr(self, key + '_var'))
+            elif hasattr(self, key + '_smiles'):
+                return len(getattr(self, key + '_smiles'))
+            else:
+                return 0
+
+
+    def __getitem__(self, idx):
+        item = {}
+        for key in self.data_keys:
+            if hasattr(self, key + '_graphs'):
+                item[key + '_graphs'] = getattr(self, key + '_graphs')[idx]
+                if hasattr(self, key + '_node_feature'):
+                    nf = getattr(self, key + '_node_feature')[idx]
+                    if type(nf) is not torch.Tensor:
+                        nf = torch.tensor(nf, dtype=torch.float32)
+                    item[key + '_node_feature'] = nf
+                if hasattr(self, key + '_edge_feature'):
+                    ef = getattr(self, key + '_edge_feature')[idx]
+                    if type(ef) is not torch.Tensor:
+                        ef = torch.tensor(ef, dtype=torch.float32)
+                    item[key + '_edge_feature'] = ef
+                if hasattr(self, key + '_smiles'):
+                    item[key + '_smiles'] = getattr(self, key + '_smiles')[idx]
+            elif hasattr(self, key + '_var'):
+                var = getattr(self, key + '_var')[idx]
+                if torch.is_tensor(var):
+                    var = var.float()
+                else:
+                    var = torch.tensor(var, dtype=torch.float32)
+                item[key + '_var'] = var
+            elif hasattr(self, key + '_smiles'):
+                item[key + '_smiles'] = getattr(self, key + '_smiles')[idx]
+        if hasattr(self, 'target') and self.target is not None:
+            item['target'] = self.target[idx]
+        if self.original_row_index is not None:
+            item['original_row_index'] = self.original_row_index[idx]
+        return item
+
+    def reload(self, data):
+        if len(data) == 0:
+            self.target = None
+            self.data_keys = []
+            return
+
+        new_data_keys = []
+        for key in data[0]:
+            values = [d[key] for d in data]
+            if key == 'target':
+                if len(values) == 0:
+                    setattr(self, key, None)
+                elif torch.is_tensor(values[0]):
+                    target = torch.stack(values, dim=0).float()
+                    if target.dim() == 1:
+                        target = target.unsqueeze(-1)
+                    setattr(self, key, target)
+                else:
+                    target = torch.tensor(values, dtype=torch.float32)
+                    if target.dim() == 1:
+                        target = target.unsqueeze(-1)
+                    setattr(self, key, target)
+            elif key == 'original_row_index':
+                setattr(self, key, torch.as_tensor(values, dtype=torch.long))
+            elif key.endswith('_var'):
+                if len(values) == 0:
+                    setattr(self, key, torch.empty((0,), dtype=torch.float32))
+                elif torch.is_tensor(values[0]):
+                    setattr(self, key, torch.stack(values, dim=0).float())
+                else:
+                    setattr(self, key, torch.tensor(values, dtype=torch.float32))
+            else:
+                setattr(self, key, values)
+            if key.endswith('_graphs'):
+                new_data_keys.append(key[:-7])  # Remove '_graphs'
+            elif key.endswith('_var'):
+                new_data_keys.append(key[:-4])  # Remove '_var'
+            elif key == 'target':
+                new_data_keys.append(key)
+        self.data_keys = new_data_keys
+
+
+    def subDataset(self, idx):
+        for key in self.data_keys:
+            if hasattr(self, key + '_graphs'):
+                setattr(self, key + '_graphs', [getattr(self, key + '_graphs')[i] for i in idx])
+                setattr(self, key + '_node_feature', [getattr(self, key + '_node_feature')[i] for i in idx])
+                setattr(self, key + '_edge_feature', [getattr(self, key + '_edge_feature')[i] for i in idx])
+            elif hasattr(self, key + '_var'):   
+                var = getattr(self, key + '_var')
+                if torch.is_tensor(var):
+                    setattr(self, key + '_var', var[torch.tensor(idx, dtype=torch.long)])
+                else:
+                    setattr(self, key + '_var', var[np.array(idx, dtype=int)])
+            elif hasattr(self, key + '_smiles'):
+                setattr(self, key + '_smiles', [getattr(self, key + '_smiles')[i] for i in idx])
+        if self.target is not None:
+            self.target = self.target[np.array(idx, dtype=int)]
+        if self.original_row_index is not None:
+            self.original_row_index = self.original_row_index[torch.as_tensor(idx, dtype=torch.long)]
+
+    def get_subDataset(self, idx):
+        dataset = GraphDataset()
+        dataset.reload([self[i] for i in idx])
+        return dataset
+
+
+    @staticmethod
+    def collate(samples):
+        batched_data = {}
+        for key in samples[0].keys():
+            if key.endswith('_graphs'):
+                batched_data[key] = Batch.from_data_list([s[key] for s in samples])
+            elif key.endswith('_node_feature'):
+                batched_data[key] = torch.concat([s[key] for s in samples], dim=0)
+            elif key.endswith('_edge_feature'):
+                batched_data[key] = torch.concat([s[key] for s in samples], dim=0)
+            elif key.endswith('_var'):
+                batched_data[key] = torch.stack([s[key] for s in samples], dim=0)
+            elif key.endswith('_smiles'):
+                batched_data[key] = [s[key] for s in samples]
+            elif key == 'target':
+                batched_data[key] = torch.stack([s[key].reshape(-1) for s in samples], dim=0)
+            elif key == 'original_row_index':
+                batched_data[key] = torch.stack([s[key].reshape(()) for s in samples], dim=0)
+        return batched_data
+
+    @staticmethod
+    def unwrapper(device='cpu', **batched_data):
+        for key in batched_data.keys():
+            if key.endswith('_graphs'):
+                batched_data[key] = batched_data[key].to(device=device)
+            elif key.endswith('_node_feature') or key.endswith('_edge_feature'):
+                batched_data[key] = batched_data[key].float().to(device=device)
+            elif key.endswith('_var'):
+                batched_data[key] = batched_data[key].float().to(device=device)
+            elif key == 'target':
+                batched_data[key] = batched_data[key].float().to(device=device)
+            elif key == 'original_row_index':
+                batched_data[key] = batched_data[key].long().to(device=device)
+        return batched_data
+
+    
+