physioex.models.linear_probe

physioex.models.linear_probe(model_name, dataset_name, n_folds=5, max_epochs=100, lr=0.001, weight_decay=0.0001, batch_size=512, device='cpu', upload=False, cache_dir=None, save_predictions=False)

5-fold subject-wise cross-validated linear probe on cached embeddings.

Loads per-subject embeddings and labels from the local cache, splits subjects into n_folds folds, trains a nn.Linear classifier per fold, and reports pooled and per-fold metrics.

Metrics reported (standard in sleep staging literature):

• Accuracy (ACC) — overall classification accuracy.

• Macro F1 (MF1) — unweighted mean of per-class F1; handles class imbalance.

• Cohen’s Kappa (κ) — chance-corrected agreement; the standard metric in sleep medicine.

• Per-class F1 — shows per-stage performance (N1 is typically the hardest).

Results are saved to linear_probe_results.json in the embeddings cache directory. If upload is True, the entire directory (embeddings + results) is uploaded to HuggingFace Hub.

Parameters:

• model_name (str) – Model identifier matching cache directory.

• dataset_name (str) – Dataset identifier matching cache directory.

• n_folds (int) – Number of cross-validation folds (default 5).

• max_epochs (int) – Training epochs per fold (default 100).

• lr (float) – Learning rate (default 1e-3).

• weight_decay (float) – L2 regularization (default 1e-4).

• batch_size (int) – Training batch size (default 512).

• device (str) – Device string ("cpu" or "cuda:0").

• upload (bool) – If True, upload results to HuggingFace Hub.

• cache_dir (str | None) – Override cache root directory.

• save_predictions (bool) – If True, save per-subject softmax probabilities and labels to linear_probe_predictions.json. Each entry contains subject_id, fold index, proba (n_epochs, n_classes), and labels (n_epochs,).

Returns:

Results dict with per_fold, pooled, and mean_std keys.

Return type:

dict