OCDocker.OCScore.Optimization.ModelExport module

Export and reload best OCScore Optuna models for inference and retraining.

After a PDBbind or DUDEz study completes, export_best_model_bundle() writes a best_model/ directory containing weights, architecture, retraining config, feature metadata, and a compact trial summary.

Usage:

from OCDocker.OCScore.Optimization.ModelExport import load_exported_model

OCDocker.OCScore.Optimization.ModelExport.export_best_model_bundle(export_dir, task, model, model_config, selected_features, best_trial_number, best_objective_value, validation_metrics, test_metrics, stage_config, splits, objective_metric, direction, best_params, random_seed=None, source_checkpoint_path=None, training_metrics=None, extra=None, calibrator=None, validate=True, source_dataframe=None)[source]

Export the best completed trial model into a reloadable bundle.

Parameters:

  • export_dir (str | Path) – Directory that will contain best_model.pt and companion metadata.

  • task (str) – "pdbbind_regression" or "dudez_screening".

  • model (nn.Module) – Best trained model instance.

  • model_config (Mapping[str, Any]) – Resolved model configuration after conditional search-space logic.

  • selected_features (Sequence[str]) – Feature names in training column order.

  • best_trial_number (int) – Optuna trial number for the exported model.

  • best_objective_value (float) – Final objective value on validation data.

  • validation_metrics (Mapping[str, Any]) – Validation metrics for the best model.

  • test_metrics (Mapping[str, Any]) – Test metrics for the best model.

  • stage_config (Mapping[str, Any]) – Serialized stage configuration (Optuna settings, splits, pruning, etc.).

  • splits (Mapping[str, Any]) – Prepared split payload including indices and optional scaler.

  • objective_metric (str) – Effective objective metric name.

  • direction (str) – Optuna optimization direction.

  • best_params (Mapping[str, Any]) – Raw Optuna trial parameters.

  • random_seed (int | None, optional) – Random seed used by the stage, by default None.

  • source_checkpoint_path (str | None, optional) – Path to the source *_best.pt checkpoint, by default None.

  • training_metrics (Mapping[str, Any] | None, optional) – Optional training diagnostics, by default None.

  • extra (Mapping[str, Any] | None, optional) – Additional export metadata, by default None.

  • calibrator (Any | None, optional) – Fitted ProbabilityCalibrator for DUDEz exports (saved as probability_calibrator.joblib).

  • validate (bool, optional) – Run validate_export_bundle() before returning, by default True.

  • source_dataframe (pd.DataFrame | None, optional) – Source reduced dataframe used to compute forbidden blind-evaluation hashes.

Returns:

Absolute paths for exported artifacts.

Return type:

dict[str, str]

OCDocker.OCScore.Optimization.ModelExport.load_exported_model(export_dir, device='cpu', transferred_extractor=None, pdbbind_export_dir=None)[source]

Load an exported best-model bundle for inference or evaluation.

Parameters:

  • export_dir (str | Path) – Exported best_model/ directory.

  • device (torch.device | str | None, optional) – Target device, by default CPU.

  • transferred_extractor (FeatureExtractor | None, optional) – Optional transferred extractor for DUDEz transfer exports. Takes precedence over pdbbind_export_dir when both are given.

  • pdbbind_export_dir (str | Path | None, optional) – Path to the linked PDBbind export. Used to resolve the transferred feature extractor for DUDEz transfer models whose recorded extra.pdbbind_best_model_export_dir no longer resolves (e.g. after moving the bundle to another machine), and as a fallback source for scaler when export_dir has none of its own.

Returns:

Loaded model and metadata keys including model, scaler, selected_features, architecture, retrain_config, and summary.

Return type:

dict[str, Any]

OCDocker.OCScore.Optimization.ModelExport.predict_from_export(export_dir, dataframe, *, device='cpu', pdbbind_export_dir=None)[source]

Score rows from a wide feature table using an exported best-model bundle.

Parameters:

  • export_dir (str or Path) – Exported best_model/ directory.

  • dataframe (pd.DataFrame) – Wide pipeline feature table containing export selected_features.

  • device (torch.device or str, optional) – Torch device for inference, by default CPU.

  • pdbbind_export_dir (str or Path, optional) – Override path to the linked PDBbind export for DUDEz transfer models.

Returns:

Input metadata with ocscore_prediction and, for DUDEz exports, ocscore_probability.

Return type:

pd.DataFrame

OCDocker.OCScore.Optimization.ModelExport.retrain_from_export(export_dir, pdbbind_df=None, dudez_df=None, device=None, use_saved_split_indices=True)[source]

Prepare data splits and a fresh model for retraining from an export bundle.

This does not run training; it returns the model, optimizer-related settings, prepared splits, and metadata needed to launch a training loop.

Parameters:

  • export_dir (str | Path) – Exported best_model/ directory.

  • pdbbind_df (pd.DataFrame | None, optional) – Reduced PDBbind dataframe for regression retraining.

  • dudez_df (pd.DataFrame | None, optional) – Reduced DUDEz dataframe for screening retraining.

  • device (torch.device | str | None, optional) – Target device, by default CPU.

  • use_saved_split_indices (bool, optional) – Reuse exported split indices when available, by default True.

Returns:

Retraining payload with model, splits, model_config, and stage_config.

Return type:

dict[str, Any]

OCDocker.OCScore.Optimization.ModelExport.transform_export_features(dataframe, selected_features, scaler)[source]

Extract and optionally scale exported feature columns.

Parameters:

  • dataframe (pd.DataFrame) – Input rows containing selected_features.

  • selected_features (Sequence[str]) – Feature names in model input order.

  • scaler (Any or None) – Optional fitted scaler (PDBbind exports).

Returns:

Feature matrix ready for model forward pass.

Return type:

np.ndarray

OCDocker.OCScore.Optimization.ModelExport.validate_export_bundle(export_dir, device=None, pdbbind_export_dir=None)[source]

Rebuild the exported model and verify weights load successfully.

Parameters:

  • export_dir (str | Path) – Exported best_model/ directory.

  • device (torch.device | None, optional) – Device used for reconstruction smoke test, by default CPU.

  • pdbbind_export_dir (str | Path | None, optional) – Override path to the linked PDBbind export for DUDEz transfer models, for bundles whose recorded extra.pdbbind_best_model_export_dir no longer resolves (e.g. after moving the bundle to another machine).

Returns:

Validation metadata including parameter counts.

Return type:

dict[str, Any]

OCDocker.OCScore.Optimization.ModelExport.validate_export_features(dataframe, selected_features)[source]

Ensure a dataframe contains all exported selected feature columns.

Parameters:

  • dataframe (pd.DataFrame) – Input feature table.

  • selected_features (Sequence[str]) – Feature names required by the export bundle.

Raises:

ValueError – If any selected feature column is missing.

Return type:

None