diff --git a/.gitignore b/.gitignore
index cba1731..61afef0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,5 @@ __pycache__/
build/
methylseqnet/configs/paths.toml
.DS_Store
+*.claude
+*.pytest*
diff --git a/README.md b/README.md
index 7e69284..7b23195 100644
--- a/README.md
+++ b/README.md
@@ -11,6 +11,48 @@ To reproduce results presented in Dixon-Luinenburg et al, 2026, reference reprod
To run the model, you can run `pip install git+https://github.com/OberonDixon/methylseq-net` and use the command line entry points or the exposed submodules.
### Entry points
+#### Predictions with a trained checkpoint
+
+Trained checkpoints can either be stored locally, as for newly trained models or models downloaded per the reproducibility instructions, or loaded automatically from [huggingface](https://huggingface.co/OberonDixon/methylseqnet).
+
+The `predict::Predictor` class implements both `from_release(**kwargs)` and `from_local(run_id,checkpoints_dir,**kwargs)` constructors to create `predictor` objects from huggingface or local models respectively. These can then be used to run inference on individual loci, as with `predict_locus`, or on whole datasets, as with `predict_dataset`. The latter can be run from CLI:
+
+```
+usage: methylseqnet-predict [-h] --dataset-keys DATASET_KEYS [DATASET_KEYS ...] --dataset-files DATASET_FILES [DATASET_FILES ...] [--model-source {local,huggingface}] [--model-identifier MODEL_IDENTIFIER] [--checkpoints-dir CHECKPOINTS_DIR] [--hf-base HF_BASE] [--hf-version HF_VERSION] [--hf-repo HF_REPO]
+ [--true-conditioning-state-weight TRUE_CONDITIONING_STATE_WEIGHT] [--no-targets] [--supplemental-outputs [SUPPLEMENTAL_OUTPUTS ...]] [--synthetic-cpg] [--variable-input-length] [--center-methyl-frac CENTER_METHYL_FRAC] [--gpus GPUS] [--num-workers NUM_WORKERS]
+
+Run predictions with a specified model.
+
+options:
+ -h, --help show this help message and exit
+ --dataset-keys DATASET_KEYS [DATASET_KEYS ...]
+ Dataset keys (e.g., atlas, longread) corresponding to the datasets being predicted on (must match length of --dataset-files)
+ --dataset-files DATASET_FILES [DATASET_FILES ...]
+ Paths to dataset H5 files (must match length of --dataset-keys)
+ --model-source {local,huggingface}
+ Source from which to load model checkpoint.
+ --model-identifier MODEL_IDENTIFIER
+ e.g. slurm24807693task2; will reference checkpoints dir
+ --checkpoints-dir CHECKPOINTS_DIR
+ Directory to load trained checkpoints.
+ --hf-base HF_BASE [hf] model base, e.g. borzoi-rep0
+ --hf-version HF_VERSION
+ [hf] release tag, e.g. v1.0
+ --hf-repo HF_REPO [hf] repo id
+ --true-conditioning-state-weight TRUE_CONDITIONING_STATE_WEIGHT
+ If set, override the model's true_conditioning_state_weight with this value for prediction.
+ --no-targets If set, do not include target tracks in the output H5 files.
+ --supplemental-outputs [SUPPLEMENTAL_OUTPUTS ...]
+ Supplemental outputs to include in predictions. Default: ['conditional_seq_rep', 'unconditional_seq_rep', 'true_conditioning_state_rep', 'imputed_conditioning_state_rep', 'cpg_density']
+ --synthetic-cpg If set, add synthetic CpG data.
+ --variable-input-length
+ If set, sequence length can be any integer multiple of 128 that is >=16384.
+ --center-methyl-frac CENTER_METHYL_FRAC
+ Fraction of CpGs methylated in the center window.
+ --gpus GPUS Number of GPUs to use
+ --num-workers NUM_WORKERS
+ Number of data loader workers
+```
#### Preprocess data
Preprocessing scripts take in standard bioinformatic files aligned to a reference. Preprocessing is configured using gin config files.
```
@@ -62,42 +104,7 @@ options:
--logging-level {CRITICAL,ERROR,WARNING,INFO,DEBUG,NOTSET}
Set the logging level
```
-#### Predictions with a trained checkpoint
-```
-usage: methylseqnet-predict [-h] --model-identifier MODEL_IDENTIFIER [--checkpoints-dir CHECKPOINTS_DIR]
- [--true-conditioning-state-weight TRUE_CONDITIONING_STATE_WEIGHT] [--no-targets] --dataset-keys DATASET_KEYS
- [DATASET_KEYS ...] --dataset-files DATASET_FILES [DATASET_FILES ...]
- [--supplemental-outputs [SUPPLEMENTAL_OUTPUTS ...]] [--synthetic-cpg] [--variable-input-length]
- [--center-methyl-frac CENTER_METHYL_FRAC] [--gpus GPUS] [--num-workers NUM_WORKERS]
-Run predictions with a specified model.
-
-options:
- -h, --help show this help message and exit
- --model-identifier MODEL_IDENTIFIER
- e.g. slurm24807693task2; will reference checkpoints dir
- --checkpoints-dir CHECKPOINTS_DIR
- Directory to load trained checkpoints.
- --true-conditioning-state-weight TRUE_CONDITIONING_STATE_WEIGHT
- If set, override the model's true_conditioning_state_weight with this value for prediction.
- --no-targets If set, do not include target tracks in the output H5 files.
- --dataset-keys DATASET_KEYS [DATASET_KEYS ...]
- Dataset keys (e.g., atlas, longread) corresponding to the datasets being predicted on (must match length of --dataset-
- files)
- --dataset-files DATASET_FILES [DATASET_FILES ...]
- Paths to dataset H5 files (must match length of --dataset-keys)
- --supplemental-outputs [SUPPLEMENTAL_OUTPUTS ...]
- Supplemental outputs to include in predictions. Default: ['conditional_seq_rep', 'unconditional_seq_rep',
- 'true_conditioning_state_rep', 'imputed_conditioning_state_rep', 'cpg_density']
- --synthetic-cpg If set, add synthetic CpG data.
- --variable-input-length
- If set, sequence length can be any integer multiple of 128 that is >=16384.
- --center-methyl-frac CENTER_METHYL_FRAC
- Fraction of CpGs methylated in the center window.
- --gpus GPUS Number of GPUs to use
- --num-workers NUM_WORKERS
- Number of data loader workers
-```
#### Create peak files from preprocessed data
```
usage: methylseqnet-peaks [-h] --dataset-paths DATASET_PATHS [DATASET_PATHS ...] --output-directory OUTPUT_DIRECTORY --label-substrings
diff --git a/methylseqnet/hub.py b/methylseqnet/hub.py
new file mode 100644
index 0000000..39330b9
--- /dev/null
+++ b/methylseqnet/hub.py
@@ -0,0 +1,34 @@
+from typing import Literal
+import logging
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download, try_to_load_from_cache
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_REPO = "OberonDixon/methylseqnet"
+DEFAULT_BASE = "borzoi-rep0"
+DEFAULT_VERSION = "v1.0"
+
+def release_checkpoint_path(
+ base: Literal["borzoi-rep0"] = DEFAULT_BASE,
+ version: str = DEFAULT_VERSION,
+ *,
+ repo_id: str = DEFAULT_REPO,
+ filename: str | None = None,
+):
+ filename = f"factorized-{base}.ckpt" if filename is None else filename
+
+ cached = try_to_load_from_cache(repo_id, filename, revision=version) # no network
+ if isinstance(cached, str):
+ ckpt_path = cached
+ method = "Cached"
+ else: # None (not cached) or _CACHED_NO_EXIST sentinel -> go fetch
+ ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename, revision=version)
+ method = "Downloaded"
+
+ sha = Path(ckpt_path).parent.name
+
+ logger.info("%s %s from %s @ %s (%s)", method, filename, repo_id, version, sha)
+
+ return ckpt_path
\ No newline at end of file
diff --git a/methylseqnet/model.py b/methylseqnet/model.py
index 639ecd7..055747a 100644
--- a/methylseqnet/model.py
+++ b/methylseqnet/model.py
@@ -4,7 +4,7 @@
import math
import logging
import warnings
-from typing import Callable, Any, Set, Type
+from typing import Callable, Any, Set, Type, Literal
logger = logging.getLogger(__name__)
import torch
@@ -23,6 +23,7 @@
from methylseqnet.losses import MaskedLoss, PoissonLoss, LogL1Loss, BCELoss, OrthogonalityLoss, MSELoss
from methylseqnet.pretrained import basenji2_pytorch, borzoi_pytorch
from methylseqnet.tensor_ops import FEATURE_MODULATION_OPS
+from methylseqnet.hub import release_checkpoint_path, DEFAULT_REPO, DEFAULT_BASE, DEFAULT_VERSION
gin.register(nn.Softplus)
gin.register(nn.Sigmoid)
@@ -409,6 +410,8 @@ def on_load_checkpoint(self, checkpoint):
@classmethod
def load_from_checkpoint(cls, checkpoint_path, *args, **kwargs):
+ # lazy imports to avoid circular dependencies while giving the gin config what it needs
+ from methylseqnet import callbacks
# Load the checkpoint to extract the gin config
checkpoint = torch.load(checkpoint_path,map_location=torch.device('cpu'))
# Parse the gin configuration from the checkpoint
@@ -434,6 +437,38 @@ def load_from_checkpoint(cls, checkpoint_path, *args, **kwargs):
*args,
**kwargs
)
+
+ @classmethod
+ def from_release(
+ cls,
+ base: Literal["borzoi-rep0"] = DEFAULT_BASE,
+ version: str = DEFAULT_VERSION,
+ *,
+ repo_id: str = DEFAULT_REPO,
+ filename: str | None = None,
+ **kwargs,
+ ):
+ ckpt_path = release_checkpoint_path(base, version, repo_id=repo_id, filename=filename)
+ return cls.load_from_checkpoint(ckpt_path, **kwargs)
+
+ @classmethod
+ def from_local(
+ cls, run_id,
+ *,
+ checkpoints_dir,
+ **kwargs,
+ ):
+ ckpt_path = max(Path(checkpoints_dir, run_id, "checkpoints").glob("best*.ckpt"),
+ key=lambda p: p.stat().st_mtime)
+ return cls.load_from_checkpoint(ckpt_path, **kwargs)
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ *args,
+ **kwargs,
+ ):
+ return cls.from_release(*args,**kwargs)
def set_io_mappings(self,io_mappings_str):
self.io_mappings_str = io_mappings_str
diff --git a/methylseqnet/predict.py b/methylseqnet/predict.py
index cc35e20..754ab02 100644
--- a/methylseqnet/predict.py
+++ b/methylseqnet/predict.py
@@ -12,6 +12,7 @@
import warnings
from typing import Type, Set
from functools import partial
+import logging
import torch
from torch import nn
@@ -30,11 +31,12 @@
from methylseqnet.builders import SingleFastaHandler, MultiFileCpGHandler
from methylseqnet.transforms import LoaderTransform, DinucShuffleSyntheticCpG, InsertSyntheticCpG
from methylseqnet.peaks import selected_peaks_from_target
+from methylseqnet.hub import release_checkpoint_path, DEFAULT_REPO, DEFAULT_BASE, DEFAULT_VERSION
class Predictor:
def __init__(
self,
- model: str | Path | nn.Module,
+ model: str | Path | nn.Module | None = None,
true_conditioning_state_weight: float | None = None,
device: str = 'auto',
supplemental_outputs: set = set(),
@@ -45,12 +47,16 @@ def __init__(
else:
self.device = torch.device(device)
if isinstance(model, nn.Module):
+ # Pre-loaded model
self.model = model
if remove_crop_for_variable_input_length:
warnings.warn("Model provided directly as nn.Module; it may be unsafe to change crop settings so this will be skipped.")
else:
- from methylseqnet.callbacks import ValidationMetricsLogger, GPUMemoryLogger, CPUMemoryLogger, HaplotypedPredLogger
- self.model = ConditionedSeqNN.load_from_checkpoint(model, map_location = self.device)
+ # Load from checkpoint
+ if model is None:
+ self.model = ConditionedSeqNN.from_pretrained(map_location=self.device)
+ else:
+ self.model = ConditionedSeqNN.load_from_checkpoint(model, map_location = self.device)
if remove_crop_for_variable_input_length:
self.model.crop_off_output = 0
self.model.crop_off_conditioning_input = 0
@@ -67,6 +73,19 @@ def __init__(
self.model.true_conditioning_state_weight = true_conditioning_state_weight
self.model.to(self.device)
+ @classmethod
+ def from_release(cls, base=DEFAULT_BASE, version=DEFAULT_VERSION, *,
+ repo_id=DEFAULT_REPO, filename=None, **predictor_kwargs):
+ path = release_checkpoint_path(base=base, version=version,
+ repo_id=repo_id, filename=filename)
+ return cls(model=path, **predictor_kwargs)
+
+ @classmethod
+ def from_local(cls, run_id, *, checkpoints_dir, **predictor_kwargs):
+ ckpt = max(Path(checkpoints_dir, run_id, "checkpoints").glob("best*.ckpt"),
+ key=lambda p: p.stat().st_mtime)
+ return cls(model=ckpt, **predictor_kwargs)
+
def to(self, device):
self.device = torch.device(device)
self.model.to(self.device)
@@ -485,6 +504,7 @@ def __del__(self):
pass
def main():
+ logging.basicConfig(level=logging.INFO)
DEFAULT_SUPPLEMENTAL_OUTPUTS = [
"conditional_seq_rep",
"unconditional_seq_rep",
@@ -493,12 +513,16 @@ def main():
"cpg_density",
]
parser = argparse.ArgumentParser(description="Run predictions with a specified model.")
- parser.add_argument("--model-identifier", required=True, help="e.g. slurm24807693task2; will reference checkpoints dir")
+ parser.add_argument("--dataset-keys", nargs='+', required=True, help="Dataset keys (e.g., atlas, longread) corresponding to the datasets being predicted on (must match length of --dataset-files)")
+ parser.add_argument("--dataset-files", nargs='+', required=True, help="Paths to dataset H5 files (must match length of --dataset-keys)")
+ parser.add_argument("--model-source", choices=["local", "huggingface"], default="huggingface", help="Source from which to load model checkpoint.")
+ parser.add_argument("--model-identifier", default=None, help="e.g. slurm24807693task2; will reference checkpoints dir")
parser.add_argument('--checkpoints-dir', type=str, required=False, default='', help='Directory to load trained checkpoints.')
+ parser.add_argument("--hf-base", default=DEFAULT_BASE, help="[hf] model base, e.g. borzoi-rep0")
+ parser.add_argument("--hf-version", default=DEFAULT_VERSION, help="[hf] release tag, e.g. v1.0")
+ parser.add_argument("--hf-repo", default=DEFAULT_REPO, help="[hf] repo id")
parser.add_argument("--true-conditioning-state-weight", type=float, default=None, help="If set, override the model's true_conditioning_state_weight with this value for prediction.")
parser.add_argument("--no-targets", action='store_true', help="If set, do not include target tracks in the output H5 files.")
- parser.add_argument("--dataset-keys", nargs='+', required=True, help="Dataset keys (e.g., atlas, longread) corresponding to the datasets being predicted on (must match length of --dataset-files)")
- parser.add_argument("--dataset-files", nargs='+', required=True, help="Paths to dataset H5 files (must match length of --dataset-keys)")
parser.add_argument("--supplemental-outputs", nargs="*", required=False, default=DEFAULT_SUPPLEMENTAL_OUTPUTS, help=f"Supplemental outputs to include in predictions. Default: {DEFAULT_SUPPLEMENTAL_OUTPUTS}")
parser.add_argument("--synthetic-cpg", action='store_true', help="If set, add synthetic CpG data.")
parser.add_argument("--variable-input-length", action='store_true', help="If set, sequence length can be any integer multiple of 128 that is >=16384.")
@@ -506,20 +530,38 @@ def main():
parser.add_argument("--gpus", type=int, default=1, help="Number of GPUs to use")
parser.add_argument("--num-workers", type=int, default=8, help="Number of data loader workers")
- args = parser.parse_args()
-
- if args.checkpoints_dir == '':
- try:
- from methylseqnet_repro.paths import model_checkpoints
- except Exception as e:
- print(f"Failed to import model_checkpoints path from methylseqnet_repro.paths: {e}. Using hardcoded UC Berkeley HPC directory instead.")
- model_checkpoints = '/clusterfs/nilah/oberon/lightning/'
- args.checkpoints_dir = model_checkpoints
-
+ args = parser.parse_args()
if len(args.dataset_keys) != len(args.dataset_files):
parser.error(f"--dataset-keys ({len(args.dataset_keys)}) and --dataset-files ({len(args.dataset_files)}) must have the same number of arguments")
+ common = dict(
+ true_conditioning_state_weight=args.true_conditioning_state_weight,
+ supplemental_outputs=set(args.supplemental_outputs),
+ remove_crop_for_variable_input_length=args.variable_input_length,
+ )
+ if args.model_source == "huggingface":
+ if args.model_identifier is not None:
+ parser.error("--model-identifier cannot be used when --model-source is 'huggingface'")
+ predictor = Predictor.from_release(
+ base=args.hf_base, version=args.hf_version, repo_id=args.hf_repo, **common)
+ output_label = f"{args.hf_base}-{args.hf_version}"
+ elif args.model_source == "local":
+ if not args.model_identifier:
+ parser.error("--model-identifier is required when --model-source is 'local'")
+ if args.checkpoints_dir == "":
+ try:
+ from methylseqnet_repro.paths import model_checkpoints
+ except Exception as e:
+ print(f"Failed to import model_checkpoints path: {e}. Using hardcoded dir.")
+ model_checkpoints = "/clusterfs/nilah/oberon/lightning/"
+ args.checkpoints_dir = model_checkpoints
+ predictor = Predictor.from_local(
+ args.model_identifier, checkpoints_dir=args.checkpoints_dir, **common)
+ output_label = args.model_identifier
+ else:
+ parser.error(f"Unsupported model source: {args.model_source}")
+
dataset_paths = [
{dataset_key:dataset_file} for dataset_key, dataset_file in zip(args.dataset_keys, args.dataset_files)
]
@@ -545,22 +587,12 @@ def main():
dataset_dir = Path(list(dataset_path.values())[0]).parent
print(f"Running through {dataset_name}.")
if args.synthetic_cpg:
- output_path = dataset_dir / args.model_identifier / f"{dataset_name}_synthetic_{args.center_methyl_frac}"
+ output_path = dataset_dir / output_label / f"{dataset_name}_synthetic_{args.center_methyl_frac}"
elif args.true_conditioning_state_weight is not None:
- output_path = dataset_dir / args.model_identifier / f"{dataset_name}_truecondweight_{args.true_conditioning_state_weight}"
+ output_path = dataset_dir / output_label / f"{dataset_name}_truecondweight_{args.true_conditioning_state_weight}"
else:
- output_path = dataset_dir / args.model_identifier / dataset_name
- best_ckpt = max(
- Path(Path(args.checkpoints_dir) / args.model_identifier / "checkpoints").glob('best*.ckpt'),
- key=lambda p: p.stat().st_mtime
- )
-
- predictor = Predictor(
- model=best_ckpt,
- true_conditioning_state_weight=args.true_conditioning_state_weight,
- supplemental_outputs = set(args.supplemental_outputs),
- remove_crop_for_variable_input_length = args.variable_input_length,
- )
+ output_path = dataset_dir / output_label / dataset_name
+
predictor.predict_dataset(
dataset_path=dataset_path,
output_path=output_path,
diff --git a/pyproject.toml b/pyproject.toml
index 267c1cf..a78801c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,6 +25,8 @@ dependencies = [
"numpy>=1.26.4",
"psutil>=5.9.8",
"tqdm>=4.66.4",
+ "huggingface_hub<1.0",
+ "transformers<5",
]
[project.optional-dependencies]
diff --git a/tests/test_hub.py b/tests/test_hub.py
new file mode 100644
index 0000000..0ea0b35
--- /dev/null
+++ b/tests/test_hub.py
@@ -0,0 +1,160 @@
+from unittest import mock
+
+import pytest
+import torch
+
+import methylseqnet.hub as hub
+from methylseqnet.hub import (
+ release_checkpoint_path,
+ DEFAULT_REPO,
+ DEFAULT_BASE,
+ DEFAULT_VERSION,
+)
+
+DEFAULT_FILENAME = f"factorized-{DEFAULT_BASE}.ckpt"
+
+
+def test_release_checkpoint_path_uses_cache_when_available():
+ """A cached file (str path) is returned directly without downloading."""
+ cached = "/fake/cache/models--repo/snapshots/abc123/" + DEFAULT_FILENAME
+ with mock.patch.object(hub, "try_to_load_from_cache", return_value=cached) as mock_cache, \
+ mock.patch.object(hub, "hf_hub_download") as mock_download:
+ path = release_checkpoint_path()
+
+ assert path == cached
+ mock_cache.assert_called_once_with(DEFAULT_REPO, DEFAULT_FILENAME, revision=DEFAULT_VERSION)
+ mock_download.assert_not_called()
+
+
+def test_release_checkpoint_path_downloads_when_not_cached():
+ """None from the cache lookup means not cached, so we download."""
+ downloaded = "/fake/hub/models--repo/snapshots/def456/" + DEFAULT_FILENAME
+ with mock.patch.object(hub, "try_to_load_from_cache", return_value=None), \
+ mock.patch.object(hub, "hf_hub_download", return_value=downloaded) as mock_download:
+ path = release_checkpoint_path()
+
+ assert path == downloaded
+ mock_download.assert_called_once_with(
+ repo_id=DEFAULT_REPO, filename=DEFAULT_FILENAME, revision=DEFAULT_VERSION
+ )
+
+
+def test_release_checkpoint_path_downloads_on_cache_no_exist_sentinel():
+ """A non-str sentinel (_CACHED_NO_EXIST) should download, not be used as a path."""
+ sentinel = object() # stand-in for huggingface_hub's _CACHED_NO_EXIST
+ downloaded = "/fake/hub/x/" + DEFAULT_FILENAME
+ with mock.patch.object(hub, "try_to_load_from_cache", return_value=sentinel), \
+ mock.patch.object(hub, "hf_hub_download", return_value=downloaded) as mock_download:
+ path = release_checkpoint_path()
+
+ assert path == downloaded
+ mock_download.assert_called_once()
+
+
+def test_release_checkpoint_path_default_filename_derived_from_base():
+ """With no filename given, it is built as factorized-.ckpt."""
+ with mock.patch.object(hub, "try_to_load_from_cache", return_value=None) as mock_cache, \
+ mock.patch.object(hub, "hf_hub_download", return_value="/x.ckpt") as mock_download:
+ release_checkpoint_path(base="borzoi-rep0")
+
+ # try_to_load_from_cache is called positionally as (repo_id, filename)
+ assert mock_cache.call_args.args[1] == "factorized-borzoi-rep0.ckpt"
+ assert mock_download.call_args.kwargs["filename"] == "factorized-borzoi-rep0.ckpt"
+
+
+def test_release_checkpoint_path_passes_through_custom_args():
+ """Explicit base/version/repo_id/filename are forwarded to the hub calls."""
+ with mock.patch.object(hub, "try_to_load_from_cache", return_value=None) as mock_cache, \
+ mock.patch.object(hub, "hf_hub_download", return_value="/custom.ckpt") as mock_download:
+ path = release_checkpoint_path(
+ "borzoi-rep0", "v9.9", repo_id="someone/else", filename="custom.ckpt"
+ )
+
+ assert path == "/custom.ckpt"
+ mock_cache.assert_called_once_with("someone/else", "custom.ckpt", revision="v9.9")
+ mock_download.assert_called_once_with(
+ repo_id="someone/else", filename="custom.ckpt", revision="v9.9"
+ )
+
+
+# Real released checkpoint: fetched via the hub.py DEFAULTs so that bumping
+# DEFAULT_VERSION (and publishing a matching checkpoint) keeps these tests pointed
+# at the current release without edits here. Borzoi-style models take a fixed
+# 524288 bp input window; reuse the fake genome / zeroed methylation track.
+RELEASE_INPUT_LENGTH = 524288
+FAKE_GENOME = "./tests/data/chr1_fake1M.fa.gz"
+ZEROED_METHYL = "./tests/data/hg38_test_zeros.hg38.bigwig"
+
+
+def _network_errors():
+ """Exceptions meaning 'couldn't reach the Hub' rather than a real compat break."""
+ errs = [OSError] # ConnectionError, timeouts, etc. subclass OSError
+ try:
+ from requests.exceptions import RequestException
+ errs.append(RequestException)
+ except Exception:
+ pass
+ try:
+ from huggingface_hub.errors import HfHubHTTPError, LocalEntryNotFoundError
+ errs.extend([HfHubHTTPError, LocalEntryNotFoundError])
+ except Exception:
+ pass
+ return tuple(errs)
+
+
+NETWORK_ERRORS = _network_errors()
+
+
+@pytest.fixture(scope="module")
+def released_checkpoint_path():
+ """Current released checkpoint (DEFAULT base/version), or skip if the Hub is unreachable."""
+ try:
+ return release_checkpoint_path()
+ except NETWORK_ERRORS as e:
+ pytest.skip(f"Could not fetch released checkpoint from the Hub (offline?): {e}")
+
+
+def test_model_from_release_returns_conditioned_seqnn(released_checkpoint_path):
+ """ConditionedSeqNN.from_release loads the published checkpoint into a model."""
+ from methylseqnet.model import ConditionedSeqNN
+
+ model = ConditionedSeqNN.from_release()
+ assert isinstance(model, ConditionedSeqNN)
+
+
+def test_model_from_pretrained_is_alias_for_from_release(released_checkpoint_path):
+ """from_pretrained is a thin alias of from_release and yields the same type."""
+ from methylseqnet.model import ConditionedSeqNN
+
+ model = ConditionedSeqNN.from_pretrained()
+ assert isinstance(model, ConditionedSeqNN)
+
+
+@pytest.mark.skipif(
+ not torch.cuda.is_available() or torch.cuda.device_count() < 1,
+ reason="Forward pass over the full input window requires a GPU",
+)
+def test_released_checkpoint_runs_forward_pass(released_checkpoint_path):
+ """Load the published checkpoint and run a forward pass, guarding against future code changes silently breaking checkpoint compatibility."""
+ from methylseqnet.predict import Predictor
+
+ # Uses the cache populated by the released_checkpoint_path fixture
+ predictor = Predictor.from_release()
+
+ prediction = predictor.predict_locus(
+ chromosome="chr1",
+ start=0,
+ end=RELEASE_INPUT_LENGTH,
+ sequence_path=FAKE_GENOME,
+ methylation_paths=[ZEROED_METHYL],
+ capture_attributions=False,
+ )
+
+ preds = prediction["predictions"]
+ assert isinstance(preds, torch.Tensor)
+ assert preds.ndim >= 2, f"expected a multi-dim prediction tensor, got shape {preds.shape}"
+ # Output is binned at the model stride; length must match the per-bin coordinates
+ assert preds.shape[-1] > 0
+ assert preds.shape[-1] == prediction["output_coordinates"].shape[-1], \
+ "prediction length does not match the number of output bins"
+ assert torch.isfinite(preds).all(), "released checkpoint produced non-finite predictions"
diff --git a/tests/test_train_predict_integration.py b/tests/test_train_predict_integration.py
index 0cd15d5..d107350 100644
--- a/tests/test_train_predict_integration.py
+++ b/tests/test_train_predict_integration.py
@@ -1,4 +1,6 @@
import os
+import sys
+import shutil
import tempfile
from pathlib import Path
import warnings
@@ -9,8 +11,10 @@
from captum.attr import IntegratedGradients
import methylseqnet.train as train
+import methylseqnet.predict as predict
from methylseqnet.model import ConditionedSeqNN
from methylseqnet.predict import Predictor
+from methylseqnet.hub import DEFAULT_BASE, DEFAULT_VERSION
from test_model import get_config_files_with_names, nuke_gin_config
import gin
@@ -21,7 +25,7 @@
reason="Test requires at least one GPU",
)
@pytest.mark.parametrize("config_file", get_config_files_with_names())
-def test_train_predict_integration(config_file):
+def test_train_predict_integration(config_file, monkeypatch):
nuke_gin_config()
wandb.finish()
with tempfile.TemporaryDirectory() as temp_dir:
@@ -104,4 +108,77 @@ def test_train_predict_integration(config_file):
)
except RuntimeError as e:
if "used in the graph" in str(e) and "allow_unused" in str(e):
- warnings.warn(f"RuntimeError during predict_locus with attributions: {e}. This may be due to the attribution method not being compatible with the model architecture, and can be ignored for the purposes of this integration test.")
\ No newline at end of file
+ warnings.warn(f"RuntimeError during predict_locus with attributions: {e}. This may be due to the attribution method not being compatible with the model architecture, and can be ignored for the purposes of this integration test.")
+
+ # CLI integration: exercise methylseqnet-predict for both model sources
+ # The CLI writes prediction outputs alongside each input file, so copy the
+ # datasets into the temp dir to avoid polluting tests/data. Datasets may be
+ # listed per-key as a single path or a list; normalize to one file per key.
+ cli_data_dir = Path(temp_dir) / "cli_data"
+ cli_data_dir.mkdir(exist_ok=True)
+ cli_keys, cli_files = [], []
+ for key, value in trainer.datamodule.train_dataset_dict.items():
+ src = value[0] if isinstance(value, (list, tuple)) else value
+ dst = cli_data_dir / f"{key}_{Path(src).name}"
+ shutil.copy(src, dst)
+ cli_keys.append(key)
+ cli_files.append(str(dst))
+
+ # Empty --supplemental-outputs keeps the test model-agnostic across the
+ # parametrized configs (factorized-only reps aren't available everywhere).
+ common_argv = [
+ "--dataset-keys", *cli_keys,
+ "--dataset-files", *cli_files,
+ "--supplemental-outputs",
+ "--gpus", "1",
+ "--num-workers", "1",
+ ]
+
+ # local source: load the checkpoint just trained above via the run identifier.
+ # In local mode the output label is the model identifier ("test").
+ monkeypatch.setattr(sys, "argv", [
+ "methylseqnet-predict",
+ "--model-source", "local",
+ "--model-identifier", "test",
+ "--checkpoints-dir", temp_dir,
+ *common_argv,
+ ])
+ predict.main()
+ assert list((cli_data_dir / "test").rglob("predictions.h5")), \
+ "local CLI run produced no predictions"
+
+ # local source requires --model-identifier; omitting it is a usage error
+ monkeypatch.setattr(sys, "argv", [
+ "methylseqnet-predict",
+ "--model-source", "local",
+ "--checkpoints-dir", temp_dir,
+ *common_argv,
+ ])
+ with pytest.raises(SystemExit):
+ predict.main()
+
+ # huggingface source: mock the hub download to return a local checkpoint so the
+ # CLI from_release plumbing is exercised without any network access.
+ fake_ckpt = str(checkpoints_dir / "temp-checkpoint.ckpt")
+ monkeypatch.setattr(predict, "release_checkpoint_path", lambda *args, **kwargs: fake_ckpt)
+
+ # --model-identifier is rejected when the source is huggingface
+ monkeypatch.setattr(sys, "argv", [
+ "methylseqnet-predict",
+ "--model-source", "huggingface",
+ "--model-identifier", "test",
+ *common_argv,
+ ])
+ with pytest.raises(SystemExit):
+ predict.main()
+
+ # valid huggingface invocation; output label is "-"
+ monkeypatch.setattr(sys, "argv", [
+ "methylseqnet-predict",
+ "--model-source", "huggingface",
+ *common_argv,
+ ])
+ predict.main()
+ hf_label = f"{DEFAULT_BASE}-{DEFAULT_VERSION}"
+ assert list((cli_data_dir / hf_label).rglob("predictions.h5")), \
+ "huggingface CLI run produced no predictions"
\ No newline at end of file