ExoShorkie integration: head port + distillation pipeline (WIP) - #20
Draft
tdsone wants to merge 8 commits into
Draft
ExoShorkie integration: head port + distillation pipeline (WIP)#20tdsone wants to merge 8 commits into
tdsone wants to merge 8 commits into
Conversation
ExoShorkie = the Shorkie trunk + a single linear Dense(1) per-bin coverage head (no softplus), squeezed to (B, T_out). Make the head and input encoding config-driven so one trunk implementation serves both models: - head activation/squeeze from the head config (Shorkie: softplus/5215; ExoShorkie: linear/1) - input encoding via species_channel/encode_n_channel (Shorkie: host at global col 114, no N channel; ExoShorkie: col 119 + explicit N channel, matching build_shorkie_features) - _load_tf_weights auto-detects the per_bin_f<F> head group (matched by name, not dense index) - exoshorkie_config() derives the ExoShorkie config from Shorkie's Shorkie path is bit-for-bit unchanged (defaults reproduce it; full suite + characterization tests pass). tests/test_exoshorkie_head.py: exact weight-identity vs the probe .h5 + golden lock.
Distill each genome's 40-model teacher ensemble into one PyTorch student, run entirely through the verified port (no TensorFlow): - gen_targets.py: sample genomic windows, build AlphaGenome-style synthetic windows (RC/SNV/SV, the authors' recipe), run the 40 teachers and average to (50k, 896) soft targets. Feeds host channel 119 + N channel (ExoShorkie's convention). - train_student.py: student = Shorkie fold-0 trunk + fresh Dense(1); MSE, Adam 2e-5, 20 epochs; gates on held-out Pearson r > 0.98. - distill_all.sh: unattended serial driver over the 6 genomes (download teachers + FASTA -> target-gen -> train -> gate -> delete teachers on pass; halts on any failure). M. pneumoniae validated end-to-end: student r = 0.9947 > 0.98.
Wire the 6 distilled per-genome students into the benchmark as the "exoshorkie" model. The students predict in the authors' log-z space; the new ExoShorkie wrapper inverts to count space (count = expm1(z*sigma + mu) per student) so the Yorzoi/Shorkie-style readouts transfer. - models/exoshorkie/wrapper.py: 6-student ensemble. Denorm-then-mean (the inverse is nonlinear and each student has its own mu/sigma), RC averaging in count space, ÷16 repeat-unbin to match the Shorkie/Yorzoi per-base scale. forward_perbase / forward_count_bins, plus a forward_track_mean_perbase alias so the Shorkie coverage scaffolds drop in. - adapters/_exoshorkie_constants.py: geometry + per-genome (mu, sigma), reproduced from the public figshare coverage NPZ by compute_logz_stats.py. - adapters/exoshorkie_brooks.py: CoverageTrackPredictor (native task). - scripts/exoshorkie/compute_logz_stats.py: recompute the 12 constants. - scripts/exoshorkie/validate_denorm.py: closed-loop check vs figshare truth (Spearman 0.80-0.97, magnitude ratio 0.74-0.95 across the 6 genomes). - registry: MODELS["exoshorkie"] + EXOSHORKIE_ADAPTERS (Brooks only for now; scalar tasks raise a clear "no adapter" until their adapters land). Shorkie path untouched (bit-exact). New wrapper unit tests run without a checkpoint or GPU.
Wire ExoShorkie into every task Shorkie runs, not just Brooks. Each scalar/
variant adapter is a thin subclass of the corresponding Shorkie adapter with an
ExoShorkie model injected via from_students — the readout is inherited verbatim,
so the only experimental variable between the Shorkie and ExoShorkie rows is the
model. Works because the ExoShorkie wrapper aliases forward_track_mean_perbase
(returning its count-space coverage, ignoring the track argument).
- adapters/exoshorkie_adapters.py: ExoShorkie{Wu,Hong,Chen,MPRAMarginalized,
Shalem}Predictor + ExoShorkieVariantScorer (eqtl), all subclassing the Shorkie
adapters with track_subset=[].
- registry: EXOSHORKIE_ADAPTERS now covers all 7 protocols (parity with Shorkie);
6 new _exoshorkie_*_adapter build fns.
- configs/default.yaml: exoshorkie run block (all Shorkie tasks).
- configs/exoshorkie_all.yaml: dedicated full run, cheap→expensive ordering,
rafi_mpra (71k seqs) last. eqtl omitted — the shared shorkie_eqtl adapter
rejects the caudal/kita distributions on a single-base REF check (multi-base
ref alleles), so Shorkie has no eqtl baseline either; pre-existing, orthogonal.
- configs/exoshorkie.yaml: Brooks-only demo.
Smoke-validated end-to-end (n_sample-capped): mpra/shalem (marginalized engine)
and hong (direct-coverage scalar) run clean.
gen_plot_preds.py (ExoShorkie + Shorkie, 16384 win) and gen_plot_preds_yorzoi.py (Yorzoi, 4992 win, run from the sibling yorzoi env on CUDA) produce per-base coverage for a few Brooks loci, paired with the measured Nanopore coverage cropped to each model's predicted genomic span. Feeds notebooks/exoshorkie_coverage_tracks.ipynb (true vs predicted, plotted with tdsone/s2fview, aligned in genomic coordinates).
tdsone
commented
Jun 3, 2026
| CROP_BP = 1024 | ||
| BIN_BP = 16 | ||
| CACHE = Path("/home/tds122/exoshorkie-weights/coverage") | ||
| MENEU_CACHE = Path("/home/tds122/yeast-meneu/data/tasks/meneu_foreign_dna/_cache") |
Member
Author
There was a problem hiding this comment.
Points to personal path, should not land on dev (same in lineabove)
# Conflicts: # src/yeastbench/registry.py
The benchmark should benchmark ready-to-use models, not train/distill them; the distillation pipeline is out of place here.
# Conflicts: # src/yeastbench/registry.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Draft / WIP — getting ExoShorkie (Mandl & Orenstein 2026) into the benchmark.
Done
models/shorkie/nn.py): ExoShorkie = the Shorkie trunk + a linearDense(1)per-bin coverage head (no softplus), squeezed to(B, T_out). Made config-driven so one trunk implementation serves both models; the Shorkie path is bit-for-bit unchanged (full suite + characterization tests pass).tests/test_exoshorkie_head.pychecks exact weight-identity against a real teacher.h5+ a golden regression lock.scripts/exoshorkie/): per genome, the 40-model teacher ensemble is distilled into one PyTorch student, run entirely through the verified port (no TensorFlow).distill_all.shdrives all 6 genomes unattended: download teachers + FASTA → target-gen → train → gate r>0.98 → delete teachers on pass (halts on any failure).Note on the species channel
ExoShorkie's
build_shorkie_featuresfeeds the S. cerevisiae host one-hot at global input column 119 (5-wide DNA incl. N +species[114]), whereas the pretrained Shorkie trunk uses column 114 (4-widedna_1hot, host at 114). It's a consistent shift baked into their reimplementation and applied in both their training and inference, so we feed 119 to run their published weights faithfully. Our port keeps 114 as the Shorkie default (config-gated). Reported upstream.TODO