Spaces:
Running
Running
| """TabICLv2 classifiers for the twelve Tox21 endpoints.""" | |
| from __future__ import annotations | |
| import gzip | |
| import os | |
| import pickle | |
| import shutil | |
| from typing import Any | |
| import numpy as np | |
| from tabicl import TabICLClassifier | |
| from .utils import TASKS | |
| CHECKPOINT_NAME = "model.tabicl.pkl.gz" | |
| class Tox21TabICL: | |
| """A separate TabICLv2 in-context classifier for each Tox21 endpoint.""" | |
| def __init__( | |
| self, | |
| seed: int = 42, | |
| device: str | None = None, | |
| model_config: dict[str, Any] | None = None, | |
| ): | |
| self.tasks = TASKS | |
| self.device = device | |
| self.model_config = model_config or {} | |
| self.model = { | |
| task: TabICLClassifier( | |
| random_state=seed, | |
| device=device, | |
| **self.model_config, | |
| ) | |
| for task in self.tasks | |
| } | |
| self._shared_model = None | |
| def _share_model_weights(self, estimator: TabICLClassifier) -> None: | |
| """Keep one immutable TabICLv2 network shared by all task estimators.""" | |
| if self._shared_model is None: | |
| self._shared_model = estimator.model_ | |
| else: | |
| estimator.model_ = self._shared_model | |
| def load_model(self, path: str) -> None: | |
| """Load all fitted task estimators from ``path``.""" | |
| for task in self.tasks: | |
| model_path = os.path.join(path, task, CHECKPOINT_NAME) | |
| with gzip.open(model_path, "rb") as checkpoint_file: | |
| estimator = pickle.load(checkpoint_file) | |
| if self.device is not None: | |
| estimator.device = self.device | |
| estimator._resolve_device() | |
| estimator.model_.to(estimator.device_) | |
| estimator._build_inference_config() | |
| self._share_model_weights(estimator) | |
| self.model[task] = estimator | |
| def save_model(self, path: str) -> None: | |
| """Save fitted state while reusing the public TabICLv2 base checkpoint.""" | |
| for task in self.tasks: | |
| model_path = os.path.join(path, task, CHECKPOINT_NAME) | |
| uncompressed_path = model_path.removesuffix(".gz") | |
| estimator = self.model[task] | |
| # Persist with automatic device selection so checkpoints remain portable. | |
| configured_device = estimator.device | |
| estimator.device = None | |
| try: | |
| estimator.save( | |
| uncompressed_path, | |
| save_model_weights=False, | |
| save_training_data=True, | |
| save_kv_cache=False, | |
| ) | |
| with open(uncompressed_path, "rb") as source: | |
| with gzip.open(model_path, "wb", compresslevel=1) as destination: | |
| shutil.copyfileobj(source, destination) | |
| os.remove(uncompressed_path) | |
| finally: | |
| estimator.device = configured_device | |
| def fit(self, task: str, input_features: np.ndarray, labels: np.ndarray) -> None: | |
| """Fit the in-context state for one Tox21 endpoint.""" | |
| if task not in self.tasks: | |
| raise ValueError(f"Unknown task: {task}") | |
| if labels.ndim != 1: | |
| raise ValueError("Function only accepts one-dimensional labels.") | |
| estimator = self.model[task].fit(input_features, labels) | |
| self._share_model_weights(estimator) | |
| def predict(self, task: str, features: np.ndarray) -> np.ndarray: | |
| """Return the probability of the positive class for one endpoint.""" | |
| if task not in self.tasks: | |
| raise ValueError(f"Unknown task: {task}") | |
| if features.ndim != 2: | |
| raise ValueError( | |
| f"Function expects a two-dimensional array; got {features.shape}." | |
| ) | |
| return self.model[task].predict_proba(features)[:, 1] | |