Skip to content

Commit ac17f0c

Browse files
lucapinellolp698claude
authored
feat: consolidate Enformer + Borzoi + LegNet weights under chorus HF org (#68)
* feat(hf-mirror): consolidate weights for Enformer + Borzoi + LegNet under chorus org Mirror the canonical weights for three more oracles to chorus-controlled HuggingFace repos so the chorus install path doesn't depend on third-party hosts that have shown lifecycle volatility (TFHub deprecation in particular). Each oracle's loader now prefers the chorus mirror and falls back to the original source on any failure — no behavior change in the happy path, redundancy on the unhappy. Mirror repos: - lucapinello/chorus-enformer (~961 MB) Mirrors deepmind/enformer/1, originally on TFHub (now redirecting to Kaggle). saved_model.pb + variables/ at byte-identical contents to the TFHub-cached SavedModel. - lucapinello/chorus-borzoi (~6 GB across 4 folds) Mirrors johahi/borzoi-replicate-{0..3} — the canonical flashzoi PyTorch port — at lucapinello/chorus-borzoi/fold_{0..3}/. Each fold has both model.safetensors and pytorch_model.bin for compatibility. - lucapinello/chorus-legnet (~38 MB) Mirrors the chorus-bundled `chorus_small_legnet.zip` from Zenodo 17863550. Zenodo stays canonical; HF mirror consolidates with the rest of the chorus oracle weights for a single download surface. Each mirror's README follows the alphagenome-pytorch upstream's attribution pattern: separately identifies (a) where the weights came from, (b) who owns the weights, (c) which model terms apply to the weights regardless of mirror, (d) which code license applies to the chorus loader. Original sources retain canonical status; chorus mirrors are durability hedges. Code changes per oracle (HF-first, original-source-fallback): - chorus/oracles/enformer.py default_model_path: "https://tfhub.dev/deepmind/enformer/1" → "lucapinello/chorus-enformer". Adds tfhub_fallback_url for fall-through. - chorus/oracles/enformer_source/templates/{load,predict}_template.py _load_with_recovery and _load helpers detect HF-repo-id-shaped weights and use snapshot_download + tf.saved_model.load; on any failure fall back to the original tfhub_hub.load path. Unwrap the loaded SavedModel's .model attribute when present (TFHub wraps; HF mirror returns the bare graph). - chorus/oracles/borzoi_source/templates/{load,predict}_template.py Add _load_borzoi(fold) helper: snapshot_download the right fold subdir from chorus-borzoi, hand local_path to Borzoi.from_pretrained; fall back to johahi/borzoi-replicate-{fold} on failure. - chorus/oracles/legnet.py Add _try_hf_mirror(dest_zip) helper: hf_hub_download from chorus-legnet, copy to expected zip path; _download_legnet_model tries it first before falling back to download_with_resume of the Zenodo URL. Env yml updates so users have huggingface_hub available in the env for the HF-first path: - environments/chorus-enformer.yml — pip: huggingface_hub>=0.20 - environments/chorus-borzoi.yml — conda: huggingface_hub>=0.20 - environments/chorus-legnet.yml — conda: huggingface_hub>=0.20 Older installs that predate these yml additions don't break — the HF helpers all catch ImportError and fall back to the original source, with a one-line info log so users know what happened. Verified: Enformer HF mirror loads cleanly via tf.saved_model.load on macOS Metal (1.0s post-download). Borzoi + LegNet HF-first paths not yet smoke-tested locally (both fall back cleanly to johahi / Zenodo if HF is unreachable, so the breakage surface is minimal). Fast suite: 368 passed, 0 failed. Sei mirror is in flight as a separate effort — needs ~3.5 GB re-download (the first attempt got truncated at 363 MB). Will land in a follow-up commit. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * feat(hf-mirror): wire chorus-sei mirror with Zenodo fallback Add Sei to the HF-mirror-first weight-resolution pattern landed in this PR for Enformer/Borzoi/LegNet. Same shape: - chorus/oracles/sei.py: new _try_hf_mirror(dest_tar) helper that calls hf_hub_download from lucapinello/chorus-sei and copies the result into the expected zenodo download path. _download_sei_model tries it first; on any failure (no huggingface_hub, network, repo missing), falls back to the existing Zenodo URL with download_with_resume. - environments/chorus-sei.yml: add huggingface_hub>=0.20 conda dep so users get the HF path automatically. Older installs predating this catch ImportError and fall back to Zenodo cleanly. Mirror provenance: - Source: https://zenodo.org/record/4906997 (Troyanskaya lab) - HF: https://huggingface.co/lucapinello/chorus-sei - Size: 3,281,639,255 bytes (verified vs Zenodo metadata) - MD5: 4297aafb711aec4ecccb645b8928ea26 (verified vs Zenodo) - License: CC BY-NC 4.0 (per the original Sei release) The first upload to chorus-sei was based on a truncated curl download (363 MB instead of 3.28 GB — Zenodo's CDN closed the connection mid- stream and curl reported success anyway). The HF mirror now has the full byte-and-md5-verified tarball. README on chorus-sei follows the alphagenome-pytorch attribution pattern: separately identifies (a) where the weights came from, (b) who owns them, (c) which terms govern the weights regardless of mirror, (d) which code license applies to the chorus loader. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * docs(hf-mirror): audit report + README/CHANGELOG entry for the chorus weight mirrors Document the four-mirror consolidation landed in this PR (Enformer, Borzoi, Sei, LegNet) so users can see at a glance where each oracle's weights come from and how the HF-first / original-fallback pattern behaves. audits/2026-04-29_hf_mirror_consolidation/report.md — full inventory: - which oracles got mirrored vs deliberately skipped (and why — AlphaGenome JAX is gated, the upstream PT port is third-party- maintained and shouldn't be cached at chorus or we risk weight drift on upstream upgrades) - upload sizes, md5 verifications (Sei against Zenodo's published `4297aafb711aec4ecccb645b8928ea26`) - the loader pattern (try HF first, catch ImportError + Exception, fall back cleanly to original source) so future mirrors follow the same shape - issues encountered: Zenodo CDN truncated the first Sei download silently (curl exit 0, file 363 MB instead of 3.28 GB); a `curl -C -` resume produced corrupt bytes; clean fix was wipe + fresh download with strict size + md5 verification against the Zenodo API metadata - older chrombpnet env missing huggingface_hub forced the CDF rebuild script to fall back to ENCODE tarballs; HANDOFF.md now has an explicit pip install step - HF CAS dedup made re-uploads instant + initially confusing ("Skipping to prevent empty commit" — verified the file size on HF after the upload to confirm it was really there) README.md — new "Where the oracle weights come from" subsection in the install/setup chapter, with the chorus mirror / original source table for all 6 oracles + 7 weight bundles. Notes that license terms on the weights are unchanged by mirroring (e.g. Sei stays CC-BY-NC). CHANGELOG.md — Unreleased entry above the existing AlphaGenome PT backend entry, with the same four-row table + a one-line summary of the loader pattern. Fast suite still 368 passed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: lp698 <lp698@dimm2fv07n65x.partners.org> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 8c2243c commit ac17f0c

14 files changed

Lines changed: 332 additions & 30 deletions

File tree

CHANGELOG.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,17 @@ project adheres to [Semantic Versioning](https://semver.org/).
88

99
### Added
1010

11+
- **HuggingFace mirror consolidation for Enformer, Borzoi, Sei, and LegNet weights.** Chorus now ships a chorus-controlled HF mirror for each non-AlphaGenome oracle, so the install path doesn't depend on third-party hosts that have shown lifecycle volatility (TFHub deprecation in particular). Each loader prefers the chorus mirror and falls back to the original source on any failure — no behavior change in the happy path, redundancy on the unhappy.
12+
13+
| Oracle | New chorus mirror | Original source |
14+
|---|---|---|
15+
| Enformer | [`lucapinello/chorus-enformer`](https://huggingface.co/lucapinello/chorus-enformer) | TFHub `deepmind/enformer/1` (redirects to Kaggle) |
16+
| Borzoi | [`lucapinello/chorus-borzoi`](https://huggingface.co/lucapinello/chorus-borzoi) (4 folds) | `johahi/borzoi-replicate-{0..3}` |
17+
| Sei | [`lucapinello/chorus-sei`](https://huggingface.co/lucapinello/chorus-sei) | Zenodo `4906997` |
18+
| LegNet | [`lucapinello/chorus-legnet`](https://huggingface.co/lucapinello/chorus-legnet) | Zenodo `17863550` |
19+
20+
Each mirror's README explicitly identifies (a) where the weights came from, (b) who owns the weights, (c) which model terms apply to the weights regardless of mirror, (d) which code license applies to the chorus loader. License terms applying to the *weights* are unchanged by mirroring (Sei stays CC-BY-NC, etc.). `huggingface_hub>=0.20` added to the four corresponding env yamls so users get the HF path by default; older installs predating this fall back to the original source via `ImportError`-handling. Sei tarball verified byte-and-md5 against Zenodo's published `4297aafb711aec4ecccb645b8928ea26`. See `audits/2026-04-29_hf_mirror_consolidation/report.md` for the full inventory and verification trail.
21+
1122
- **AlphaGenome PyTorch backend (second AlphaGenome oracle, installed by default)**
1223
`AlphaGenomePTOracle` wraps the upstream
1324
[`genomicsxai/alphagenome-pytorch`](https://github.com/genomicsxai/alphagenome-pytorch)

README.md

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,22 @@ Chorus converts every raw prediction into an **effect percentile** and **activit
290290

291291
**Nothing to configure.** `chorus setup` pre-downloads the relevant backgrounds for every oracle. If you skipped that step, on the first variant analysis for a given oracle the backgrounds are automatically fetched from the public HuggingFace dataset [`lucapinello/chorus-backgrounds`](https://huggingface.co/datasets/lucapinello/chorus-backgrounds) and cached at `~/.chorus/backgrounds/`.
292292

293+
#### Where the oracle weights come from
294+
295+
Chorus mirrors every oracle's weights to chorus-controlled HuggingFace repos so the install path stays stable even if upstream sources move (TFHub deprecation, Zenodo single-link records, third-party HF accounts). Each loader prefers the chorus mirror and falls back to the original source on any failure — so chorus stays bootable even if the mirror is unreachable.
296+
297+
| Oracle | Chorus mirror | Original source | Mirror size |
298+
|---|---|---|---|
299+
| AlphaGenome (JAX) | not mirrored — official source | [`google/alphagenome-all-folds`](https://huggingface.co/google/alphagenome-all-folds) (gated) ||
300+
| AlphaGenome (PyTorch) | not mirrored — upstream port | [`gtca/alphagenome_pytorch`](https://huggingface.co/gtca/alphagenome_pytorch) ||
301+
| Enformer | [`lucapinello/chorus-enformer`](https://huggingface.co/lucapinello/chorus-enformer) | TFHub `deepmind/enformer/1` (now redirects to Kaggle) | 961 MB |
302+
| Borzoi | [`lucapinello/chorus-borzoi`](https://huggingface.co/lucapinello/chorus-borzoi) | [`johahi/borzoi-replicate-{0..3}`](https://huggingface.co/johahi/borzoi-replicate-0) | ~6 GB (4 folds) |
303+
| ChromBPNet | [`lucapinello/chorus-chrombpnet-slim`](https://huggingface.co/lucapinello/chorus-chrombpnet-slim) | ENCODE per-experiment tarballs | 1.49 GB (786 h5's) |
304+
| Sei | [`lucapinello/chorus-sei`](https://huggingface.co/lucapinello/chorus-sei) | Zenodo [4906997](https://zenodo.org/record/4906997) | 3.28 GB |
305+
| LegNet | [`lucapinello/chorus-legnet`](https://huggingface.co/lucapinello/chorus-legnet) | Zenodo [17863550](https://zenodo.org/records/17863550) | 38 MB |
306+
307+
The chorus mirrors are byte-identical to the originals (verified via md5 / size against upstream metadata where published). License terms applying to the *weights* are unchanged by mirroring — see each mirror's README on HuggingFace for explicit attribution and the upstream model terms.
308+
293309
| Oracle | File size | Tracks covered |
294310
|---|---|---|
295311
| AlphaGenome | ~260 MB | 5,168 |
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# HF mirror consolidation — chorus-controlled mirrors for Enformer, Borzoi, Sei, LegNet
2+
3+
**Date**: 2026-04-29
4+
**Branch**: `feat/hf-mirror-consolidation` → PR [#68](https://github.com/pinellolab/chorus/pull/68)
5+
**Trigger**: chorus's install path was depending on third-party hosts that have shown lifecycle volatility — most pressingly TFHub (deprecated by Google in 2024, redirecting to Kaggle Models). Plus `johahi/borzoi-replicate-*` and Zenodo single-link weights. User asked to consolidate weight provenance under the chorus org for resilience.
6+
7+
## Inventory after consolidation
8+
9+
| Oracle | Mirror (chorus) | Original source | Size | Verified |
10+
|---|---|---|---|---|
11+
| **Enformer** | [`lucapinello/chorus-enformer`](https://huggingface.co/lucapinello/chorus-enformer) | TFHub `deepmind/enformer/1` (now redirects to Kaggle) | 961 MB | ✅ load round-trip on macOS Metal |
12+
| **Borzoi** | [`lucapinello/chorus-borzoi`](https://huggingface.co/lucapinello/chorus-borzoi) (4 folds) | [`johahi/borzoi-replicate-{0..3}`](https://huggingface.co/johahi/borzoi-replicate-0) | ~6 GB | ✅ upload round-trip |
13+
| **Sei** | [`lucapinello/chorus-sei`](https://huggingface.co/lucapinello/chorus-sei) | Zenodo [4906997](https://zenodo.org/record/4906997) | 3,281,639,255 bytes | ✅ md5 `4297aafb711aec4ecccb645b8928ea26` matches Zenodo metadata |
14+
| **LegNet** | [`lucapinello/chorus-legnet`](https://huggingface.co/lucapinello/chorus-legnet) | Zenodo [17863550](https://zenodo.org/records/17863550) (chorus-pinned) | ~38 MB | ✅ size match |
15+
| **ChromBPNet** (default fold-0 nobias) | [`lucapinello/chorus-chrombpnet-slim`](https://huggingface.co/lucapinello/chorus-chrombpnet-slim) | ENCODE per-experiment tarballs | 1.49 GB (786 h5's) | already shipped in 0.3.0 |
16+
| **AlphaGenome JAX** | `google/alphagenome-all-folds` (gated, official) | Google DeepMind | ~700 MB sharded | not mirrored — official source |
17+
| **AlphaGenome PT** | `gtca/alphagenome_pytorch` (third-party, public) | Conversion of JAX checkpoint | 878 MB | not mirrored — upstream port maintainer; mirroring would risk weight drift if upstream updates |
18+
| **Per-track CDFs** | [`lucapinello/chorus-backgrounds`](https://huggingface.co/datasets/lucapinello/chorus-backgrounds) | chorus-built | ~2 GB across 6 oracles | already shipped |
19+
20+
## What's mirrored vs not
21+
22+
**Mirrored** (chorus-controlled durability):
23+
- Enformer / Borzoi / Sei / LegNet weights — under `lucapinello/chorus-*`
24+
- ChromBPNet slim mirror — under `lucapinello/chorus-chrombpnet-slim` (predates this consolidation)
25+
- Per-track CDFs — under `lucapinello/chorus-backgrounds` (predates this consolidation)
26+
27+
**Deliberately NOT mirrored**:
28+
- **AlphaGenome JAX** (`google/alphagenome-all-folds`): Google DeepMind is the authoritative source; the model is gated under non-commercial terms and chorus shouldn't host a copy that bypasses Google's gating mechanism.
29+
- **AlphaGenome PT** (`gtca/alphagenome_pytorch`): the upstream port author (gtca) maintains the conversion and may update it. Mirroring under chorus would risk users on stale weights when the port upgrades. The non-commercial terms still apply regardless of mirror — see PR #62 README.
30+
31+
## Loader pattern (every chorus mirror)
32+
33+
Each oracle's loader prefers the chorus mirror and falls back to the original source on any failure:
34+
35+
```python
36+
def _try_hf_mirror(...):
37+
try:
38+
from huggingface_hub import hf_hub_download # or snapshot_download
39+
local = hf_hub_download(repo_id="lucapinello/chorus-<oracle>", filename=...)
40+
# use local
41+
return True
42+
except ImportError:
43+
# huggingface_hub not in env — log + fall back
44+
return False
45+
except Exception as exc:
46+
# network / repo missing / HF unreachable — log + fall back
47+
return False
48+
49+
def _download_<oracle>_model(self):
50+
if not self._try_hf_mirror(...):
51+
# Fall back to existing TFHub / johahi / Zenodo / etc. path
52+
...
53+
```
54+
55+
This means:
56+
- Happy path: chorus mirror is fast + chorus-stable.
57+
- HF-down: fall back to original source. No breakage.
58+
- `huggingface_hub` not installed: fall back. Older installs predating this PR continue to work via the original source.
59+
60+
## Attribution + licensing
61+
62+
Each mirror README explicitly identifies four things separately:
63+
64+
1. **Where the weights came from** (original source URL + paper citation)
65+
2. **Who owns the weights** (upstream lab / company)
66+
3. **Which model terms apply to the weights regardless of mirror** (the mirror does NOT override license terms — e.g. Sei stays CC-BY-NC even on HF)
67+
4. **Which code license applies to the chorus loader** (chorus's Apache 2.0)
68+
69+
This format follows the [`gtca/alphagenome_pytorch`](https://huggingface.co/gtca/alphagenome_pytorch) attribution template the user pointed to. Concrete language used per mirror:
70+
71+
> The weights were copied from [original source]. Those weights were created by [upstream] and are the property of [upstream owner]. The model parameters, outputs, and any derivatives thereof remain subject to [upstream license/terms]. The chorus loader code is released under [chorus license]. These terms are consistent with the terms for the reference code and the model weights.
72+
73+
## Issues encountered + resolutions
74+
75+
### 1. Sei truncated download (1× retry needed)
76+
77+
First curl of the Sei tarball completed with exit 0 but only 363 MB landed (vs 3.28 GB expected). Curl didn't detect Zenodo's mid-stream connection close as a failure. Caught when the user noted the size felt small.
78+
79+
**Fix**: re-download with strict size + md5 verification against Zenodo's API metadata (`https://zenodo.org/api/records/4906997` exposes `size` and `checksum: md5:...`). The second download (using the API content URL `/files/.../content` instead of the public file URL) landed clean.
80+
81+
**Note**: a `curl -C -` resume attempt in between produced a corrupt 3.29 GB file (8.5 MB longer than expected, totally different md5) — likely byte-overlap from the resume offset. **Don't trust `curl -C -` resumes for long-running mid-stream-truncated downloads** without re-verifying. Better path: delete and re-download fresh, then verify.
82+
83+
### 2. `huggingface_hub` not in older chrombpnet env (CDF build path)
84+
85+
The CDF rebuild script falls back to ENCODE tarballs (~700 MB per model × 786 models = 100s of GB) instead of the slim mirror (25 MB per model) when `huggingface_hub` isn't importable. PR #60 added it to `chorus-chrombpnet.yml`, but envs created before that PR don't have it.
86+
87+
**Fix in HANDOFF.md** (PR #67): added an explicit `pip install` step at the top of the CUDA-box runbook so older envs work. Caught this on the M3 Ultra during a sanity run before handing off.
88+
89+
### 3. Duplicate file detection on HF re-upload
90+
91+
When re-uploading the corrected Sei tarball, the `HfApi.upload_file` reported "Skipping to prevent empty commit" — initially confusing because the file content had changed. Verified via `repo_info(files_metadata=True)` that the actual file size on HF matched the new 3.28 GB (HF's CAS dedup pool just happened to have the new content cached).
92+
93+
## Time + bandwidth
94+
95+
Total upload to HF over residential connection:
96+
- Enformer: 16 s for 1 GB
97+
- LegNet: 6 s for 38 MB
98+
- Borzoi: 114 s for 6 GB across 4 folds
99+
- Sei: ~6 s for 3.28 GB (HF's CAS dedup made this fast)
100+
101+
Total mirror size on HF: ~11.3 GB across 4 new repos.
102+
103+
## Follow-ups
104+
105+
- **CDF rebuild on CUDA box** — separate workflow, in flight on the user's lab GPU, gated by the CUDA agent. Output is a fresh `chrombpnet_pertrack.npz` rebuild against `chrombpnet_nobias` (the 0.3.0 default), uploaded to the existing `lucapinello/chorus-backgrounds` dataset.
106+
- **Smoke-test Borzoi / LegNet HF-first paths** locally — Enformer was verified end-to-end via `tf.saved_model.load`; the other two have only the upload round-trip and the fall-back path verified. Adding integration smoke tests is a small follow-up if the equivalence-test infra is reused.
107+
- **Document mirror pattern in CONTRIBUTING / docs** — the HF-first / original-fallback pattern should be the standard for any future weight-fetching code.
108+
109+
## Critical files modified in PR #68
110+
111+
- `chorus/oracles/enformer.py``default_model_path` flipped to HF repo id
112+
- `chorus/oracles/enformer_source/templates/{load,predict}_template.py` — HF path detection + fallback
113+
- `chorus/oracles/borzoi_source/templates/{load,predict}_template.py``_load_borzoi(fold)` helper
114+
- `chorus/oracles/sei.py``_try_hf_mirror` helper
115+
- `chorus/oracles/legnet.py``_try_hf_mirror` helper
116+
- `environments/chorus-{enformer,borzoi,sei,legnet}.yml``huggingface_hub>=0.20`

chorus/oracles/borzoi_source/templates/load_template.py

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,33 @@
1-
import json
2-
import torch
1+
import json
2+
import os
3+
import torch
34
from borzoi_pytorch import Borzoi
45

5-
with open("__ARGS_FILE_NAME__") as inp: # to be formatted by calling script
6+
with open("__ARGS_FILE_NAME__") as inp: # to be formatted by calling script
67
args = json.load(inp)
78

8-
flashzoi = Borzoi.from_pretrained(f'johahi/borzoi-replicate-{args["fold"]}')
9+
10+
def _load_borzoi(fold: int):
11+
"""Prefer the chorus-controlled HF mirror at lucapinello/chorus-borzoi
12+
(per-fold subdirs); fall back to johahi/borzoi-replicate-{fold} on
13+
any failure. Both contain the same flashzoi PyTorch weights."""
14+
try:
15+
from huggingface_hub import snapshot_download
16+
local_dir = snapshot_download(
17+
repo_id="lucapinello/chorus-borzoi",
18+
repo_type="model",
19+
allow_patterns=[f"fold_{fold}/*"],
20+
)
21+
fold_path = os.path.join(local_dir, f"fold_{fold}")
22+
if not os.path.isdir(fold_path):
23+
raise FileNotFoundError(f"fold_{fold}/ missing in chorus-borzoi mirror")
24+
return Borzoi.from_pretrained(fold_path)
25+
except Exception as exc:
26+
print(f"chorus-borzoi mirror unavailable ({exc}); falling back to johahi")
27+
return Borzoi.from_pretrained(f'johahi/borzoi-replicate-{fold}')
28+
29+
30+
flashzoi = _load_borzoi(args["fold"])
931
device = args['device']
1032
if device is None or device == 'auto':
1133
if torch.cuda.is_available():
@@ -25,4 +47,4 @@
2547
'model_class': str(type(flashzoi)),
2648
'description': 'Borzoi model loaded successfully',
2749
'device': args['device']
28-
}
50+
}

chorus/oracles/borzoi_source/templates/predict_template.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,28 @@
1212
args = json.load(inp)
1313

1414
from borzoi_pytorch import Borzoi
15-
flashzoi = Borzoi.from_pretrained(f'johahi/borzoi-replicate-{args["fold"]}')
15+
16+
17+
def _load_borzoi(fold: int):
18+
"""Same HF-mirror-first loader as load_template.py — prefer the
19+
chorus-controlled mirror at lucapinello/chorus-borzoi, fall back to
20+
johahi's original repos."""
21+
try:
22+
from huggingface_hub import snapshot_download
23+
local_dir = snapshot_download(
24+
repo_id="lucapinello/chorus-borzoi",
25+
repo_type="model",
26+
allow_patterns=[f"fold_{fold}/*"],
27+
)
28+
fold_path = os.path.join(local_dir, f"fold_{fold}")
29+
if not os.path.isdir(fold_path):
30+
raise FileNotFoundError(f"fold_{fold}/ missing in chorus-borzoi mirror")
31+
return Borzoi.from_pretrained(fold_path)
32+
except Exception:
33+
return Borzoi.from_pretrained(f'johahi/borzoi-replicate-{fold}')
34+
35+
36+
flashzoi = _load_borzoi(args["fold"])
1637

1738
device = args['device']
1839
if device is None or device == 'auto':

chorus/oracles/enformer.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,8 +89,14 @@ def __init__(self,
8989
self._enformer_model = None
9090
self._track_dict = None
9191

92-
# Default model path
93-
self.default_model_path = "https://tfhub.dev/deepmind/enformer/1"
92+
# Default model path. Chorus mirrors the original DeepMind
93+
# SavedModel at huggingface.co/lucapinello/chorus-enformer for
94+
# resilience — the original TFHub URL now redirects through
95+
# Kaggle. The HF mirror is byte-identical to the TFHub-cached
96+
# version (saved_model.pb + variables/) and chorus prefers it
97+
# at load time.
98+
self.default_model_path = "lucapinello/chorus-enformer"
99+
self.tfhub_fallback_url = "https://tfhub.dev/deepmind/enformer/1"
94100

95101
# Reference genome
96102
self.reference_fasta = reference_fasta

chorus/oracles/enformer_source/templates/load_template.py

Lines changed: 41 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
import tensorflow as tf
22
import tensorflow_hub as hub
33
import os
4-
import json
4+
import json
55

6-
with open("__ARGS_FILE_NAME__") as inp: # to be formatted by calling script
6+
with open("__ARGS_FILE_NAME__") as inp: # to be formatted by calling script
77
args = json.load(inp)
88

99
# Set TFHub progress tracking
@@ -32,12 +32,35 @@
3232
else:
3333
print("No GPU detected, using CPU")
3434

35-
# Load the model. If tfhub's on-disk cache is corrupt (incomplete download
36-
# from a previous session — missing saved_model.pb), hub.load raises
37-
# "contains neither 'saved_model.pb' nor 'saved_model.pbtxt'". Detect this,
38-
# clear the bad cache directory, and retry once.
39-
def _load_with_tfhub_recovery(weights: str):
35+
36+
# Load the model. Chorus prefers the HF mirror at
37+
# `lucapinello/chorus-enformer` (mirror of the original DeepMind
38+
# SavedModel) for resilience — the original TFHub URL now redirects
39+
# through Kaggle and the migration trail has caused breakage. We fall
40+
# back to the original TFHub URL if the HF mirror is unreachable.
41+
def _load_with_recovery(weights: str):
42+
"""Load via huggingface_hub when weights is a HF repo id, else via
43+
tensorflow_hub. On any failure for the HF path, fall back to the
44+
original TFHub URL the load template was designed for. On TFHub
45+
cache corruption (incomplete download from a previous session —
46+
missing saved_model.pb), clear the bad cache directory and retry."""
4047
import re, shutil
48+
49+
# HF repo path: not a URL, looks like "user/repo-name".
50+
if "/" in weights and not weights.startswith("http"):
51+
try:
52+
from huggingface_hub import snapshot_download
53+
local_dir = snapshot_download(
54+
repo_id=weights,
55+
repo_type="model",
56+
allow_patterns=["saved_model.pb", "variables/*"],
57+
)
58+
print(f"Loading Enformer SavedModel from HF mirror at {{local_dir}}")
59+
return tf.saved_model.load(local_dir)
60+
except Exception as exc:
61+
print(f"HF mirror load failed ({{exc}}); falling back to TFHub")
62+
weights = "https://tfhub.dev/deepmind/enformer/1"
63+
4164
try:
4265
return hub.load(weights)
4366
except Exception as exc:
@@ -53,9 +76,16 @@ def _load_with_tfhub_recovery(weights: str):
5376
shutil.rmtree(bad_dir, ignore_errors=True)
5477
return hub.load(weights)
5578

56-
enformer = _load_with_tfhub_recovery(args['model_weights'])
57-
# Get the actual model from the enformer object
58-
model = enformer.model
79+
80+
loaded = _load_with_recovery(args['model_weights'])
81+
82+
# The HF SavedModel is the *underlying* tf.saved_model.load output, while
83+
# the TFHub `hub.load` wraps it in an "enformer" object with a `.model`
84+
# attribute. Detect which we got and unwrap.
85+
if hasattr(loaded, "model"):
86+
model = loaded.model
87+
else:
88+
model = loaded
5989

6090
# Get device info
6191
if device == 'cpu' or not tf.config.list_physical_devices('GPU'):
@@ -70,4 +100,4 @@ def _load_with_tfhub_recovery(weights: str):
70100
'has_predict': hasattr(model, 'predict_on_batch'),
71101
'description': 'Enformer model loaded successfully',
72102
'device': actual_device
73-
}
103+
}

0 commit comments

Comments
 (0)