14. Predict with the score model¶
This page runs the trained AP-1 score model (previous page) forward. It is the same
prediction_task as the accessibility page — DNA
window + cell state → base-resolution accessibility — with one addition: each cell state
also carries a chromVAR AP-1 activity score, and the model predicts the accessibility
that matches it. The goal here is the I/O contract: how the score is supplied and what
the task writes out — not benchmarking.
This uses the
bolerodatamodel zoo (the paper route, like the variant-effect page) and needs the data lake. The AP-1 model is registered as an ATAC score model (Score Input = AP1, no gene-count head).
Setup¶
from bolero import init
from bolerodata import MODELS
init(num_cpus=24)
2026-07-17 12:14:28,125 INFO worker.py:1781 -- Started a local Ray instance.
Pick the model and the dataset to predict for¶
MODELS["Bolero10M+AP1+260104"] is the trained AP-1 ATAC score model. Because it is a
multi-dataset model, you must tell it which dataset's cell states you are predicting
(model.dataset) — that selects the per-dataset conditioning encoder and the peak set.
MODEL_NAME = "Bolero10M+AP1+260104"
DATASET_NAME = "HumanBrainDev"
model = MODELS[MODEL_NAME]
model.dataset = DATASET_NAME # multi-dataset model: set the dataset before predicting
print("model group:", model.model_group, "| score input:", model._tf_score_input)
model group: BorzoiSignalMultiDS | score input: AP1
The score rides in with the cell states¶
A prediction is made for a (region, cell state) pair. For the score model each cell
state is a pseudobulk record that carries two things: its cell-state embedding and
its __shared_data__ — the AP-1 score for that state. These prebuilt records are what
supply the score at inference; the model reads __shared_data__ and folds it into the same
conditioning vector it was trained with.
import joblib
# prebuilt per-pseudobulk records: {embedding, __shared_data__ (the AP-1 score)} per state
PSEUDOBULK_RECORDS = "PseudobulkRecords.HumanBrainDev.None.1.AP1.joblib"
records = joblib.load(PSEUDOBULK_RECORDS)
records = records.get("pseudobulk_records", records)
_pid, _rec = next(iter(records.items()))
print("a cell-state record holds:", list(_rec))
print(" embedding shape :", _rec["embedding"].shape)
print(" __shared_data__ :", _rec["__shared_data__"], "(the AP-1 activity score)")
a cell-state record holds: ['embedding', '__shared_data__'] embedding shape : (30,) __shared_data__ : [-1.5] (the AP-1 activity score)
Build the predictor and run the task¶
create_predictor assembles the model + data manager from the config; prediction_task
does the GPU forward pass and writes the outputs. The toggles below match a
signal-free, embedding-driven run:
nosignal=True— predict from DNA + conditioning only (no reference-coverage input). The AP-1 score is still applied; only the reference ATAC track branch is off.use_ref_bw=False— no reference bigWig needed (paired withnosignal).embedding_only_mode=True— drive purely from the pseudobulk embeddings/scores; skip loading ground-truth signal.peak_bed_path— the dataset's peaks, so the base track is also summarized per peak.
regions="test_regions" predicts on the model's held-out chromosomes; we keep just a handful
(downsample_regions=8) so the demo runs quickly. Drop that argument to predict the full set.
predictor = model.create_predictor(
subset_name=None,
pseudobulk_records_path=PSEUDOBULK_RECORDS,
sample_n_pseudobulks=None,
nosignal=True,
use_ref_bw=False,
embedding_only_mode=True,
peak_bed_path=model.dataset.peak_bed_path,
)
output_dir = f"predictions/{MODEL_NAME}:{DATASET_NAME}"
predictor.prediction_task(
output_dir,
regions="test_regions",
downsample_regions=8, # demo: a handful of held-out regions (drop to predict all)
mode="prediction",
)
print("wrote predictions to", output_dir)
Saving batches to /large_storage/zhoulab/hanliu/260102-ScoreModel/whole_genome/HumanBrainDev/predictions/Bolero10M+AP1+260104:HumanBrainDev/batch Saving stats to /large_storage/zhoulab/hanliu/260102-ScoreModel/whole_genome/HumanBrainDev/predictions/Bolero10M+AP1+260104:HumanBrainDev/summary_stats.joblib.gz Using temporary directory /tmp/tmpktz288jl 0 regions has finished in /large_storage/zhoulab/hanliu/260102-ScoreModel/whole_genome/HumanBrainDev/predictions/Bolero10M+AP1+260104:HumanBrainDev/batch 8 regions to compute Data loader batch size 16
Loading base model weights from: /home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_0.renamed.pt
Model does not have gene_count_output_head, skip Loading checkpoint from /large_storage/zhoulab/hanliu/260102-ScoreModel/model/260102-Bolero10M_AP1_atac_fold0.lora.best_checkpoint.pt
========== Dataloader Batch Schema: - __conditionemb__: <class 'torch.Tensor'> torch.Size([52, 1, 3]) torch.float32 cuda:0 - __conditionemb__:cond0: <class 'torch.Tensor'> torch.Size([26, 1, 3]) torch.float32 cuda:0 - __conditionemb__:cond1: <class 'torch.Tensor'> torch.Size([26, 1, 3]) torch.float32 cuda:0 - __conditionemb__:delta: <class 'torch.Tensor'> torch.Size([26, 1, 3]) torch.float32 cuda:0 - __dna__: <class 'torch.Tensor'> torch.Size([8, 4, 524288]) torch.uint8 cuda:0 - __embedding__: <class 'torch.Tensor'> torch.Size([52, 60]) torch.float32 cuda:0 - __embedding__:cond0: <class 'torch.Tensor'> torch.Size([26, 60]) torch.float32 cuda:0 - __embedding__:cond1: <class 'torch.Tensor'> torch.Size([26, 60]) torch.float32 cuda:0 - __embedding__:delta: <class 'torch.Tensor'> torch.Size([26, 60]) torch.float32 cuda:0 - __embedding_only__:cov_scale: <class 'torch.Tensor'> torch.Size([52]) torch.float64 cuda:0 - __shared_data__: <class 'torch.Tensor'> torch.Size([52, 1]) torch.float64 cuda:0 - __shared_data__:cond0: <class 'torch.Tensor'> torch.Size([26, 1]) torch.float64 cuda:0 - __shared_data__:cond1: <class 'torch.Tensor'> torch.Size([26, 1]) torch.float64 cuda:0 - __shared_data__:delta: <class 'torch.Tensor'> torch.Size([26, 1]) torch.float64 cuda:0 - __ypred__:cond1: <class 'torch.Tensor'> torch.Size([8, 26, 16352]) torch.float32 cuda:0 - __ypred__:peak:cond1: <class 'torch.Tensor'> torch.Size([26, 1793]) torch.float32 cuda:0 - condition_pairs: <class 'numpy.ndarray'> (26, 2) <U8 - peak: <class 'pandas.core.frame.DataFrame'> (1793, 7) - pseudobulk_ids: <class 'numpy.ndarray'> (52,) <U25 - region: <class 'numpy.ndarray'> (8,) <U25 - region_name: <class 'numpy.ndarray'> (8,) <U25 ========== Total time: 8.88s Inference time: 8.41s or 8.406s per batch Callback time: 0.28s or 0.277s per batch (total 1 batches) ========== Saved Batch Schema: - __embedding__: <class 'numpy.ndarray'> (52, 60) float32 - __ypred__:peak:cond1: <class 'numpy.ndarray'> (26, 1793) float32 - condition_pairs: <class 'numpy.ndarray'> (26, 2) <U8 - peak: <class 'pandas.core.frame.DataFrame'> (1793, 7) - pseudobulk_ids: <class 'numpy.ndarray'> (52,) <U25 - region: <class 'numpy.ndarray'> (8,) <U25 - region_name: <class 'numpy.ndarray'> (8,) <U25 ========== ========== Final Stats Batch Schema: - condition_pairs: <class 'numpy.ndarray'> (26, 2) <U8 - peak: <class 'pandas.core.frame.DataFrame'> (1793, 7) - pseudobulk_ids: <class 'numpy.ndarray'> (52,) <U25 - region: <class 'numpy.ndarray'> (8,) <U25 ========== Removed temporary files in /tmp/tmpktz288jl 1 batches in /large_storage/zhoulab/hanliu/260102-ScoreModel/whole_genome/HumanBrainDev/predictions/Bolero10M+AP1+260104:HumanBrainDev True peak data not found in batch, skip true peak data saving. wrote predictions to predictions/Bolero10M+AP1+260104:HumanBrainDev
What you get¶
prediction_task writes an output_dir/ with the per-batch predictions plus aggregated
files: a per-base 32 bp accessibility track for each region × cell-state, a per-peak
accessibility value, and summary_stats. Rows are the cell states you supplied (their
pseudobulk ids); each was predicted with its own AP-1 score baked in.
That is the whole contract: supply regions + score-carrying cell states, get score-aware
accessibility back. To measure a TF's sequence effect you would predict the same regions
and states while changing __shared_data__ and compare — but that (and benchmarking the
predictions against held-out activity data) is beyond this I/O walkthrough.