Huggingface

.gitignore

@@ -6,3 +6,5 @@ __pycache__/
 build/
 methylseqnet/configs/paths.toml
 .DS_Store
+*.claude
+*.pytest*

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

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

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

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,33 +513,55 @@ 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.")
     parser.add_argument("--center-methyl-frac", type=float, default=0.05, help="Fraction of CpGs methylated in the center window.")
     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,

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]