torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/base/old_predictor.py

@@ -0,0 +1,904 @@
+import os
+import json
+import datetime
+import tempfile
+import warnings
+import torch
+import numpy as np
+from rdkit import Chem
+from ..utils import (
+    roc_auc_score,
+    accuracy_score,
+    mean_absolute_error,
+    mean_squared_error,
+    root_mean_squared_error,
+    r2_score,
+)
+from ..utils.format import sanitize_config
+
+from dataclasses import dataclass, field
+from abc import ABC, abstractmethod
+from typing import Optional, ClassVar, Union, List, Dict, Any, Tuple, Callable, Type
+
+
+@dataclass
+class BaseMolecularPredictor(ABC):
+    """Base class for molecular discovery estimators.
+
+    This class provides a skeleton for implementing custom molecular discovery
+    models following scikit-learn's conventions. It includes basic functionality
+    for parameter management, fitting, prediction, and validation.
+
+    Attributes
+    ----------
+    num_task : int
+        Number of prediction tasks
+    task_type : str
+        Type of task ('classification' or 'regression')
+    model_name : str
+        Name of the model
+    device : Optional[torch.device]
+        Device to run the model on
+    model : Optional[torch.nn.Module]
+        Model instance
+    """
+    
+    num_task: int
+    task_type: str
+    model_name: str = field(default="BaseMolecularPredictor")
+    device: Optional[torch.device] = field(default=None)
+    model: Optional[torch.nn.Module] = field(default=None, init=False)
+    
+    # Class variables
+    DEFAULT_METRICS: ClassVar[Dict] = {
+        "classification": {
+            "default": ("roc_auc", roc_auc_score, True),  # (name, function, higher_better)
+        },
+        "regression": {"default": ("mae", mean_absolute_error, False)},
+    }
+    
+    # Instance variables with default values
+    is_fitted_: bool = field(default=False, init=False)
+    model_class: Optional[Type[torch.nn.Module]] = field(default=None, init=False)
+    
+    def __post_init__(self):
+        """Validate inputs and set device after initialization."""
+        if self.task_type not in ["classification", "regression"]:
+            raise ValueError(
+                f"task_type must be one of ['classification', 'regression'], got {self.task_type}"
+            )
+
+        if not isinstance(self.num_task, (int, np.integer)) or self.num_task <= 0:
+            raise ValueError(f"num_task must be a positive integer, got {self.num_task}")
+            
+        # Set device if not provided
+        if self.device is None:
+            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        elif isinstance(self.device, str):
+            self.device = torch.device(self.device)
+    
+    def get_params(self, deep: bool = True) -> Dict[str, Any]:
+        """Get parameters for this estimator.
+
+        Parameters
+        ----------
+        deep : bool, default=True
+            If True, will return the parameters for this estimator and
+            contained subobjects.
+
+        Returns
+        -------
+        params : dict
+            Parameter names mapped to their values.
+        """
+        out = dict()
+        for key in self._get_param_names():
+            value = getattr(self, key)
+            if deep and hasattr(value, "get_params"):
+                deep_items = value.get_params().items()
+                out.update((key + "__" + k, val) for k, val in deep_items)
+            out[key] = value
+        return out
+
+    def set_params(self, **params) -> "BaseMolecularPredictor":
+        """Set the parameters of this predictor.
+
+        Parameters
+        ----------
+        **params : dict
+            predictor parameters.
+
+        Returns
+        -------
+        self : object
+            predictor instance.
+        """
+        if not params:
+            return self
+
+        valid_params = self.get_params(deep=True)
+        for key, value in params.items():
+            if key not in valid_params:
+                raise ValueError(f"Invalid parameter {key} for predictor {self}")
+            setattr(self, key, value)
+        return self
+
+    def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+        """Setup optimization components including optimizer and learning rate scheduler.
+
+        Returns
+        -------
+        Tuple[optim.Optimizer, Optional[Any]]
+            A tuple containing:
+            - The configured optimizer
+            - The learning rate scheduler (if enabled, else None)
+        """
+        raise NotImplementedError("optimizers and schedulers must be implemented by child class")
+
+    def _setup_evaluation(
+        self,
+        evaluate_criterion: Optional[Union[str, Callable]],
+        evaluate_higher_better: Optional[bool],
+    ) -> None:
+        """Setup evaluation metric and direction.
+
+        Parameters
+        ----------
+        evaluate_criterion : Union[str, Callable, None]
+            Evaluation metric specification
+        evaluate_higher_better : bool, optional
+            Whether higher values are better
+        """
+        if evaluate_criterion is None:
+            # Use default metric based on task type
+            default_metric = self.DEFAULT_METRICS[self.task_type]["default"]
+            self.evaluate_name = default_metric[0]
+            self.evaluate_criterion = default_metric[1]
+            self.evaluate_higher_better = default_metric[2]
+
+        else:
+            if isinstance(evaluate_criterion, str):
+                # Map string to predefined metric
+                metric_map = {
+                    "accuracy": (accuracy_score, True),
+                    "roc_auc": (roc_auc_score, True),
+                    "rmse": (root_mean_squared_error, False),
+                    "mse": (mean_squared_error, False),
+                    "mae": (mean_absolute_error, False),
+                    "r2": (r2_score, True),
+                }
+
+                if evaluate_criterion not in metric_map:
+                    raise ValueError(
+                        f"Unknown metric: {evaluate_criterion}. "
+                        f"Available metrics: {list(metric_map.keys())}"
+                    )
+
+                self.evaluate_name = evaluate_criterion
+                self.evaluate_criterion = metric_map[evaluate_criterion][0]
+                self.evaluate_higher_better = (
+                    metric_map[evaluate_criterion][1]
+                    if evaluate_higher_better is None
+                    else evaluate_higher_better
+                )
+
+            else:
+                # Custom metric function
+                if evaluate_higher_better is None:
+                    raise ValueError(
+                        "evaluate_higher_better must be specified "
+                        "when using a custom evaluation function"
+                    )
+
+                self.evaluate_name = "custom"
+                self.evaluate_criterion = evaluate_criterion
+                self.evaluate_higher_better = evaluate_higher_better
+
+    @abstractmethod
+    def fit(self, X_train, y_train, X_val=None, y_val=None, X_unlbl=None) -> "BaseMolecularPredictor":
+        """Fit the model to the data.
+
+        Parameters
+        ----------
+        X_train / X_val / X_unlbl: list of SMILES
+        y_train / y_val: array-like of shape (n_samples,), optional
+            Target values.
+        Returns
+        -------
+        self : object
+            Fitted estimator.
+        """
+        pass
+
+    @abstractmethod
+    def autofit(
+        self,
+        X_train,
+        y_train,
+        X_val=None,
+        y_val=None,
+        search_parameters: Optional[dict] = None,
+        n_trials: int = 10,
+    ) -> "BaseMolecularPredictor":
+        """AutoFit the model to the data.
+
+        Parameters
+        ----------
+        X_train / X_val : list of SMILES
+        y_train / y_val: array-like of shape (n_samples,), optional
+            Target values.
+        Returns
+        -------
+        self : object
+            Fitted estimator.
+        """
+        pass
+
+    @abstractmethod
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        """Make predictions using the fitted model.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data.
+
+        Returns
+        -------
+        y_pred : ndarray of shape (n_samples,)
+            Predicted values.
+        """
+        pass
+
+    def _train_epoch(self, train_loader, optimizer):
+        """Training logic for one epoch.
+
+        Should be implemented by child classes based on specific model architecture.
+        """
+        raise NotImplementedError("Training epoch logic must be implemented by child class")
+
+    def _evaluation_epoch(self, evaluate_loader):
+        """Training logic for one evaluation.
+
+        Should be implemented by child classes based on specific model architecture.
+        """
+        raise NotImplementedError("Evaluation epoch logic must be implemented by child class")
+
+    def _check_is_fitted(self) -> None:
+        """Check if the estimator is fitted.
+
+        Raises
+        ------
+        NotFittedError
+            If the estimator is not fitted.
+        """
+        if not self.is_fitted_:
+            raise AttributeError(
+                "This MolecularBaseEstimator instance is not fitted yet. "
+                "Call 'fit' before using this estimator."
+            )
+
+    def __repr__(self, N_CHAR_MAX=700):
+        """Get string representation of the estimator.
+
+        Parameters
+        ----------
+        N_CHAR_MAX : int, default=700
+            Maximum number of non-blank characters to show
+
+        Returns
+        -------
+        str
+            String representation of the estimator
+        """
+        # Get all relevant attributes (excluding private ones and callables)
+        attributes = {
+            name: value
+            for name, value in sorted(self.__dict__.items())
+            if not name.startswith("_") and not callable(value)
+        }
+
+        def format_value(v):
+            """Format a value for representation."""
+            if isinstance(v, (float, np.float32, np.float64)):
+                return f"{v:.3g}"
+            elif isinstance(v, (list, tuple, np.ndarray)) and len(v) > 6:
+                return f"{str(v[:3])[:-1]}, ..., {str(v[-3:])[1:]}"
+            elif isinstance(v, str) and len(v) > 50:
+                return f"'{v[:25]}...{v[-22:]}'"
+            elif isinstance(v, dict) and len(v) > 6:
+                items = list(v.items())
+                return f"{dict(items[:3])}...{dict(items[-3:])}"
+            elif isinstance(v, torch.nn.Module):
+                return f"{v.__class__.__name__}(...)"
+            elif v is None:
+                return "None"
+            else:
+                return repr(v)
+
+        # Build the representation string
+        class_name = self.__class__.__name__
+        attributes_str = []
+
+        # First add important attributes that should appear first
+        important_attrs = ["num_task", "task_type", "model_name", "is_fitted_"]
+        for attr in important_attrs:
+            if attr in attributes:
+                value = attributes.pop(attr)
+                attributes_str.append(f"{attr}={format_value(value)}")
+
+        # Then add the rest
+        sorted_attrs = sorted(attributes.items())
+        attributes_str.extend(f"{name}={format_value(value)}" for name, value in sorted_attrs)
+
+        # Join all parts
+        content = ",\n    ".join(attributes_str)
+        repr_ = f"{class_name}(\n    {content}\n)"
+
+        # Handle length limit
+        if len(repr_) > N_CHAR_MAX:
+            lines = repr_.split("\n")
+
+            # Keep header
+            result = [lines[0]]
+            current_length = len(lines[0])
+
+            # Process middle lines
+            for line in lines[1:-1]:
+                if current_length + len(line) + 5 > N_CHAR_MAX:  # 5 for "...\n"
+                    result.append("    ...")
+                    break
+                result.append(line)
+                current_length += len(line) + 1  # +1 for newline
+
+            # Add closing parenthesis
+            result.append(lines[-1])
+            repr_ = "\n".join(result)
+
+        return repr_
+
+    @abstractmethod
+    def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+        """Get model parameters either from checkpoint or current instance.
+
+        Parameters
+        ----------
+        checkpoint : Optional[Dict]
+            Checkpoint containing model hyperparameters
+
+        Returns
+        -------
+        Dict[str, Any]
+            Dictionary of model parameters
+
+        Raises
+        ------
+        NotImplementedError
+            If the child class doesn't implement this method
+        ValueError
+            If checkpoint contains invalid parameters
+        """
+        raise NotImplementedError(
+            "Method _get_model_params must be implemented by child class. "
+            "This method should return a dictionary of model parameters either "
+            "from a checkpoint or current instance values."
+        )
+
+    def _initialize_model(
+        self, model_class: Type[torch.nn.Module], device, checkpoint: Optional[Dict] = None
+    ) -> None:
+        """Initialize the GNN model with given parameters or checkpoint.
+
+        Parameters
+        ----------
+        checkpoint : Optional[Dict]
+            If provided, should contain:
+                - hyperparameters: Dict of model configuration
+                - model_state_dict: Dict of model state
+            If None, initializes model with current instance parameters.
+
+        Raises
+        ------
+        ValueError
+            If checkpoint is provided but missing required keys or contains invalid parameters
+        RuntimeError
+            If device initialization fails
+        """
+        try:
+            model_params = self._get_model_params(checkpoint)
+
+            # Initialize model with parameters
+            self.model = model_class(**model_params)
+
+            # Move model to device
+            try:
+                self.model = self.model.to(device)
+            except Exception as e:
+                raise RuntimeError(f"Failed to move model to device {device}: {str(e)}")
+
+            # Load state dict if provided in checkpoint
+            if checkpoint is not None:
+                try:
+                    self.model.load_state_dict(checkpoint["model_state_dict"])
+                except Exception as e:
+                    raise ValueError(f"Failed to load model state dictionary: {str(e)}")
+
+        except Exception as e:
+            raise RuntimeError(f"Model initialization failed: {str(e)}")
+
+    def _load_default_criterion(self):
+        """Sets default criterion based on task_type."""
+        if self.task_type == "regression":
+            return torch.nn.L1Loss()
+        elif self.task_type == "classification":
+            return torch.nn.BCEWithLogitsLoss()
+        else:
+            warnings.warn(
+                "Unknown task type. Using MSE Loss as the default criterion. "
+                "Please specify a valid task type ('regression' or 'classification') or provide a criterion."
+            )
+            return torch.nn.L1Loss()
+
+    def _validate_inputs(
+        self, X: List[str], y: Optional[Union[List, np.ndarray]] = None, return_rdkit_mol: bool = True
+    ) -> Tuple[Union[List[str], List["Mol"]], Optional[np.ndarray]]:
+        """Validate input SMILES strings and target values.
+
+        Parameters
+        ----------
+        X : List[str]
+            List of SMILES strings
+        y : Union[List, np.ndarray], optional
+            Target values. Can be:
+            - 1D array/list for single task (num_task must be 1)
+            - 2D array/list for multiple tasks
+        return_rdkit_mol : bool, default=False
+            If True, returns RDKit Mol objects instead of SMILES strings
+
+        Returns
+        -------
+        X : Union[List[str], List[Mol]]
+            Validated SMILES strings or RDKit Mol objects if return_rdkit_mol=True
+        y : np.ndarray or None
+            Validated target values as numpy array, reshaped to 2D if needed
+
+        Raises
+        ------
+        ValueError
+            If inputs don't meet requirements
+        """
+        # Validate X
+        if not isinstance(X, list):
+            raise ValueError("X must be a list of SMILES strings")
+
+        if not all(isinstance(s, str) for s in X):
+            raise ValueError("All elements in X must be strings")
+
+        # Validate SMILES using RDKit and store Mol objects
+        invalid_smiles = []
+        rdkit_mols = []
+        for i, smiles in enumerate(X):
+            is_valid, error_msg, mol = self._validate_smiles(smiles, i)
+            if not is_valid:
+                invalid_smiles.append(error_msg)
+            else:
+                rdkit_mols.append(mol)
+
+        if invalid_smiles:
+            raise ValueError("Invalid SMILES found:\n" + "\n".join(invalid_smiles))
+
+        # Validate y if provided
+        if y is not None:
+            try:
+                # Convert to numpy array if it's a list
+                y = np.asarray(y, dtype=np.float32)
+            except Exception as e:
+                raise ValueError(f"Could not convert y to numpy array: {str(e)}")
+
+            # Handle 1D array case
+            if len(y.shape) == 1:
+                if self.num_task != 1:
+                    raise ValueError(
+                        f"1D target array provided but num_task is {self.num_task}. "
+                        "For multiple tasks, y must be 2D array."
+                    )
+                # Reshape to 2D
+                y = y.reshape(-1, 1)
+
+            # Validate dimensions
+            if len(y.shape) != 2:
+                raise ValueError(
+                    f"y must be 1D (for single task) or 2D (for multiple tasks), "
+                    f"got shape {y.shape}"
+                )
+
+            # Check sample dimension
+            if y.shape[0] != len(X):
+                raise ValueError(
+                    f"First dimension of y ({y.shape[0]}) must match " f"length of X ({len(X)})"
+                )
+
+            # Check task dimension
+            if y.shape[1] != self.num_task:
+                raise ValueError(
+                    f"Number of tasks in y ({y.shape[1]}) must match "
+                    f"num_task ({self.num_task})"
+                )
+
+            # Handle infinite values
+            inf_mask = ~np.isfinite(y)
+            if np.any(inf_mask):
+                inf_indices = np.where(inf_mask)
+                warnings.warn(
+                    f"Infinite values found in y at indices: {list(zip(*inf_indices))}. "
+                    "Converting to NaN.",
+                    RuntimeWarning,
+                )
+                y = y.astype(float)  # Ensure float type for NaN support
+                y[inf_mask] = np.nan
+
+            # Check for NaN values after conversion
+            nan_mask = np.isnan(y)
+            if np.any(nan_mask):
+                # Count NaNs per task
+                nan_counts = np.sum(nan_mask, axis=0)
+                nan_percentages = (nan_counts / len(X)) * 100
+
+                # Create detailed warning message
+                task_warnings = []
+                for task_idx, (count, percentage) in enumerate(zip(nan_counts, nan_percentages)):
+                    if count > 0:
+                        task_warnings.append(f"Task {task_idx}: {count} NaNs ({percentage:.1f}%)")
+
+                warnings.warn(
+                    "NaN values present in y:\n"
+                    + "\n".join(task_warnings)
+                    + "\nSamples with NaN values may be ignored during training.",
+                    RuntimeWarning,
+                )
+
+        return rdkit_mols if return_rdkit_mol else X, y
+
+    def _validate_smiles(self, smiles: str, idx: int) -> Tuple[bool, Optional[str], Optional["Mol"]]:
+        """Validate a single SMILES string using RDKit.
+
+        Parameters
+        ----------
+        smiles : str
+            SMILES string to validate
+        idx : int
+            Index of the SMILES in the input list (for error reporting)
+
+        Returns
+        -------
+        is_valid : bool
+            Whether the SMILES is valid
+        error_msg : str or None
+            Error message if invalid, None if valid
+        mol : Mol or None
+            RDKit Mol object if valid, None if invalid
+        """
+        if not smiles or not smiles.strip():
+            return False, f"Empty SMILES at index {idx}", None
+
+        try:
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is None:
+                return False, f"Invalid SMILES structure at index {idx}: {smiles}", None
+            return True, None, mol
+        except Exception as e:
+            return False, f"RDKit error at index {idx}: {str(e)}", None
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        """Get parameter names for the estimator."""
+        return ["num_task", "task_type", "is_fitted_"]
+
+    def save_model(self, path: str) -> None:
+        """Save the model to disk.
+
+        Parameters
+        ----------
+        path : str
+            Path where to save the model
+        """
+        if not self.is_fitted_:
+            raise ValueError("Model must be fitted before saving")
+
+        if not path.endswith((".pt", ".pth")):
+            raise ValueError("Save path should end with '.pt' or '.pth'")
+
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+
+        model_name = os.path.splitext(os.path.basename(path))[0]
+        save_dict = {
+            "model_state_dict": self.model.state_dict(),
+            "hyperparameters": self.get_params(),
+            "model_name": model_name,
+            "date_saved": datetime.datetime.now().isoformat(),
+            "version": getattr(self, "__version__", "1.0.0"),
+        }
+        torch.save(save_dict, path)
+
+    def load_model(self, path: str, repo_id: Optional[str] = None):
+        """Load a saved model from disk or download from HuggingFace hub.
+
+        Parameters
+        ----------
+        path : str
+            Path to the saved model file or desired local path for downloaded model
+        repo_id : str, optional
+            HuggingFace model repository ID (e.g., 'username/model-name')
+            If provided and local path doesn't exist, will attempt to download from hub
+
+        Returns
+        -------
+        self : GNNMolecularPredictor
+            Updated model instance with loaded weights and parameters
+
+        Raises
+        ------
+        FileNotFoundError
+            If the model file doesn't exist locally and no repo is provided,
+            or if download from repo fails
+        ValueError
+            If the saved file is corrupted or incompatible
+        """
+        # First try to load from local path
+        if os.path.exists(path):
+            try:
+                checkpoint = torch.load(path, map_location=self.device)
+            except Exception as e:
+                raise ValueError(f"Error loading model from {path}: {str(e)}")
+
+        # If local file doesn't exist and repo is provided, try downloading
+        elif repo_id is not None:
+            try:
+                from huggingface_hub import hf_hub_download
+
+                # Create directory if it doesn't exist
+                os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+
+                # Download from HuggingFace using specified parameters
+                model_name = os.path.splitext(os.path.basename(path))[0]
+
+                # Download model file
+                downloaded_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename=f"{model_name}.pt",
+                    local_dir=os.path.dirname(path),
+                    # local_dir_use_symlinks=False
+                )
+                hf_hub_download(
+                    repo_id=repo_id,
+                    filename="config.json",
+                    local_dir=os.path.dirname(path),
+                )
+                # Load the downloaded model
+                checkpoint = torch.load(downloaded_path, map_location=self.device)
+
+            except Exception as e:
+                if os.path.exists(path):
+                    os.remove(path)  # Clean up partial downloads
+                raise FileNotFoundError(
+                    f"Failed to download model from repository '{repo_id}': {str(e)}"
+                )
+
+        # If neither local file exists nor repo provided
+        else:
+            raise FileNotFoundError(f"No model file found at '{path}' and no repository provided")
+
+        try:
+            # Validate checkpoint contents
+            required_keys = {
+                "model_state_dict",
+                "hyperparameters",
+                "model_name",
+            }
+            if not all(key in checkpoint for key in required_keys):
+                missing_keys = required_keys - set(checkpoint.keys())
+                raise ValueError(f"Checkpoint missing required keys: {missing_keys}")
+
+            # Validate and track all hyperparameter comparisons
+            parameter_status = []
+            for key, new_value in checkpoint["hyperparameters"].items():
+                if hasattr(self, key):
+                    old_value = getattr(self, key)
+                    is_changed = old_value != new_value
+                    parameter_status.append(
+                        {
+                            "Parameter": key,
+                            "Old Value": old_value,
+                            "New Value": new_value,
+                            "Status": "Changed" if is_changed else "Unchanged",
+                        }
+                    )
+                    if is_changed:
+                        setattr(self, key, new_value)
+
+            # Print comprehensive parameter status table
+            if parameter_status:
+                print("\nHyperparameter Status:")
+                print("-" * 80)
+                print(f"{'Parameter':<20} {'Old Value':<20} {'New Value':<20} {'Status':<10}")
+                print("-" * 80)
+
+                # Sort by status (Changed first, then Unchanged)
+                parameter_status.sort(key=lambda x: (x["Status"] != "Changed", x["Parameter"]))
+
+                for param in parameter_status:
+                    status_color = (
+                        "\033[91m" if param["Status"] == "Changed" else "\033[92m"
+                    )  # Red for changed, green for unchanged
+                    print(
+                        f"{param['Parameter']:<20} "
+                        f"{str(param['Old Value']):<20} "
+                        f"{str(param['New Value']):<20} "
+                        f"{status_color}{param['Status']}\033[0m"
+                    )
+                print("-" * 80)
+
+                # Print summary
+                changes_count = sum(1 for p in parameter_status if p["Status"] == "Changed")
+                print(
+                    f"\nSummary: {changes_count} parameters changed, "
+                    f"{len(parameter_status) - changes_count} unchanged"
+                )
+
+            self._initialize_model(self.model_class, self.device, checkpoint)
+            self.model_name = checkpoint["model_name"]
+            self.is_fitted_ = True
+
+            # Move model to correct device
+            self.model = self.model.to(self.device)
+
+            print(f"Model successfully loaded from {'repository' if repo_id else 'local path'}")
+            return self
+
+        except Exception as e:
+            raise ValueError(f"Error loading model checkpoint: {str(e)}")
+
+        finally:
+            # Clean up any temporary files
+            if (
+                repo_id is not None
+                and os.path.exists(downloaded_path)
+                and os.path.abspath(downloaded_path) != os.path.abspath(path)
+            ):
+                try:
+                    os.remove(downloaded_path)
+                except Exception:
+                    pass
+
+    def push_to_huggingface(
+        self,
+        repo_id: str,
+        task_id: str = "default",
+        metrics: Optional[Dict[str, float]] = None,
+        commit_message: str = "Update model",
+        token: Optional[str] = None,
+        private: bool = False,
+    ) -> None:
+        """Push a task-specific model checkpoint to Hugging Face Hub.
+
+        Parameters
+        ----------
+        repo_id : str
+            The ID of the model repository (format: '<username>/<model_name>')
+        task_id : str
+            Task identifier (e.g., 'O2', 'N2')
+        metrics : Optional[Dict[str, float]]
+            Optional dictionary of metric names and values for this task
+        commit_message : str, default="Update model"
+            Message to commit with
+        token : Optional[str], default=None
+            HuggingFace token
+        private : bool, default=False
+            Whether to make the repository private
+        """
+        try:
+            from huggingface_hub import HfApi, create_repo, metadata_update
+            from .utils.hf_hub import merge_task_configs, get_existing_repo_data, create_model_card
+        except ImportError:
+            raise ImportError(
+                "huggingface_hub package is required to push to Hugging Face Hub. "
+                "Install it with: pip install huggingface_hub"
+            )
+
+        # Validate inputs
+        if not isinstance(repo_id, str) or "/" not in repo_id:
+            raise ValueError("repo_id must be in format '<username>/<model_name>'")
+        if not task_id:
+            raise ValueError("task_id must be provided")
+
+        self._check_is_fitted()
+
+        try:
+            # Initialize API and check existing repo
+            api = HfApi(token=token)
+            repo_exists, existing_config, existing_readme = get_existing_repo_data(repo_id, token)
+
+            # Create repo if it doesn't exist
+            create_repo(repo_id, private=private, token=token, exist_ok=True)
+
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                # Save task-specific model checkpoint
+                model_path = os.path.join(tmp_dir, f"{self.model_name}.pt")
+                self.save_model(model_path)
+
+                # Prepare task-specific configuration
+                task_config = {
+                    "task_type": self.task_type,
+                    "config": sanitize_config(self.get_params(deep=True)),
+                    "metrics": metrics or {},
+                }
+
+                # Create or update config
+                base_config = (
+                    {
+                        "model_name": self.model_name,
+                        "framework": "torch_molecule",
+                        "date_created": existing_config.get(
+                            "date_created", datetime.datetime.now().isoformat()
+                        ),
+                    }
+                    if not repo_exists
+                    else existing_config
+                )
+
+                # Merge task configuration
+                final_config = merge_task_configs(
+                    task_id=task_id,
+                    existing_config=base_config,
+                    new_task_config=task_config,
+                    num_params=sum(p.numel() for p in self.model.parameters()),
+                )
+
+                # Save config
+                config_path = os.path.join(tmp_dir, "config.json")
+                with open(config_path, "w") as f:
+                    json.dump(final_config, f, indent=2)
+
+                # Create and save README
+                readme_content = create_model_card(
+                    model_class=self.__class__.__name__,
+                    model_name=self.model_name,
+                    tasks_config=final_config.get("tasks", {}),
+                    model_config=final_config,
+                    repo_id=repo_id,
+                    existing_readme=existing_readme,
+                )
+
+                readme_path = os.path.join(tmp_dir, "README.md")
+                with open(readme_path, "w") as f:
+                    f.write(readme_content)
+
+                # Push files to hub
+                api.upload_folder(
+                    repo_id=repo_id,
+                    folder_path=tmp_dir,
+                    commit_message=f"{commit_message} - Task: {task_id}",
+                )
+
+                # Update repository metadata
+                metadata = {
+                    "tags": [
+                        "torch_molecule",
+                        "molecular-property-prediction",
+                    ],
+                    "library_name": "torch_molecule",
+                    "pipeline_tag": "graph-ml",
+                }
+                metadata_update(repo_id=repo_id, metadata=metadata, token=token, overwrite=True)
+
+                task_config = final_config["tasks"][task_id]
+                print(f"Successfully pushed model for task {task_id} to {repo_id}")
+                print(f"Task version: {task_config['current_version']}")
+                if metrics:
+                    print("Metrics:")
+                    for metric, value in metrics.items():
+                        print(f"  {metric}: {value:.4f}")
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to push to Hugging Face Hub: {str(e)}")