Add ZeroGPU-enabled BioLM-NET workbench

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+.pytest_cache/
+.venv/
+*.py[cod]
+*.zip
+*.csv
+!tests/fixtures/*.csv

LICENSE

@@ -0,0 +1,94 @@
+Creative Commons Attribution-NonCommercial 4.0 International Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.
+
+Section 1 – Definitions.
+
+    Adapted Material means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image.
+    Adapter's License means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License.
+    Copyright and Similar Rights means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights.
+    Effective Technological Measures means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements.
+    Exceptions and Limitations means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material.
+    Licensed Material means the artistic or literary work, database, or other material to which the Licensor applied this Public License.
+    Licensed Rights means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license.
+    Licensor means the individual(s) or entity(ies) granting rights under this Public License.
+    NonCommercial means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange.
+    Share means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them.
+    Sui Generis Database Rights means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world.
+    You means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning.
+
+Section 2 – Scope.
+
+    License grant.
+        Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to:
+            reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and
+            produce, reproduce, and Share Adapted Material for NonCommercial purposes only.
+        Exceptions and Limitations. For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions.
+        Term. The term of this Public License is specified in Section 6(a).
+        Media and formats; technical modifications allowed. The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material.
+        Downstream recipients.
+            Offer from the Licensor – Licensed Material. Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License.
+            No downstream restrictions. You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material.
+        No endorsement. Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i).
+
+    Other rights.
+        Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise.
+        Patent and trademark rights are not licensed under this Public License.
+        To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes.
+
+Section 3 – License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the following conditions.
+
+    Attribution.
+
+        If You Share the Licensed Material (including in modified form), You must:
+            retain the following if it is supplied by the Licensor with the Licensed Material:
+                identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated);
+                a copyright notice;
+                a notice that refers to this Public License;
+                a notice that refers to the disclaimer of warranties;
+                a URI or hyperlink to the Licensed Material to the extent reasonably practicable;
+            indicate if You modified the Licensed Material and retain an indication of any previous modifications; and
+            indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License.
+        You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information.
+        If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable.
+        If You Share Adapted Material You produce, the Adapter's License You apply must not prevent recipients of the Adapted Material from complying with this Public License.
+
+Section 4 – Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material:
+
+    for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only;
+    if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material; and
+    You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights.
+
+Section 5 – Disclaimer of Warranties and Limitation of Liability.
+
+    Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.
+    To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.
+
+    The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability.
+
+Section 6 – Term and Termination.
+
+    Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates:
+        automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or
+        upon express reinstatement by the Licensor.
+    For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License.
+    For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License.
+    Sections 1, 5, 6, 7, and 8 survive termination of this Public License.
+
+Section 7 – Other Terms and Conditions.
+
+    The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed.
+    Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License.
+
+Section 8 – Interpretation.
+
+    For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License.
+    To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions.
+    No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor.
+    Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority.

README.md

@@ -1,13 +1,127 @@
 ---
-title: BioLM NET Workbench
-emoji: 😻
-colorFrom: purple
-colorTo: red
+title: BioLM-NET Workbench
+emoji: 🧬
+colorFrom: green
+colorTo: yellow
 sdk: gradio
-sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
+python_version: "3.12"
+suggested_hardware: zero-a10g
+license: cc-by-nc-4.0
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# BioLM-NET Workbench
+
+BioLM-NET Workbench is a guided Hugging Face Spaces interface for assembling,
+training, exporting, and applying the interpretable dual-omics architecture
+described by Rifat et al.
+
+The interface supports:
+
+- the BRCA, COAD, GBM, and scTrioseq2 datasets from
+  [`bozdaglab/BioLM-NET`](https://github.com/bozdaglab/BioLM-NET);
+- a GitHub dataset folder containing the standard BioLM-NET filenames;
+- user uploads for paired gene-expression and DNA-methylation matrices, labels,
+  and gene-to-pathway mappings;
+- automatic retrieval of the upstream `PDI.csv` and `PPI.csv`, with optional
+  custom overrides;
+- paper-defined trainable biological masking (`W_masked = W ⊙ M`);
+- GenePT-guided, membership-masked pathway attention;
+- on-demand ZeroGPU allocation only while model training runs;
+- stratified validation, balanced loss, downloadable model artifacts, and a
+  separate prediction workflow.
+
+## Data contract
+
+Samples are rows. Omics features are columns named with HGNC gene symbols.
+Paired omics files and the label file must have the same row order and count.
+
+| File | Required columns / orientation |
+| --- | --- |
+| `Gene_Expression.csv` | samples × gene-symbol columns |
+| `DNA_Methylation.csv` | samples × gene-symbol columns |
+| `label.csv` | exactly one label column |
+| `PDI.csv` | `TF`, `Target` |
+| `PPI.csv` | `protein1`, `protein2`, `combined_score` |
+| gene/pathway mapping | `SYMBOL`, `PathwayID` |
+
+For a GitHub folder, the app expects:
+
+```text
+Gene_Expression.csv
+DNA_Methylation.csv
+label.csv
+ge_target_to_KEGG_significant.csv
+dna_target_to_KEGG_significant.csv
+```
+
+If pathway files are already the significant results of an enrichment analysis,
+leave the checkbox enabled. For a complete gene/pathway annotation catalog,
+disable it; the app will run over-representation testing and retain pathways
+with Benjamini-Hochberg adjusted p-value below 0.05.
+
+## Architecture choices
+
+The implementation follows the paper's mathematical description:
+
+1. PDI edges are binary TF-to-target connections.
+2. STRING PPI scores are normalized, filtered above 0.7, and restricted to the
+   top score decile. Because PPI is undirected, an input gene can connect to its
+   partner from either endpoint.
+3. The PDI/PPI matrix masks the trainable first-layer weights directly.
+4. Only hidden genes represented by the selected GenePT file are retained.
+5. Gene-to-pathway attention scores are masked to actual memberships before the
+   softmax; nonmember genes cannot receive pathway attention.
+6. Gene-expression and DNA-methylation pathway projections are fused for
+   multiclass prediction.
+
+Training scalers are fit only on the training partition and are included in the
+exported artifact. Model tensors use `safetensors`; the artifact loader reads
+only the four expected files and does not unpickle user content.
+
+## Run locally
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements-local.txt
+python app.py
+```
+
+For local development, install `requirements-local.txt` instead. Hugging Face's
+Gradio runtime supplies the platform-managed `spaces` package automatically;
+the local requirements file installs its no-op implementation so the same app
+also runs on CPU or a conventional GPU.
+
+## Deploy to Hugging Face
+
+Create a Gradio Space, push every file in this folder to the Space repository,
+and select **ZeroGPU** in **Settings → Hardware**. The Space runtime installs
+`requirements.txt` and starts `app.py`. The training callback uses a dynamic
+30–300 second reservation based on dataset size, biological graph size, batch
+size, and epochs. Data preparation and prediction remain on CPU and do not
+consume GPU quota.
+
+ZeroGPU users have tier-specific daily quotas and may queue during busy periods.
+The epoch control is capped at 200 so paper-scale runs can fit within a single
+free-tier reservation when the workload permits. If a cohort cannot finish in
+five minutes, reduce the epochs or switch the Space to dedicated hardware.
+
+## Scientific and licensing note
+
+This is research software, not a clinical decision system. Users remain
+responsible for cohort design, preprocessing, leakage checks, external
+validation, and interpretation.
+
+BioLM-NET source and datasets are provided by the Bozdag Lab under
+[CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/). Cite:
+
+> Rifat, J. I. M. et al. “BioLM-NET: an interpretable deep learning model
+> combining prior biological knowledge and contextual LLM gene embeddings on
+> multi-omics data to predict disease.” *Pacific Symposium on Biocomputing
+> 2026* (2025).
+
+GenePT embeddings are retrieved from
+[`honicky/genept-composable-embeddings`](https://huggingface.co/honicky/genept-composable-embeddings);
+review that repository's license and terms for your intended use.

app.py

@@ -0,0 +1,984 @@
+from __future__ import annotations
+
+import html
+import math
+import tempfile
+from functools import lru_cache
+from pathlib import Path
+
+import spaces
+import gradio as gr
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+
+from biolmnet.artifacts import load_bundle, save_bundle
+from biolmnet.data import (
+    PreparedWorkspace,
+    attach_embeddings_and_pathways,
+    build_biological_mask,
+    github_dataset_sources,
+    load_genept_embeddings,
+    read_csv,
+    upstream_example_sources,
+    upstream_interaction_sources,
+    validate_and_align_omics,
+)
+from biolmnet.training import (
+    Hyperparameters,
+    ModelBundle,
+    pathway_importance,
+    predict,
+    train,
+)
+
+
+GENEPT_OPTIONS = {
+    "Auto — bulk or single-cell based on source": "auto",
+    "Bulk · large-3 context": "embedding_original_large_3.parquet",
+    "Bulk · original ada-text context": "embedding_original_ada_text.parquet",
+    "Single-cell · cell type, tissue, drug & pathway": (
+        "embedding_associations_cell_type_tissue_drug_pathway_openai_large.parquet"
+    ),
+    "Single-cell · age, cell type, drugs & pathways": (
+        "embedding_associations_age_cell_type_drugs_pathways_openai_large.parquet"
+    ),
+}
+
+
+CSS = """
+:root {
+  --ink: #102824;
+  --muted: #60726e;
+  --line: #d7e2de;
+  --mint: #e8f4ee;
+  --leaf: #167c5a;
+  --leaf-dark: #0e5b43;
+  --amber: #e5a63c;
+}
+.gradio-container {
+  max-width: 1260px !important;
+  margin: 0 auto !important;
+  background:
+    radial-gradient(circle at 92% 4%, rgba(209, 235, 224, .7), transparent 26rem),
+    #f8fbf9 !important;
+  color: var(--ink) !important;
+}
+.biolm-hero {
+  position: relative;
+  overflow: hidden;
+  padding: 34px 36px 30px;
+  margin: 12px 0 18px;
+  border: 1px solid #cfe0d9;
+  border-radius: 24px;
+  background: linear-gradient(135deg, #0b2f28 0%, #124b3d 70%, #17694f 100%);
+  box-shadow: 0 18px 50px rgba(19, 63, 52, .12);
+}
+.biolm-hero::after {
+  content: "";
+  position: absolute;
+  width: 250px;
+  height: 250px;
+  right: -60px;
+  top: -90px;
+  border: 1px solid rgba(255,255,255,.2);
+  border-radius: 50%;
+  box-shadow: 0 0 0 34px rgba(255,255,255,.035), 0 0 0 68px rgba(255,255,255,.025);
+}
+.eyebrow {
+  color: #a9e2ca;
+  font: 700 12px/1.2 ui-monospace, SFMono-Regular, Menlo, monospace;
+  letter-spacing: .14em;
+  text-transform: uppercase;
+}
+.biolm-hero h1 {
+  color: white;
+  font-size: clamp(34px, 5vw, 58px);
+  line-height: .98;
+  letter-spacing: -.045em;
+  margin: 12px 0 14px;
+}
+.biolm-hero p {
+  max-width: 780px;
+  color: #d6e8e1;
+  font-size: 17px;
+  line-height: 1.55;
+  margin: 0;
+}
+.hero-meta {
+  display: flex;
+  gap: 10px;
+  flex-wrap: wrap;
+  margin-top: 22px;
+}
+.hero-chip {
+  color: #e9f7f1;
+  border: 1px solid rgba(255,255,255,.23);
+  background: rgba(255,255,255,.07);
+  border-radius: 999px;
+  padding: 7px 11px;
+  font: 600 12px/1 ui-monospace, SFMono-Regular, Menlo, monospace;
+}
+.phase-card {
+  border: 1px solid var(--line) !important;
+  border-radius: 18px !important;
+  background: rgba(255,255,255,.86) !important;
+  box-shadow: 0 9px 26px rgba(27, 68, 57, .05) !important;
+}
+.phase-intro {
+  border-left: 3px solid var(--leaf);
+  padding: 2px 0 2px 15px;
+  color: var(--muted);
+}
+.phase-intro strong { color: var(--ink); }
+.status-box {
+  border-radius: 15px;
+  padding: 14px 16px;
+  background: var(--mint);
+  border: 1px solid #cbe3d8;
+  color: #254d41;
+}
+.error-box {
+  border-radius: 15px;
+  padding: 14px 16px;
+  background: #fff1ee;
+  border: 1px solid #f0cac2;
+  color: #793b31;
+}
+.metric-grid {
+  display: grid;
+  grid-template-columns: repeat(4, minmax(0,1fr));
+  gap: 10px;
+}
+.metric {
+  padding: 13px 14px;
+  border: 1px solid #d4e2dd;
+  border-radius: 14px;
+  background: white;
+}
+.metric span {
+  display:block;
+  color: var(--muted);
+  font-size: 11px;
+  font-weight: 700;
+  letter-spacing: .07em;
+  text-transform: uppercase;
+}
+.metric b {
+  display:block;
+  margin-top: 3px;
+  color: var(--ink);
+  font-size: 22px;
+}
+button.primary {
+  background: var(--leaf) !important;
+  border-color: var(--leaf) !important;
+}
+button.primary:hover { background: var(--leaf-dark) !important; }
+.footnote {
+  color: #6d7e79;
+  font-size: 12px;
+  line-height: 1.55;
+}
+@media (max-width: 720px) {
+  .biolm-hero { padding: 26px 22px; border-radius: 18px; }
+  .metric-grid { grid-template-columns: repeat(2, minmax(0,1fr)); }
+}
+"""
+
+
+def _status(message: str, error: bool = False) -> str:
+    class_name = "error-box" if error else "status-box"
+    return f'<div class="{class_name}">{html.escape(message)}</div>'
+
+
+def _source_visibility(mode: str):
+    return (
+        gr.update(visible=mode == "BioLM-NET examples"),
+        gr.update(visible=mode == "GitHub folder"),
+        gr.update(visible=mode == "Upload files"),
+    )
+
+
+@lru_cache(maxsize=1)
+def _upstream_interactions() -> tuple[pd.DataFrame, pd.DataFrame]:
+    pdi_url, ppi_url = upstream_interaction_sources()
+    return read_csv(pdi_url), read_csv(ppi_url)
+
+
+def _read_required_upload(path: str | None, label: str) -> pd.DataFrame:
+    if not path:
+        raise ValueError(f"Upload {label}.")
+    return read_csv(path)
+
+
+def _resolve_embedding_file(
+    option: str, source_mode: str, example_dataset: str
+) -> str:
+    selected = GENEPT_OPTIONS[option]
+    if selected != "auto":
+        return selected
+    if source_mode == "BioLM-NET examples" and example_dataset == "scTrioseq2":
+        return (
+            "embedding_associations_cell_type_tissue_drug_pathway_"
+            "openai_large.parquet"
+        )
+    return "embedding_original_large_3.parquet"
+
+
+def prepare_workspace(
+    source_mode: str,
+    example_dataset: str,
+    github_folder: str,
+    uploaded_gene: str | None,
+    uploaded_dna: str | None,
+    uploaded_labels: str | None,
+    uploaded_gene_pathway: str | None,
+    uploaded_dna_pathway: str | None,
+    uploaded_pdi: str | None,
+    uploaded_ppi: str | None,
+    pathway_files_are_significant: bool,
+    embedding_option: str,
+    progress=gr.Progress(track_tqdm=False),
+):
+    try:
+        progress(0.03, desc="Reading omics data")
+        if source_mode == "BioLM-NET examples":
+            sources = upstream_example_sources(example_dataset)
+            frames = {name: read_csv(url) for name, url in sources.items()}
+            source_name = f"BioLM-NET / {example_dataset}"
+            precomputed_significant = True
+            allow_preset_trim = True
+        elif source_mode == "GitHub folder":
+            if not github_folder.strip():
+                raise ValueError("Enter a GitHub dataset folder URL.")
+            sources = github_dataset_sources(github_folder)
+            frames = {name: read_csv(url) for name, url in sources.items()}
+            source_name = github_folder.strip()
+            precomputed_significant = pathway_files_are_significant
+            allow_preset_trim = False
+        else:
+            frames = {
+                "gene": _read_required_upload(
+                    uploaded_gene, "Gene_Expression.csv"
+                ),
+                "dna": _read_required_upload(
+                    uploaded_dna, "DNA_Methylation.csv"
+                ),
+                "labels": _read_required_upload(
+                    uploaded_labels, "label.csv"
+                ),
+                "gene_pathways": _read_required_upload(
+                    uploaded_gene_pathway,
+                    "the gene-expression pathway mapping CSV",
+                ),
+                "dna_pathways": _read_required_upload(
+                    uploaded_dna_pathway,
+                    "the DNA-methylation pathway mapping CSV",
+                ),
+            }
+            source_name = "Uploaded dataset"
+            precomputed_significant = pathway_files_are_significant
+            allow_preset_trim = False
+
+        (
+            gene_frame,
+            dna_frame,
+            labels,
+            label_names,
+            warnings,
+        ) = validate_and_align_omics(
+            frames["gene"],
+            frames["dna"],
+            frames["labels"],
+            allow_preset_trim=allow_preset_trim,
+        )
+
+        progress(0.18, desc="Loading PDI and PPI priors")
+        if uploaded_pdi or uploaded_ppi:
+            if not uploaded_pdi or not uploaded_ppi:
+                raise ValueError(
+                    "To override the repository priors, upload both PDI and PPI files."
+                )
+            pdi_frame = read_csv(uploaded_pdi)
+            ppi_frame = read_csv(uploaded_ppi)
+        else:
+            pdi_frame, ppi_frame = _upstream_interactions()
+
+        progress(0.38, desc="Constructing sparse biological masks")
+        gene_branch = build_biological_mask(
+            list(gene_frame.columns), pdi_frame, ppi_frame
+        )
+        dna_branch = build_biological_mask(
+            list(dna_frame.columns), pdi_frame, ppi_frame
+        )
+
+        embedding_file = _resolve_embedding_file(
+            embedding_option, source_mode, example_dataset
+        )
+        progress(0.53, desc="Retrieving GenePT embeddings")
+        embeddings = load_genept_embeddings(embedding_file)
+
+        progress(0.72, desc="Building enriched pathway connections")
+        gene_enrichment = attach_embeddings_and_pathways(
+            gene_branch,
+            embeddings,
+            frames["gene_pathways"],
+            precomputed_significant=precomputed_significant,
+        )
+        dna_enrichment = attach_embeddings_and_pathways(
+            dna_branch,
+            embeddings,
+            frames["dna_pathways"],
+            precomputed_significant=precomputed_significant,
+        )
+        workspace = PreparedWorkspace(
+            gene_expression=gene_frame.to_numpy(dtype="float32"),
+            dna_methylation=dna_frame.to_numpy(dtype="float32"),
+            labels=labels,
+            label_names=label_names,
+            gene_branch=gene_branch,
+            dna_branch=dna_branch,
+            source_name=source_name,
+            warnings=warnings,
+        )
+
+        architecture = pd.DataFrame(
+            [
+                {
+                    "branch": "Gene expression",
+                    "samples": len(gene_frame),
+                    "input genes": len(gene_branch.input_genes),
+                    "PDI edges": gene_branch.pdi_edges,
+                    "PPI edges": gene_branch.ppi_edges,
+                    "hidden genes": len(gene_branch.hidden_genes),
+                    "pathways": len(gene_branch.pathways),
+                    "mask density": (
+                        gene_branch.biological_mask.astype(bool).mean()
+                    ),
+                },
+                {
+                    "branch": "DNA methylation",
+                    "samples": len(dna_frame),
+                    "input genes": len(dna_branch.input_genes),
+                    "PDI edges": dna_branch.pdi_edges,
+                    "PPI edges": dna_branch.ppi_edges,
+                    "hidden genes": len(dna_branch.hidden_genes),
+                    "pathways": len(dna_branch.pathways),
+                    "mask density": (
+                        dna_branch.biological_mask.astype(bool).mean()
+                    ),
+                },
+            ]
+        )
+        enrichments = pd.concat(
+            [
+                gene_enrichment.assign(branch="Gene expression"),
+                dna_enrichment.assign(branch="DNA methylation"),
+            ],
+            ignore_index=True,
+        )
+        warning_text = (
+            "<br><small>" + " · ".join(html.escape(item) for item in warnings) + "</small>"
+            if warnings
+            else ""
+        )
+        summary = (
+            '<div class="status-box"><strong>Architecture ready.</strong> '
+            f"{len(gene_frame):,} paired samples · {len(label_names)} classes · "
+            f"{len(gene_branch.pathways) + len(dna_branch.pathways):,} "
+            f"branch-specific pathways · GenePT: {html.escape(embedding_file)}"
+            f"{warning_text}</div>"
+        )
+        progress(1.0, desc="Ready to train")
+        return workspace, summary, architecture, enrichments.head(100)
+    except Exception as exc:
+        return None, _status(str(exc), error=True), pd.DataFrame(), pd.DataFrame()
+
+
+def _training_plots(history: list[dict[str, float]], confusion, labels):
+    history_frame = pd.DataFrame(history)
+    loss_figure = go.Figure()
+    loss_figure.add_trace(
+        go.Scatter(
+            x=history_frame["epoch"],
+            y=history_frame["training_loss"],
+            mode="lines",
+            name="Training",
+            line={"color": "#167c5a", "width": 3},
+        )
+    )
+    loss_figure.add_trace(
+        go.Scatter(
+            x=history_frame["epoch"],
+            y=history_frame["validation_loss"],
+            mode="lines",
+            name="Validation",
+            line={"color": "#e5a63c", "width": 3},
+        )
+    )
+    loss_figure.update_layout(
+        title="Loss by epoch",
+        xaxis_title="Epoch",
+        yaxis_title="Cross-entropy loss",
+        template="plotly_white",
+        margin={"l": 30, "r": 15, "t": 50, "b": 35},
+        legend={"orientation": "h", "y": 1.12},
+    )
+    confusion_figure = px.imshow(
+        confusion,
+        x=labels,
+        y=labels,
+        text_auto=True,
+        color_continuous_scale=[[0, "#eef6f2"], [1, "#167c5a"]],
+        labels={"x": "Predicted", "y": "Observed", "color": "Samples"},
+        title="Validation confusion matrix",
+    )
+    confusion_figure.update_layout(
+        template="plotly_white", margin={"l": 30, "r": 15, "t": 50, "b": 35}
+    )
+    return loss_figure, confusion_figure
+
+
+def estimate_training_duration(
+    workspace: PreparedWorkspace | None,
+    epochs: int,
+    batch_size: int,
+    learning_rate: float,
+    weight_decay: float,
+    dropout: float,
+    projection_dim: int,
+    fusion_dim: int,
+    validation_fraction: float,
+    optimizer: str,
+    class_weighting: bool,
+    progress=None,
+) -> int:
+    """Estimate a conservative ZeroGPU reservation from the prepared graph.
+
+    ZeroGPU checks the declared duration against each visitor's remaining quota
+    before the call starts. Keep small jobs short for better queue priority and
+    cap a single free-tier training request at five minutes.
+    """
+
+    if workspace is None:
+        return 10
+    samples = max(int(len(workspace.labels)), 1)
+    biological_parameters = (
+        int(workspace.gene_branch.biological_mask.size)
+        + int(workspace.dna_branch.biological_mask.size)
+    )
+    sample_factor = max(samples / 875.0, 0.25)
+    graph_factor = max(math.sqrt(biological_parameters / 1_850_000.0), 0.3)
+    batch_factor = max((16.0 / max(int(batch_size), 1)) ** 0.35, 0.55)
+    seconds = 25 + int(epochs) * 0.8 * sample_factor * graph_factor * batch_factor
+    return int(min(300, max(30, math.ceil(seconds))))
+
+
+@spaces.GPU(duration=estimate_training_duration)
+def train_workspace(
+    workspace: PreparedWorkspace | None,
+    epochs: int,
+    batch_size: int,
+    learning_rate: float,
+    weight_decay: float,
+    dropout: float,
+    projection_dim: int,
+    fusion_dim: int,
+    validation_fraction: float,
+    optimizer: str,
+    class_weighting: bool,
+    progress=gr.Progress(track_tqdm=False),
+):
+    if workspace is None:
+        return (
+            None,
+            _status("Prepare data and priors in Phase 1 before training.", True),
+            None,
+            None,
+            pd.DataFrame(),
+            None,
+            pd.DataFrame(),
+        )
+    try:
+        parameters = Hyperparameters(
+            epochs=int(epochs),
+            batch_size=int(batch_size),
+            learning_rate=float(learning_rate),
+            weight_decay=float(weight_decay),
+            dropout=float(dropout),
+            projection_dim=int(projection_dim),
+            fusion_dim=int(fusion_dim),
+            validation_fraction=float(validation_fraction),
+            optimizer=optimizer,
+            class_weighting=bool(class_weighting),
+        )
+
+        def report(fraction: float, description: str) -> None:
+            progress(fraction, desc=description)
+
+        result = train(workspace, parameters, progress=report)
+        bundle = result.bundle
+        artifact = save_bundle(bundle)
+        metrics = bundle.metrics
+        metrics_html = f"""
+        <div class="status-box"><strong>Training complete.</strong>
+        Best validation checkpoint restored; the downloadable artifact includes
+        architecture, preprocessing, weights, and metrics.</div>
+        <div class="metric-grid" style="margin-top:10px">
+          <div class="metric"><span>Macro F1</span><b>{metrics['f1_macro']:.3f}</b></div>
+          <div class="metric"><span>Accuracy</span><b>{metrics['accuracy']:.3f}</b></div>
+          <div class="metric"><span>Macro precision</span><b>{metrics['precision_macro']:.3f}</b></div>
+          <div class="metric"><span>Macro recall</span><b>{metrics['recall_macro']:.3f}</b></div>
+        </div>
+        """
+        loss_plot, confusion_plot = _training_plots(
+            bundle.history, result.confusion, bundle.label_names
+        )
+        importance = pathway_importance(bundle).head(100)
+        return (
+            bundle,
+            metrics_html,
+            loss_plot,
+            confusion_plot,
+            result.validation_predictions,
+            artifact,
+            importance,
+        )
+    except Exception as exc:
+        return (
+            None,
+            _status(str(exc), error=True),
+            None,
+            None,
+            pd.DataFrame(),
+            None,
+            pd.DataFrame(),
+        )
+
+
+def run_prediction(
+    bundle: ModelBundle | None,
+    uploaded_artifact: str | None,
+    gene_file: str | None,
+    dna_file: str | None,
+):
+    try:
+        active_bundle = (
+            load_bundle(uploaded_artifact) if uploaded_artifact else bundle
+        )
+        if active_bundle is None:
+            raise ValueError(
+                "Train a model in Phase 2 or upload a BioLM-NET model artifact."
+            )
+        gene_frame = _read_required_upload(
+            gene_file, "a prediction gene-expression CSV"
+        )
+        dna_frame = _read_required_upload(
+            dna_file, "a prediction DNA-methylation CSV"
+        )
+        output = predict(gene_frame, dna_frame, active_bundle)
+        destination = (
+            Path(tempfile.mkdtemp(prefix="biolmnet-prediction-"))
+            / "biolm-net-predictions.csv"
+        )
+        output.to_csv(destination, index=False)
+        counts = (
+            output["predicted_class"]
+            .value_counts()
+            .rename_axis("class")
+            .reset_index(name="samples")
+        )
+        figure = px.bar(
+            counts,
+            x="class",
+            y="samples",
+            color="class",
+            color_discrete_sequence=[
+                "#167c5a",
+                "#e5a63c",
+                "#497f93",
+                "#8d6fa8",
+                "#be6f55",
+            ],
+            title="Predicted class distribution",
+        )
+        figure.update_layout(
+            showlegend=False,
+            template="plotly_white",
+            margin={"l": 30, "r": 15, "t": 50, "b": 35},
+        )
+        status = _status(
+            f"Predicted {len(output):,} samples. Mean confidence: "
+            f"{output['confidence'].mean():.3f}."
+        )
+        return active_bundle, status, output, figure, str(destination)
+    except Exception as exc:
+        return bundle, _status(str(exc), True), pd.DataFrame(), None, None
+
+
+THEME = gr.themes.Base(
+    primary_hue="emerald",
+    neutral_hue="slate",
+)
+
+
+with gr.Blocks(title="BioLM-NET Workbench") as demo:
+    workspace_state = gr.State(None)
+    model_state = gr.State(None)
+
+    gr.HTML(
+        """
+        <section class="biolm-hero">
+          <div class="eyebrow">Interpretable multi-omics modeling</div>
+          <h1>BioLM-NET<br>Workbench</h1>
+          <p>Build a biologically masked network from paired gene expression and
+          DNA methylation, train it with GenePT-guided pathway attention, then
+          carry the exact preprocessing and architecture into prediction.</p>
+          <div class="hero-meta">
+            <span class="hero-chip">PDI · DoRothEA</span>
+            <span class="hero-chip">PPI · STRING</span>
+            <span class="hero-chip">Pathways · KEGG</span>
+            <span class="hero-chip">Context · GenePT</span>
+            <span class="hero-chip">Compute · ZeroGPU on demand</span>
+          </div>
+        </section>
+        """
+    )
+
+    with gr.Tabs():
+        with gr.Tab("1 · Data & priors", id="data"):
+            gr.HTML(
+                """
+                <p class="phase-intro"><strong>Assemble the model graph.</strong>
+                Select an upstream example, point to a GitHub folder that follows
+                the BioLM-NET file convention, or upload paired omics files.</p>
+                """
+            )
+            with gr.Row():
+                with gr.Column(scale=7, elem_classes=["phase-card"]):
+                    source_mode = gr.Radio(
+                        [
+                            "BioLM-NET examples",
+                            "GitHub folder",
+                            "Upload files",
+                        ],
+                        value="BioLM-NET examples",
+                        label="Dataset source",
+                    )
+                    with gr.Column(visible=True) as example_group:
+                        example_dataset = gr.Dropdown(
+                            ["BRCA", "COAD", "GBM", "scTrioseq2"],
+                            value="BRCA",
+                            label="Repository dataset",
+                        )
+                    with gr.Column(visible=False) as github_group:
+                        github_folder = gr.Textbox(
+                            label="GitHub dataset folder",
+                            placeholder=(
+                                "https://github.com/owner/repo/tree/main/Dataset/BRCA"
+                            ),
+                            info=(
+                                "The folder must contain the five standard "
+                                "BioLM-NET CSV filenames."
+                            ),
+                        )
+                    with gr.Column(visible=False) as upload_group:
+                        with gr.Row():
+                            uploaded_gene = gr.File(
+                                label="Gene expression",
+                                file_types=[".csv"],
+                                type="filepath",
+                            )
+                            uploaded_dna = gr.File(
+                                label="DNA methylation",
+                                file_types=[".csv"],
+                                type="filepath",
+                            )
+                            uploaded_labels = gr.File(
+                                label="Labels",
+                                file_types=[".csv"],
+                                type="filepath",
+                            )
+                        with gr.Row():
+                            uploaded_gene_pathway = gr.File(
+                                label="Gene → pathway mapping",
+                                file_types=[".csv"],
+                                type="filepath",
+                            )
+                            uploaded_dna_pathway = gr.File(
+                                label="DNA → pathway mapping",
+                                file_types=[".csv"],
+                                type="filepath",
+                            )
+                    pathway_files_are_significant = gr.Checkbox(
+                        value=True,
+                        label="Pathway files already contain significant pathways",
+                        info=(
+                            "Turn off for a full SYMBOL/PathwayID annotation "
+                            "catalog; enrichment will use BH-adjusted p < 0.05."
+                        ),
+                    )
+                with gr.Column(scale=5, elem_classes=["phase-card"]):
+                    embedding_option = gr.Dropdown(
+                        list(GENEPT_OPTIONS),
+                        value=list(GENEPT_OPTIONS)[0],
+                        label="GenePT context",
+                    )
+                    with gr.Accordion("Interaction priors", open=False):
+                        gr.Markdown(
+                            "By default, the app retrieves `PDI.csv` and `PPI.csv` "
+                            "from `bozdaglab/BioLM-NET`. Upload both only to override."
+                        )
+                        uploaded_pdi = gr.File(
+                            label="Custom PDI.csv",
+                            file_types=[".csv"],
+                            type="filepath",
+                        )
+                        uploaded_ppi = gr.File(
+                            label="Custom PPI.csv",
+                            file_types=[".csv"],
+                            type="filepath",
+                        )
+                    prepare_button = gr.Button(
+                        "Build biological architecture",
+                        variant="primary",
+                        size="lg",
+                    )
+
+            preparation_status = gr.HTML(
+                _status("Choose a source, then build the biological architecture.")
+            )
+            with gr.Row():
+                architecture_table = gr.Dataframe(
+                    label="Sparse architecture audit",
+                    interactive=False,
+                    wrap=True,
+                )
+                enrichment_table = gr.Dataframe(
+                    label="Retained enriched pathways (first 100)",
+                    interactive=False,
+                    wrap=True,
+                )
+            gr.HTML(
+                """
+                <p class="footnote">Expected orientation: samples in rows and
+                HGNC gene symbols in columns. PDI requires <code>TF, Target</code>;
+                PPI requires <code>protein1, protein2, combined_score</code>;
+                pathways require <code>SYMBOL, PathwayID</code>. PPI is filtered
+                to score &gt; 0.7 and the retained top decile, following the paper.</p>
+                """
+            )
+
+        with gr.Tab("2 · Train", id="train"):
+            gr.HTML(
+                """
+                <p class="phase-intro"><strong>Fit and evaluate.</strong> The
+                split is stratified; scaling is fit on training samples only;
+                the best validation checkpoint is exported as a safe,
+                self-contained model artifact. A shared GPU is requested only
+                while this training callback is running.</p>
+                """
+            )
+            with gr.Row():
+                with gr.Column(scale=4, elem_classes=["phase-card"]):
+                    epochs = gr.Slider(5, 200, value=50, step=5, label="Epochs")
+                    batch_size = gr.Dropdown(
+                        [8, 16, 32, 64, 128], value=16, label="Batch size"
+                    )
+                    learning_rate = gr.Number(
+                        value=0.001, label="Learning rate", minimum=0.000001
+                    )
+                    weight_decay = gr.Number(
+                        value=0.01, label="L2 weight decay", minimum=0
+                    )
+                    dropout = gr.Slider(
+                        0, 0.8, value=0.3, step=0.05, label="Dropout"
+                    )
+                with gr.Column(scale=4, elem_classes=["phase-card"]):
+                    projection_dim = gr.Dropdown(
+                        [16, 32, 64, 128], value=64, label="Branch projection"
+                    )
+                    fusion_dim = gr.Dropdown(
+                        [8, 12, 16, 32], value=12, label="Fusion layer"
+                    )
+                    validation_fraction = gr.Slider(
+                        0.1,
+                        0.4,
+                        value=0.2,
+                        step=0.05,
+                        label="Validation fraction",
+                    )
+                    optimizer = gr.Radio(
+                        ["Adam", "SGD"], value="Adam", label="Optimizer"
+                    )
+                    class_weighting = gr.Checkbox(
+                        value=True,
+                        label="Balance classes in the loss",
+                        info="Uses N / (classes × samples in class), as in the paper.",
+                    )
+                with gr.Column(scale=4, elem_classes=["phase-card"]):
+                    gr.Markdown(
+                        """
+                        **Paper-faithful defaults**
+
+                        - First layer: trainable `W ⊙ M`
+                        - PDI weights: binary
+                        - PPI weights: normalized STRING score
+                        - Pathway attention: GenePT query attention
+                        - Fusion: dual branch → dense → softmax
+                        """
+                    )
+                    train_button = gr.Button(
+                        "Train BioLM-NET on ZeroGPU", variant="primary", size="lg"
+                    )
+                    model_download = gr.File(
+                        label="Trained model artifact", interactive=False
+                    )
+                    gr.Markdown(
+                        "ZeroGPU reserves 30–300 seconds according to dataset "
+                        "size and epochs. Visitors use their own daily quota."
+                    )
+            training_status = gr.HTML(
+                _status("Phase 2 unlocks after the architecture is prepared.")
+            )
+            with gr.Row():
+                loss_plot = gr.Plot(label="Training history")
+                confusion_plot = gr.Plot(label="Confusion matrix")
+            with gr.Row():
+                validation_table = gr.Dataframe(
+                    label="Validation predictions",
+                    interactive=False,
+                    wrap=True,
+                )
+                importance_table = gr.Dataframe(
+                    label="Pathway attention audit (first 100)",
+                    interactive=False,
+                    wrap=True,
+                )
+
+        with gr.Tab("3 · Predict", id="predict"):
+            gr.HTML(
+                """
+                <p class="phase-intro"><strong>Apply a trained model.</strong>
+                Continue with the model from this session or upload a previous
+                artifact. Feature names are validated and reordered exactly as
+                they were during training.</p>
+                """
+            )
+            with gr.Row():
+                with gr.Column(scale=4, elem_classes=["phase-card"]):
+                    prediction_artifact = gr.File(
+                        label="Optional trained model artifact",
+                        file_types=[".zip"],
+                        type="filepath",
+                    )
+                    prediction_gene = gr.File(
+                        label="Prediction gene expression",
+                        file_types=[".csv"],
+                        type="filepath",
+                    )
+                    prediction_dna = gr.File(
+                        label="Prediction DNA methylation",
+                        file_types=[".csv"],
+                        type="filepath",
+                    )
+                    predict_button = gr.Button(
+                        "Make predictions", variant="primary", size="lg"
+                    )
+                    prediction_download = gr.File(
+                        label="Prediction CSV", interactive=False
+                    )
+                with gr.Column(scale=8):
+                    prediction_status = gr.HTML(
+                        _status("Use the current trained model or upload an artifact.")
+                    )
+                    prediction_plot = gr.Plot(label="Class distribution")
+            prediction_table = gr.Dataframe(
+                label="Predictions and class probabilities",
+                interactive=False,
+                wrap=True,
+            )
+
+    gr.Markdown(
+        """
+        <p class="footnote">Research use only. This interface reproduces the
+        architecture described by Rifat et al. and uses the upstream BioLM-NET
+        repository and GenePT embeddings. Validate cohorts, preprocessing, and
+        model performance before drawing biological or clinical conclusions.</p>
+        """
+    )
+
+    source_mode.change(
+        _source_visibility,
+        inputs=[source_mode],
+        outputs=[example_group, github_group, upload_group],
+    )
+    prepare_button.click(
+        prepare_workspace,
+        inputs=[
+            source_mode,
+            example_dataset,
+            github_folder,
+            uploaded_gene,
+            uploaded_dna,
+            uploaded_labels,
+            uploaded_gene_pathway,
+            uploaded_dna_pathway,
+            uploaded_pdi,
+            uploaded_ppi,
+            pathway_files_are_significant,
+            embedding_option,
+        ],
+        outputs=[
+            workspace_state,
+            preparation_status,
+            architecture_table,
+            enrichment_table,
+        ],
+    )
+    train_button.click(
+        train_workspace,
+        inputs=[
+            workspace_state,
+            epochs,
+            batch_size,
+            learning_rate,
+            weight_decay,
+            dropout,
+            projection_dim,
+            fusion_dim,
+            validation_fraction,
+            optimizer,
+            class_weighting,
+        ],
+        outputs=[
+            model_state,
+            training_status,
+            loss_plot,
+            confusion_plot,
+            validation_table,
+            model_download,
+            importance_table,
+        ],
+    )
+    predict_button.click(
+        run_prediction,
+        inputs=[
+            model_state,
+            prediction_artifact,
+            prediction_gene,
+            prediction_dna,
+        ],
+        outputs=[
+            model_state,
+            prediction_status,
+            prediction_table,
+            prediction_plot,
+            prediction_download,
+        ],
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch(theme=THEME, css=CSS)

biolmnet/__init__.py

@@ -0,0 +1,5 @@
+"""BioLM-NET model and data utilities for the Hugging Face Space."""
+
+from .model import BioLMNet
+
+__all__ = ["BioLMNet"]

biolmnet/artifacts.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import json
+import shutil
+import tempfile
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors.torch import load_file, save_file
+
+from .model import BioLMNet
+from .training import ModelBundle
+
+
+EXPECTED_FILES = {"config.json", "arrays.npz", "model.safetensors", "metrics.json"}
+
+
+def _branch_arrays(bundle: ModelBundle) -> dict[str, np.ndarray]:
+    model = bundle.model
+    return {
+        "gene_biological_mask": model.gene_branch.biological.mask.T.cpu().numpy(),
+        "dna_biological_mask": model.dna_branch.biological.mask.T.cpu().numpy(),
+        "gene_embeddings": (
+            model.gene_branch.pathway_attention.gene_embeddings.cpu().numpy()
+        ),
+        "dna_embeddings": (
+            model.dna_branch.pathway_attention.gene_embeddings.cpu().numpy()
+        ),
+        "gene_pathway_mask": (
+            model.gene_branch.pathway_attention.pathway_mask.cpu().numpy()
+        ),
+        "dna_pathway_mask": (
+            model.dna_branch.pathway_attention.pathway_mask.cpu().numpy()
+        ),
+        "gene_mean": bundle.gene_mean,
+        "gene_scale": bundle.gene_scale,
+        "dna_mean": bundle.dna_mean,
+        "dna_scale": bundle.dna_scale,
+    }
+
+
+def save_bundle(bundle: ModelBundle, destination: str | Path | None = None) -> str:
+    if destination is None:
+        destination = Path(tempfile.mkdtemp(prefix="biolmnet-export-")) / (
+            "biolm-net-trained-model.zip"
+        )
+    destination = Path(destination)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(prefix="biolmnet-pack-") as directory:
+        root = Path(directory)
+        (root / "config.json").write_text(
+            json.dumps(bundle.config, indent=2), encoding="utf-8"
+        )
+        (root / "metrics.json").write_text(
+            json.dumps(
+                {"metrics": bundle.metrics, "history": bundle.history}, indent=2
+            ),
+            encoding="utf-8",
+        )
+        np.savez_compressed(root / "arrays.npz", **_branch_arrays(bundle))
+        state = {
+            key: value.detach().cpu().contiguous()
+            for key, value in bundle.model.state_dict().items()
+        }
+        save_file(state, root / "model.safetensors")
+        with zipfile.ZipFile(
+            destination, "w", compression=zipfile.ZIP_DEFLATED
+        ) as archive:
+            for filename in sorted(EXPECTED_FILES):
+                archive.write(root / filename, arcname=filename)
+    return str(destination)
+
+
+def _safe_extract(archive_path: Path, directory: Path) -> None:
+    with zipfile.ZipFile(archive_path) as archive:
+        names = set(archive.namelist())
+        missing = EXPECTED_FILES - names
+        if missing:
+            raise ValueError(
+                "Model artifact is incomplete; missing " + ", ".join(sorted(missing))
+            )
+        for filename in EXPECTED_FILES:
+            info = archive.getinfo(filename)
+            if info.file_size > 1_000_000_000:
+                raise ValueError(f"Artifact member {filename} is unexpectedly large.")
+            with archive.open(info) as source, (directory / filename).open(
+                "wb"
+            ) as target:
+                shutil.copyfileobj(source, target)
+
+
+def load_bundle(archive_path: str | Path) -> ModelBundle:
+    archive_path = Path(archive_path)
+    if archive_path.suffix.lower() != ".zip":
+        raise ValueError("Upload the .zip artifact produced by the training phase.")
+    with tempfile.TemporaryDirectory(prefix="biolmnet-load-") as directory_name:
+        directory = Path(directory_name)
+        _safe_extract(archive_path, directory)
+        config = json.loads((directory / "config.json").read_text("utf-8"))
+        if config.get("format_version") != 1:
+            raise ValueError("Unsupported BioLM-NET artifact version.")
+        metrics_payload = json.loads(
+            (directory / "metrics.json").read_text("utf-8")
+        )
+        with np.load(directory / "arrays.npz", allow_pickle=False) as data:
+            arrays = {key: data[key].copy() for key in data.files}
+
+        architecture = config["architecture"]
+        model = BioLMNet(
+            gene_biological_mask=torch.from_numpy(
+                arrays["gene_biological_mask"]
+            ),
+            dna_biological_mask=torch.from_numpy(
+                arrays["dna_biological_mask"]
+            ),
+            gene_embeddings=torch.from_numpy(arrays["gene_embeddings"]),
+            dna_embeddings=torch.from_numpy(arrays["dna_embeddings"]),
+            gene_pathway_mask=torch.from_numpy(arrays["gene_pathway_mask"]),
+            dna_pathway_mask=torch.from_numpy(arrays["dna_pathway_mask"]),
+            n_classes=len(config["label_names"]),
+            projection_dim=int(architecture["projection_dim"]),
+            fusion_dim=int(architecture["fusion_dim"]),
+            dropout=float(architecture["dropout"]),
+            biological_activation=architecture["biological_activation"],
+            projection_activation=architecture["projection_activation"],
+            fusion_activation=architecture["fusion_activation"],
+        )
+        state = load_file(directory / "model.safetensors", device="cpu")
+        model.load_state_dict(state, strict=True)
+    model.eval()
+    return ModelBundle(
+        model=model,
+        gene_features=list(config["gene_features"]),
+        dna_features=list(config["dna_features"]),
+        label_names=list(config["label_names"]),
+        gene_mean=arrays["gene_mean"],
+        gene_scale=arrays["gene_scale"],
+        dna_mean=arrays["dna_mean"],
+        dna_scale=arrays["dna_scale"],
+        config=config,
+        metrics=metrics_payload.get("metrics", {}),
+        history=metrics_payload.get("history", []),
+    )

biolmnet/data.py

@@ -0,0 +1,470 @@
+from __future__ import annotations
+
+import hashlib
+import io
+import re
+import urllib.parse
+import urllib.request
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import BinaryIO
+
+import numpy as np
+import pandas as pd
+from scipy.stats import hypergeom
+
+
+UPSTREAM_REPOSITORY = "https://github.com/bozdaglab/BioLM-NET"
+UPSTREAM_RAW = "https://raw.githubusercontent.com/bozdaglab/BioLM-NET/main"
+GENEPT_REPOSITORY = "honicky/genept-composable-embeddings"
+DATASET_FILES = {
+    "gene": "Gene_Expression.csv",
+    "dna": "DNA_Methylation.csv",
+    "labels": "label.csv",
+    "gene_pathways": "ge_target_to_KEGG_significant.csv",
+    "dna_pathways": "dna_target_to_KEGG_significant.csv",
+}
+
+
+@dataclass
+class BranchPriors:
+    input_genes: list[str]
+    hidden_genes: list[str]
+    biological_mask: np.ndarray
+    pathways: list[str] = field(default_factory=list)
+    pathway_mask: np.ndarray | None = None
+    embeddings: np.ndarray | None = None
+    pdi_edges: int = 0
+    ppi_edges: int = 0
+    missing_embedding_genes: int = 0
+
+
+@dataclass
+class PreparedWorkspace:
+    gene_expression: np.ndarray
+    dna_methylation: np.ndarray
+    labels: np.ndarray
+    label_names: list[str]
+    gene_branch: BranchPriors
+    dna_branch: BranchPriors
+    source_name: str
+    warnings: list[str] = field(default_factory=list)
+
+
+def _normalise_columns(frame: pd.DataFrame) -> pd.DataFrame:
+    result = frame.copy()
+    result.columns = [str(value).strip() for value in result.columns]
+    return result
+
+
+def read_csv(source: str | Path | BinaryIO) -> pd.DataFrame:
+    if hasattr(source, "read"):
+        return _normalise_columns(pd.read_csv(source))
+    text = str(source)
+    if text.startswith(("http://", "https://")):
+        request = urllib.request.Request(
+            text, headers={"User-Agent": "BioLM-NET-HuggingFace-Space/1.0"}
+        )
+        with urllib.request.urlopen(request, timeout=60) as response:
+            payload = response.read()
+        return _normalise_columns(pd.read_csv(io.BytesIO(payload)))
+    return _normalise_columns(pd.read_csv(text))
+
+
+def github_folder_to_raw_base(folder_url: str) -> str:
+    value = folder_url.strip().rstrip("/")
+    if value.startswith("https://raw.githubusercontent.com/"):
+        return value
+    match = re.match(
+        r"https://github\.com/([^/]+)/([^/]+)/(?:tree|blob)/([^/]+)(?:/(.*))?$",
+        value,
+    )
+    if match:
+        owner, repository, branch, folder = match.groups()
+        suffix = f"/{folder}" if folder else ""
+        return (
+            f"https://raw.githubusercontent.com/{owner}/{repository}/"
+            f"{branch}{suffix}"
+        )
+    match = re.match(r"https://github\.com/([^/]+)/([^/]+)$", value)
+    if match:
+        owner, repository = match.groups()
+        return f"https://raw.githubusercontent.com/{owner}/{repository}/main"
+    raise ValueError(
+        "Use a GitHub repository/folder URL such as "
+        "https://github.com/owner/repo/tree/main/Dataset/BRCA."
+    )
+
+
+def upstream_example_sources(dataset: str) -> dict[str, str]:
+    if dataset not in {"BRCA", "COAD", "GBM", "scTrioseq2"}:
+        raise ValueError(f"Unknown BioLM-NET example dataset: {dataset}")
+    base = f"{UPSTREAM_RAW}/Dataset/{dataset}"
+    return {key: f"{base}/{filename}" for key, filename in DATASET_FILES.items()}
+
+
+def github_dataset_sources(folder_url: str) -> dict[str, str]:
+    base = github_folder_to_raw_base(folder_url)
+    return {key: f"{base}/{filename}" for key, filename in DATASET_FILES.items()}
+
+
+def upstream_interaction_sources() -> tuple[str, str]:
+    return (
+        f"{UPSTREAM_RAW}/Dataset/PDI/PDI.csv",
+        f"{UPSTREAM_RAW}/Dataset/PPI/PPI.csv",
+    )
+
+
+def validate_and_align_omics(
+    gene_frame: pd.DataFrame,
+    dna_frame: pd.DataFrame,
+    labels_frame: pd.DataFrame,
+    *,
+    allow_preset_trim: bool = False,
+) -> tuple[pd.DataFrame, pd.DataFrame, np.ndarray, list[str], list[str]]:
+    warnings: list[str] = []
+    if gene_frame.columns.duplicated().any() or dna_frame.columns.duplicated().any():
+        raise ValueError("Omics files must have unique gene-name columns.")
+    if labels_frame.shape[1] != 1:
+        raise ValueError("The label file must contain exactly one column.")
+    if gene_frame.empty or dna_frame.empty or labels_frame.empty:
+        raise ValueError("Gene expression, DNA methylation, and labels cannot be empty.")
+
+    counts = [len(gene_frame), len(dna_frame), len(labels_frame)]
+    if len(set(counts)) != 1:
+        if not allow_preset_trim:
+            raise ValueError(
+                "The two omics files and label file must contain the same number "
+                f"of rows; received {counts}."
+            )
+        common = min(counts)
+        warnings.append(
+            f"The upstream example has row counts {counts}; all inputs were "
+            f"aligned to the first {common} rows, matching repository order."
+        )
+        gene_frame = gene_frame.iloc[:common].reset_index(drop=True)
+        dna_frame = dna_frame.iloc[:common].reset_index(drop=True)
+        labels_frame = labels_frame.iloc[:common].reset_index(drop=True)
+
+    for name, frame in (
+        ("gene expression", gene_frame),
+        ("DNA methylation", dna_frame),
+    ):
+        converted = frame.apply(pd.to_numeric, errors="coerce")
+        invalid = int(converted.isna().sum().sum())
+        if invalid:
+            raise ValueError(
+                f"{name.title()} contains {invalid:,} missing or non-numeric values."
+            )
+        if not np.isfinite(converted.to_numpy(dtype=np.float64)).all():
+            raise ValueError(f"{name.title()} contains infinite values.")
+        if name == "gene expression":
+            gene_frame = converted
+        else:
+            dna_frame = converted
+
+    raw_labels = labels_frame.iloc[:, 0]
+    if raw_labels.isna().any():
+        raise ValueError("Labels cannot be empty.")
+    labels = raw_labels.astype(str).str.strip()
+    if labels.eq("").any():
+        raise ValueError("Labels cannot be empty.")
+    unique_labels = sorted(labels.unique().tolist())
+    if len(unique_labels) < 2:
+        raise ValueError("Training requires at least two label classes.")
+    label_to_index = {label: index for index, label in enumerate(unique_labels)}
+    encoded = labels.map(label_to_index).to_numpy(dtype=np.int64)
+    return gene_frame, dna_frame, encoded, unique_labels, warnings
+
+
+def _clean_interactions(
+    pdi_frame: pd.DataFrame, ppi_frame: pd.DataFrame
+) -> tuple[pd.DataFrame, pd.DataFrame]:
+    required_pdi = {"TF", "Target"}
+    required_ppi = {"protein1", "protein2", "combined_score"}
+    if not required_pdi.issubset(pdi_frame.columns):
+        raise ValueError("PDI.csv must contain TF and Target columns.")
+    if not required_ppi.issubset(ppi_frame.columns):
+        raise ValueError(
+            "PPI.csv must contain protein1, protein2, and combined_score columns."
+        )
+
+    pdi = pdi_frame.loc[:, ["TF", "Target"]].dropna().copy()
+    pdi["TF"] = pdi["TF"].astype(str).str.strip()
+    pdi["Target"] = pdi["Target"].astype(str).str.strip()
+    pdi = pdi[(pdi["TF"] != "") & (pdi["Target"] != "")].drop_duplicates()
+
+    ppi = ppi_frame.loc[
+        :, ["protein1", "protein2", "combined_score"]
+    ].dropna().copy()
+    ppi["protein1"] = ppi["protein1"].astype(str).str.strip()
+    ppi["protein2"] = ppi["protein2"].astype(str).str.strip()
+    ppi["combined_score"] = pd.to_numeric(
+        ppi["combined_score"], errors="coerce"
+    )
+    ppi = ppi.dropna()
+    if ppi["combined_score"].max() > 1:
+        ppi["combined_score"] = ppi["combined_score"] / 1000.0
+    ppi = ppi[ppi["combined_score"] > 0.7]
+    if ppi.empty:
+        raise ValueError("No PPI interactions remain above score 0.7.")
+    threshold = float(ppi["combined_score"].quantile(0.9))
+    ppi = ppi[ppi["combined_score"] >= threshold].drop_duplicates(
+        ["protein1", "protein2"]
+    )
+    return pdi, ppi
+
+
+def build_biological_mask(
+    input_genes: list[str],
+    pdi_frame: pd.DataFrame,
+    ppi_frame: pd.DataFrame,
+) -> BranchPriors:
+    pdi, ppi = _clean_interactions(pdi_frame, ppi_frame)
+    input_genes = [str(gene).strip() for gene in input_genes]
+    input_index = {gene: index for index, gene in enumerate(input_genes)}
+    input_set = set(input_genes)
+
+    # The paper retains PDI targets that are DE/HVG; both TF and target must
+    # therefore be represented in the input branch.
+    pdi_selected = pdi[
+        pdi["TF"].isin(input_set) & pdi["Target"].isin(input_set)
+    ].copy()
+
+    # STRING PPI is undirected. Add the partner of every input protein,
+    # regardless of which endpoint it occupies in the source file.
+    forward = ppi[ppi["protein1"].isin(input_set)].rename(
+        columns={"protein1": "source", "protein2": "target"}
+    )
+    reverse = ppi[ppi["protein2"].isin(input_set)].rename(
+        columns={"protein2": "source", "protein1": "target"}
+    )
+    ppi_selected = pd.concat(
+        [
+            forward[["source", "target", "combined_score"]],
+            reverse[["source", "target", "combined_score"]],
+        ],
+        ignore_index=True,
+    ).drop_duplicates(["source", "target"])
+
+    hidden_genes = sorted(
+        set(pdi_selected["Target"].tolist())
+        | set(ppi_selected["target"].tolist())
+    )
+    if not hidden_genes:
+        raise ValueError(
+            "None of the input genes have retained PDI/PPI connections. "
+            "Check that columns use HGNC gene symbols."
+        )
+    hidden_index = {gene: index for index, gene in enumerate(hidden_genes)}
+    mask = np.zeros((len(input_genes), len(hidden_genes)), dtype=np.float32)
+
+    for row in pdi_selected.itertuples(index=False):
+        mask[input_index[row.TF], hidden_index[row.Target]] = 1.0
+    for row in ppi_selected.itertuples(index=False):
+        source_index = input_index[row.source]
+        target_index = hidden_index[row.target]
+        mask[source_index, target_index] = max(
+            mask[source_index, target_index], float(row.combined_score)
+        )
+
+    active = mask.sum(axis=0) > 0
+    return BranchPriors(
+        input_genes=input_genes,
+        hidden_genes=[
+            gene for gene, keep in zip(hidden_genes, active, strict=True) if keep
+        ],
+        biological_mask=mask[:, active],
+        pdi_edges=int(len(pdi_selected)),
+        ppi_edges=int(len(ppi_selected)),
+    )
+
+
+def _benjamini_hochberg(p_values: np.ndarray) -> np.ndarray:
+    count = len(p_values)
+    order = np.argsort(p_values)
+    ranked = p_values[order]
+    adjusted = ranked * count / np.arange(1, count + 1)
+    adjusted = np.minimum.accumulate(adjusted[::-1])[::-1]
+    output = np.empty_like(adjusted)
+    output[order] = np.clip(adjusted, 0.0, 1.0)
+    return output
+
+
+def build_pathway_mask(
+    hidden_genes: list[str],
+    pathway_frame: pd.DataFrame,
+    *,
+    precomputed_significant: bool,
+    adjusted_p_threshold: float = 0.05,
+) -> tuple[list[str], np.ndarray, pd.DataFrame]:
+    required = {"SYMBOL", "PathwayID"}
+    if not required.issubset(pathway_frame.columns):
+        raise ValueError("Pathway data must contain SYMBOL and PathwayID columns.")
+    mapping = pathway_frame.loc[:, ["SYMBOL", "PathwayID"]].dropna().copy()
+    mapping["SYMBOL"] = mapping["SYMBOL"].astype(str).str.strip()
+    mapping["PathwayID"] = mapping["PathwayID"].astype(str).str.strip()
+    mapping = mapping[
+        (mapping["SYMBOL"] != "") & (mapping["PathwayID"] != "")
+    ].drop_duplicates()
+
+    hidden_set = set(hidden_genes)
+    overlap = mapping[mapping["SYMBOL"].isin(hidden_set)]
+    if overlap.empty:
+        raise ValueError(
+            "No PDI/PPI hidden genes overlap the supplied pathway annotations."
+        )
+
+    rows: list[dict[str, float | int | str]] = []
+    if precomputed_significant:
+        for pathway, group in overlap.groupby("PathwayID"):
+            rows.append(
+                {
+                    "PathwayID": pathway,
+                    "overlap_genes": int(group["SYMBOL"].nunique()),
+                    "adjusted_p_value": np.nan,
+                }
+            )
+    else:
+        universe = set(mapping["SYMBOL"])
+        selected = hidden_set & universe
+        population = len(universe)
+        draws = len(selected)
+        for pathway, group in mapping.groupby("PathwayID"):
+            members = set(group["SYMBOL"])
+            successes = len(members)
+            observed = len(selected & members)
+            if observed == 0:
+                continue
+            p_value = float(
+                hypergeom.sf(observed - 1, population, successes, draws)
+            )
+            rows.append(
+                {
+                    "PathwayID": pathway,
+                    "overlap_genes": observed,
+                    "p_value": p_value,
+                }
+            )
+        if rows:
+            p_values = np.array([float(row["p_value"]) for row in rows])
+            adjusted = _benjamini_hochberg(p_values)
+            for row, value in zip(rows, adjusted, strict=True):
+                row["adjusted_p_value"] = float(value)
+            rows = [
+                row
+                for row in rows
+                if float(row["adjusted_p_value"]) < adjusted_p_threshold
+            ]
+
+    enrichment = pd.DataFrame(rows)
+    if enrichment.empty:
+        raise ValueError(
+            "No significantly enriched pathways remain at BH-adjusted p < 0.05. "
+            "Upload a broader gene-to-pathway annotation catalog or revise the "
+            "input feature selection."
+        )
+    enrichment = enrichment.sort_values(
+        ["overlap_genes", "PathwayID"], ascending=[False, True]
+    ).reset_index(drop=True)
+    pathways = enrichment["PathwayID"].astype(str).tolist()
+    gene_index = {gene: index for index, gene in enumerate(hidden_genes)}
+    pathway_index = {
+        pathway: index for index, pathway in enumerate(pathways)
+    }
+    mask = np.zeros((len(hidden_genes), len(pathways)), dtype=bool)
+    kept_mapping = overlap[overlap["PathwayID"].isin(pathway_index)]
+    for row in kept_mapping.itertuples(index=False):
+        mask[gene_index[row.SYMBOL], pathway_index[row.PathwayID]] = True
+
+    active_pathways = mask.sum(axis=0) > 0
+    pathways = [
+        pathway
+        for pathway, keep in zip(pathways, active_pathways, strict=True)
+        if keep
+    ]
+    return pathways, mask[:, active_pathways], enrichment
+
+
+def deterministic_gene_embeddings(
+    genes: list[str], dimensions: int = 64
+) -> pd.DataFrame:
+    """Deterministic test/fallback embeddings, never silently used for GenePT."""
+
+    vectors = []
+    for gene in genes:
+        digest = hashlib.sha256(gene.encode("utf-8")).digest()
+        seed = int.from_bytes(digest[:8], "little")
+        generator = np.random.default_rng(seed)
+        vector = generator.normal(0, 1, dimensions).astype(np.float32)
+        vector /= max(float(np.linalg.norm(vector)), 1e-8)
+        vectors.append(vector)
+    return pd.DataFrame(vectors, index=genes)
+
+
+def load_genept_embeddings(filename: str) -> pd.DataFrame:
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as exc:
+        raise RuntimeError(
+            "huggingface_hub is required to retrieve GenePT embeddings."
+        ) from exc
+    path = hf_hub_download(
+        repo_id=GENEPT_REPOSITORY,
+        filename=filename,
+        repo_type="model",
+    )
+    frame = pd.read_parquet(path)
+    frame.index = frame.index.astype(str).str.strip()
+    return frame
+
+
+def attach_embeddings_and_pathways(
+    branch: BranchPriors,
+    embedding_frame: pd.DataFrame,
+    pathway_frame: pd.DataFrame,
+    *,
+    precomputed_significant: bool,
+) -> pd.DataFrame:
+    embedding_index = set(embedding_frame.index.astype(str))
+    keep = np.array(
+        [gene in embedding_index for gene in branch.hidden_genes], dtype=bool
+    )
+    branch.missing_embedding_genes = int((~keep).sum())
+    if not keep.any():
+        raise ValueError(
+            "No retained PDI/PPI genes have embeddings in the selected GenePT file."
+        )
+    branch.hidden_genes = [
+        gene
+        for gene, retained in zip(branch.hidden_genes, keep, strict=True)
+        if retained
+    ]
+    branch.biological_mask = branch.biological_mask[:, keep]
+    branch.embeddings = (
+        embedding_frame.loc[branch.hidden_genes].to_numpy(dtype=np.float32)
+    )
+    (
+        branch.pathways,
+        branch.pathway_mask,
+        enrichment,
+    ) = build_pathway_mask(
+        branch.hidden_genes,
+        pathway_frame,
+        precomputed_significant=precomputed_significant,
+    )
+
+    genes_with_pathways = branch.pathway_mask.sum(axis=1) > 0
+    if not genes_with_pathways.any():
+        raise ValueError("No embedded hidden genes belong to a retained pathway.")
+    branch.hidden_genes = [
+        gene
+        for gene, retained in zip(
+            branch.hidden_genes, genes_with_pathways, strict=True
+        )
+        if retained
+    ]
+    branch.biological_mask = branch.biological_mask[:, genes_with_pathways]
+    branch.embeddings = branch.embeddings[genes_with_pathways]
+    branch.pathway_mask = branch.pathway_mask[genes_with_pathways]
+    return enrichment

biolmnet/model.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def _activation(name: str) -> nn.Module:
+    choices: dict[str, nn.Module] = {
+        "relu": nn.ReLU(),
+        "gelu": nn.GELU(),
+        "sigmoid": nn.Sigmoid(),
+        "tanh": nn.Tanh(),
+    }
+    try:
+        return choices[name.lower()]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported activation: {name}") from exc
+
+
+class BioMaskedLinear(nn.Module):
+    """Trainable linear layer whose weights are constrained by a biological mask.
+
+    The paper defines W_masked = W ⊙ M. The public notebook approximates this
+    with a dense projection followed by a fixed matrix multiplication; this
+    layer implements the paper's equation directly.
+    """
+
+    def __init__(self, mask: Tensor, bias: bool = True) -> None:
+        super().__init__()
+        if mask.ndim != 2:
+            raise ValueError("Biological mask must be [input_genes, hidden_genes].")
+        if mask.shape[0] == 0 or mask.shape[1] == 0:
+            raise ValueError("Biological mask cannot be empty.")
+
+        input_features, output_features = mask.shape
+        self.input_features = int(input_features)
+        self.output_features = int(output_features)
+        self.weight = nn.Parameter(torch.empty(output_features, input_features))
+        self.bias = nn.Parameter(torch.empty(output_features)) if bias else None
+        self.register_buffer("mask", mask.T.to(dtype=torch.float32).contiguous())
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return F.linear(inputs, self.weight * self.mask, self.bias)
+
+
+class AttentionPathwayLayer(nn.Module):
+    """GenePT-guided attention from biological hidden genes to pathways."""
+
+    def __init__(self, gene_embeddings: Tensor, pathway_mask: Tensor) -> None:
+        super().__init__()
+        if gene_embeddings.ndim != 2:
+            raise ValueError("Gene embeddings must be [hidden_genes, embedding_dim].")
+        if pathway_mask.ndim != 2:
+            raise ValueError("Pathway mask must be [hidden_genes, pathways].")
+        if gene_embeddings.shape[0] != pathway_mask.shape[0]:
+            raise ValueError("Gene embeddings and pathway mask must share gene order.")
+        if torch.any(pathway_mask.sum(dim=0) == 0):
+            raise ValueError("Every pathway must contain at least one retained gene.")
+
+        embeddings = gene_embeddings.to(dtype=torch.float32)
+        membership = pathway_mask.to(dtype=torch.bool)
+        self.register_buffer("gene_embeddings", embeddings)
+        self.register_buffer("pathway_mask", membership)
+        self.query = nn.Parameter(
+            torch.empty(pathway_mask.shape[1], gene_embeddings.shape[1])
+        )
+        nn.init.xavier_uniform_(self.query)
+
+    def attention_weights(self) -> Tensor:
+        scale = math.sqrt(self.gene_embeddings.shape[1])
+        # [hidden genes, pathways]
+        scores = self.gene_embeddings @ self.query.T / scale
+        scores = scores.masked_fill(~self.pathway_mask, torch.finfo(scores.dtype).min)
+        return torch.softmax(scores, dim=0)
+
+    def forward(self, hidden_gene_signal: Tensor) -> Tensor:
+        return hidden_gene_signal @ self.attention_weights()
+
+
+class BioBranch(nn.Module):
+    def __init__(
+        self,
+        biological_mask: Tensor,
+        gene_embeddings: Tensor,
+        pathway_mask: Tensor,
+        projection_dim: int,
+        dropout: float,
+        biological_activation: str,
+        projection_activation: str,
+    ) -> None:
+        super().__init__()
+        self.biological = BioMaskedLinear(biological_mask)
+        self.biological_activation = _activation(biological_activation)
+        self.dropout = nn.Dropout(dropout)
+        self.pathway_attention = AttentionPathwayLayer(
+            gene_embeddings, pathway_mask
+        )
+        self.projection = nn.Linear(pathway_mask.shape[1], projection_dim)
+        self.projection_activation = _activation(projection_activation)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        hidden = self.dropout(self.biological_activation(self.biological(inputs)))
+        pathways = self.pathway_attention(hidden)
+        return self.projection_activation(self.projection(pathways))
+
+
+@dataclass(frozen=True)
+class ModelDimensions:
+    gene_inputs: int
+    dna_inputs: int
+    gene_hidden: int
+    dna_hidden: int
+    gene_pathways: int
+    dna_pathways: int
+    classes: int
+
+
+class BioLMNet(nn.Module):
+    """Dual-omics BioLM-NET classifier."""
+
+    def __init__(
+        self,
+        gene_biological_mask: Tensor,
+        dna_biological_mask: Tensor,
+        gene_embeddings: Tensor,
+        dna_embeddings: Tensor,
+        gene_pathway_mask: Tensor,
+        dna_pathway_mask: Tensor,
+        n_classes: int,
+        projection_dim: int = 64,
+        fusion_dim: int = 12,
+        dropout: float = 0.3,
+        biological_activation: str = "relu",
+        projection_activation: str = "sigmoid",
+        fusion_activation: str = "tanh",
+    ) -> None:
+        super().__init__()
+        if n_classes < 2:
+            raise ValueError("BioLM-NET requires at least two label classes.")
+
+        self.gene_branch = BioBranch(
+            gene_biological_mask,
+            gene_embeddings,
+            gene_pathway_mask,
+            projection_dim,
+            dropout,
+            biological_activation,
+            projection_activation,
+        )
+        self.dna_branch = BioBranch(
+            dna_biological_mask,
+            dna_embeddings,
+            dna_pathway_mask,
+            projection_dim,
+            dropout,
+            biological_activation,
+            projection_activation,
+        )
+        self.fusion = nn.Linear(projection_dim * 2, fusion_dim)
+        self.fusion_activation = _activation(fusion_activation)
+        self.fusion_dropout = nn.Dropout(dropout)
+        self.output = nn.Linear(fusion_dim, n_classes)
+        self.dimensions = ModelDimensions(
+            gene_inputs=gene_biological_mask.shape[0],
+            dna_inputs=dna_biological_mask.shape[0],
+            gene_hidden=gene_biological_mask.shape[1],
+            dna_hidden=dna_biological_mask.shape[1],
+            gene_pathways=gene_pathway_mask.shape[1],
+            dna_pathways=dna_pathway_mask.shape[1],
+            classes=n_classes,
+        )
+
+    def forward(self, gene_expression: Tensor, dna_methylation: Tensor) -> Tensor:
+        gene_projection = self.gene_branch(gene_expression)
+        dna_projection = self.dna_branch(dna_methylation)
+        fused = torch.cat([gene_projection, dna_projection], dim=1)
+        fused = self.fusion_dropout(
+            self.fusion_activation(self.fusion(fused))
+        )
+        return self.output(fused)
+
+    def pathway_attention(self) -> dict[str, Tensor]:
+        return {
+            "gene_expression": self.gene_branch.pathway_attention.attention_weights(),
+            "dna_methylation": self.dna_branch.pathway_attention.attention_weights(),
+        }

biolmnet/training.py

@@ -0,0 +1,482 @@
+from __future__ import annotations
+
+import copy
+import random
+from dataclasses import asdict, dataclass
+from typing import Callable
+
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.metrics import (
+    accuracy_score,
+    confusion_matrix,
+    f1_score,
+    precision_score,
+    recall_score,
+)
+from sklearn.model_selection import train_test_split
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from .data import PreparedWorkspace
+from .model import BioLMNet
+
+
+ProgressCallback = Callable[[float, str], None]
+
+
+@dataclass(frozen=True)
+class Hyperparameters:
+    epochs: int = 50
+    batch_size: int = 16
+    learning_rate: float = 0.001
+    weight_decay: float = 0.01
+    dropout: float = 0.3
+    projection_dim: int = 64
+    fusion_dim: int = 12
+    validation_fraction: float = 0.2
+    optimizer: str = "Adam"
+    class_weighting: bool = True
+    early_stopping_patience: int = 12
+    seed: int = 42
+
+    def validate(self) -> None:
+        if not 1 <= self.epochs <= 1000:
+            raise ValueError("Epochs must be between 1 and 1,000.")
+        if not 2 <= self.batch_size <= 1024:
+            raise ValueError("Batch size must be between 2 and 1,024.")
+        if not 0 < self.learning_rate <= 1:
+            raise ValueError("Learning rate must be in (0, 1].")
+        if not 0 <= self.dropout < 1:
+            raise ValueError("Dropout must be in [0, 1).")
+        if not 0.05 <= self.validation_fraction <= 0.5:
+            raise ValueError("Validation fraction must be between 0.05 and 0.5.")
+        if self.optimizer.lower() not in {"adam", "sgd"}:
+            raise ValueError("Optimizer must be Adam or SGD.")
+
+
+@dataclass
+class ModelBundle:
+    model: BioLMNet
+    gene_features: list[str]
+    dna_features: list[str]
+    label_names: list[str]
+    gene_mean: np.ndarray
+    gene_scale: np.ndarray
+    dna_mean: np.ndarray
+    dna_scale: np.ndarray
+    config: dict
+    metrics: dict
+    history: list[dict[str, float]]
+
+
+@dataclass
+class TrainingResult:
+    bundle: ModelBundle
+    validation_predictions: pd.DataFrame
+    confusion: np.ndarray
+
+
+def set_reproducible_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def _fit_scaler(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    mean = values.mean(axis=0, dtype=np.float64).astype(np.float32)
+    scale = values.std(axis=0, dtype=np.float64).astype(np.float32)
+    scale[scale < 1e-8] = 1.0
+    return mean, scale
+
+
+def _scale(values: np.ndarray, mean: np.ndarray, scale: np.ndarray) -> np.ndarray:
+    return ((values.astype(np.float32) - mean) / scale).astype(np.float32)
+
+
+def _make_model(
+    workspace: PreparedWorkspace, hyperparameters: Hyperparameters
+) -> BioLMNet:
+    gene = workspace.gene_branch
+    dna = workspace.dna_branch
+    if (
+        gene.embeddings is None
+        or dna.embeddings is None
+        or gene.pathway_mask is None
+        or dna.pathway_mask is None
+    ):
+        raise ValueError("Workspace priors are incomplete; run data preparation first.")
+
+    return BioLMNet(
+        gene_biological_mask=torch.from_numpy(gene.biological_mask),
+        dna_biological_mask=torch.from_numpy(dna.biological_mask),
+        gene_embeddings=torch.from_numpy(gene.embeddings),
+        dna_embeddings=torch.from_numpy(dna.embeddings),
+        gene_pathway_mask=torch.from_numpy(gene.pathway_mask),
+        dna_pathway_mask=torch.from_numpy(dna.pathway_mask),
+        n_classes=len(workspace.label_names),
+        projection_dim=hyperparameters.projection_dim,
+        fusion_dim=hyperparameters.fusion_dim,
+        dropout=hyperparameters.dropout,
+    )
+
+
+def _evaluate(
+    model: BioLMNet,
+    gene_values: np.ndarray,
+    dna_values: np.ndarray,
+    labels: np.ndarray,
+    device: torch.device,
+) -> tuple[float, np.ndarray, np.ndarray]:
+    model.eval()
+    with torch.no_grad():
+        logits = model(
+            torch.from_numpy(gene_values).to(device),
+            torch.from_numpy(dna_values).to(device),
+        )
+        loss = nn.functional.cross_entropy(
+            logits, torch.from_numpy(labels).to(device)
+        ).item()
+        probabilities = torch.softmax(logits, dim=1).cpu().numpy()
+    predictions = probabilities.argmax(axis=1)
+    return float(loss), probabilities, predictions
+
+
+def train(
+    workspace: PreparedWorkspace,
+    hyperparameters: Hyperparameters,
+    progress: ProgressCallback | None = None,
+) -> TrainingResult:
+    hyperparameters.validate()
+    set_reproducible_seed(hyperparameters.seed)
+    labels = workspace.labels
+    classes, class_counts = np.unique(labels, return_counts=True)
+    if class_counts.min() < 2:
+        raise ValueError(
+            "Each class needs at least two samples for a stratified train/validation split."
+        )
+
+    indices = np.arange(len(labels))
+    train_index, validation_index = train_test_split(
+        indices,
+        test_size=hyperparameters.validation_fraction,
+        random_state=hyperparameters.seed,
+        stratify=labels,
+    )
+    gene_mean, gene_scale = _fit_scaler(
+        workspace.gene_expression[train_index]
+    )
+    dna_mean, dna_scale = _fit_scaler(workspace.dna_methylation[train_index])
+    gene_train = _scale(
+        workspace.gene_expression[train_index], gene_mean, gene_scale
+    )
+    gene_validation = _scale(
+        workspace.gene_expression[validation_index], gene_mean, gene_scale
+    )
+    dna_train = _scale(
+        workspace.dna_methylation[train_index], dna_mean, dna_scale
+    )
+    dna_validation = _scale(
+        workspace.dna_methylation[validation_index], dna_mean, dna_scale
+    )
+    y_train = labels[train_index]
+    y_validation = labels[validation_index]
+
+    dataset = TensorDataset(
+        torch.from_numpy(gene_train),
+        torch.from_numpy(dna_train),
+        torch.from_numpy(y_train),
+    )
+    generator = torch.Generator().manual_seed(hyperparameters.seed)
+    loader = DataLoader(
+        dataset,
+        batch_size=min(hyperparameters.batch_size, len(dataset)),
+        shuffle=True,
+        generator=generator,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = _make_model(workspace, hyperparameters).to(device)
+    if hyperparameters.optimizer.lower() == "sgd":
+        optimizer = torch.optim.SGD(
+            model.parameters(),
+            lr=hyperparameters.learning_rate,
+            weight_decay=hyperparameters.weight_decay,
+        )
+    else:
+        optimizer = torch.optim.Adam(
+            model.parameters(),
+            lr=hyperparameters.learning_rate,
+            weight_decay=hyperparameters.weight_decay,
+        )
+
+    class_weights: torch.Tensor | None = None
+    if hyperparameters.class_weighting:
+        count_by_class = np.bincount(
+            y_train, minlength=len(workspace.label_names)
+        ).astype(np.float32)
+        weights = len(y_train) / (len(count_by_class) * count_by_class)
+        class_weights = torch.from_numpy(weights).to(device)
+    criterion = nn.CrossEntropyLoss(weight=class_weights)
+
+    history: list[dict[str, float]] = []
+    best_state: dict[str, torch.Tensor] | None = None
+    best_loss = float("inf")
+    patience = 0
+    for epoch in range(hyperparameters.epochs):
+        model.train()
+        running_loss = 0.0
+        seen = 0
+        for gene_batch, dna_batch, label_batch in loader:
+            gene_batch = gene_batch.to(device)
+            dna_batch = dna_batch.to(device)
+            label_batch = label_batch.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            logits = model(gene_batch, dna_batch)
+            loss = criterion(logits, label_batch)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item() * len(label_batch)
+            seen += len(label_batch)
+
+        validation_loss, _, validation_predictions = _evaluate(
+            model,
+            gene_validation,
+            dna_validation,
+            y_validation,
+            device,
+        )
+        epoch_row = {
+            "epoch": float(epoch + 1),
+            "training_loss": float(running_loss / max(seen, 1)),
+            "validation_loss": validation_loss,
+            "validation_accuracy": float(
+                accuracy_score(y_validation, validation_predictions)
+            ),
+            "validation_f1_macro": float(
+                f1_score(
+                    y_validation,
+                    validation_predictions,
+                    average="macro",
+                    zero_division=0,
+                )
+            ),
+        }
+        history.append(epoch_row)
+        if progress:
+            progress(
+                (epoch + 1) / hyperparameters.epochs,
+                (
+                    f"Epoch {epoch + 1}/{hyperparameters.epochs} · "
+                    f"validation F1 {epoch_row['validation_f1_macro']:.3f}"
+                ),
+            )
+
+        if validation_loss < best_loss - 1e-5:
+            best_loss = validation_loss
+            best_state = copy.deepcopy(model.state_dict())
+            patience = 0
+        else:
+            patience += 1
+            if patience >= hyperparameters.early_stopping_patience:
+                break
+
+    if best_state is not None:
+        model.load_state_dict(best_state)
+    validation_loss, probabilities, predictions = _evaluate(
+        model,
+        gene_validation,
+        dna_validation,
+        y_validation,
+        device,
+    )
+    confusion = confusion_matrix(
+        y_validation,
+        predictions,
+        labels=np.arange(len(workspace.label_names)),
+    )
+    metrics = {
+        "validation_loss": validation_loss,
+        "accuracy": float(accuracy_score(y_validation, predictions)),
+        "f1_macro": float(
+            f1_score(
+                y_validation, predictions, average="macro", zero_division=0
+            )
+        ),
+        "f1_weighted": float(
+            f1_score(
+                y_validation, predictions, average="weighted", zero_division=0
+            )
+        ),
+        "precision_macro": float(
+            precision_score(
+                y_validation, predictions, average="macro", zero_division=0
+            )
+        ),
+        "recall_macro": float(
+            recall_score(
+                y_validation, predictions, average="macro", zero_division=0
+            )
+        ),
+        "epochs_completed": len(history),
+        "training_samples": int(len(train_index)),
+        "validation_samples": int(len(validation_index)),
+        "device": str(device),
+    }
+    predictions_frame = pd.DataFrame(
+        {
+            "sample_row": validation_index,
+            "observed": [
+                workspace.label_names[value] for value in y_validation
+            ],
+            "predicted": [
+                workspace.label_names[value] for value in predictions
+            ],
+        }
+    )
+    for index, label in enumerate(workspace.label_names):
+        predictions_frame[f"P({label})"] = probabilities[:, index]
+
+    gene = workspace.gene_branch
+    dna = workspace.dna_branch
+    config = {
+        "format_version": 1,
+        "source_name": workspace.source_name,
+        "gene_features": gene.input_genes,
+        "dna_features": dna.input_genes,
+        "label_names": workspace.label_names,
+        "gene_hidden_genes": gene.hidden_genes,
+        "dna_hidden_genes": dna.hidden_genes,
+        "gene_pathways": gene.pathways,
+        "dna_pathways": dna.pathways,
+        "hyperparameters": asdict(hyperparameters),
+        "architecture": {
+            "projection_dim": hyperparameters.projection_dim,
+            "fusion_dim": hyperparameters.fusion_dim,
+            "dropout": hyperparameters.dropout,
+            "biological_activation": "relu",
+            "projection_activation": "sigmoid",
+            "fusion_activation": "tanh",
+        },
+    }
+    bundle = ModelBundle(
+        model=model.cpu(),
+        gene_features=gene.input_genes,
+        dna_features=dna.input_genes,
+        label_names=workspace.label_names,
+        gene_mean=gene_mean,
+        gene_scale=gene_scale,
+        dna_mean=dna_mean,
+        dna_scale=dna_scale,
+        config=config,
+        metrics=metrics,
+        history=history,
+    )
+    return TrainingResult(
+        bundle=bundle,
+        validation_predictions=predictions_frame,
+        confusion=confusion,
+    )
+
+
+def validate_prediction_frames(
+    gene_frame: pd.DataFrame,
+    dna_frame: pd.DataFrame,
+    bundle: ModelBundle,
+) -> tuple[np.ndarray, np.ndarray]:
+    if len(gene_frame) != len(dna_frame):
+        raise ValueError(
+            "Prediction gene-expression and DNA-methylation files must have "
+            "the same number of rows."
+        )
+    missing_gene = sorted(set(bundle.gene_features) - set(gene_frame.columns))
+    missing_dna = sorted(set(bundle.dna_features) - set(dna_frame.columns))
+    if missing_gene or missing_dna:
+        details = []
+        if missing_gene:
+            details.append(
+                "gene-expression: " + ", ".join(missing_gene[:8])
+            )
+        if missing_dna:
+            details.append("DNA-methylation: " + ", ".join(missing_dna[:8]))
+        raise ValueError(
+            "Prediction files are missing trained features (" + "; ".join(details) + ")."
+        )
+    gene_values = gene_frame.loc[:, bundle.gene_features].apply(
+        pd.to_numeric, errors="coerce"
+    )
+    dna_values = dna_frame.loc[:, bundle.dna_features].apply(
+        pd.to_numeric, errors="coerce"
+    )
+    if gene_values.isna().any().any() or dna_values.isna().any().any():
+        raise ValueError("Prediction inputs contain missing or non-numeric values.")
+    return (
+        _scale(gene_values.to_numpy(), bundle.gene_mean, bundle.gene_scale),
+        _scale(dna_values.to_numpy(), bundle.dna_mean, bundle.dna_scale),
+    )
+
+
+def predict(
+    gene_frame: pd.DataFrame,
+    dna_frame: pd.DataFrame,
+    bundle: ModelBundle,
+) -> pd.DataFrame:
+    gene_values, dna_values = validate_prediction_frames(
+        gene_frame, dna_frame, bundle
+    )
+    bundle.model.eval()
+    with torch.no_grad():
+        logits = bundle.model(
+            torch.from_numpy(gene_values), torch.from_numpy(dna_values)
+        )
+        probabilities = torch.softmax(logits, dim=1).numpy()
+    predicted = probabilities.argmax(axis=1)
+    output = pd.DataFrame(
+        {
+            "sample_row": np.arange(len(gene_frame)),
+            "predicted_class": [
+                bundle.label_names[index] for index in predicted
+            ],
+            "confidence": probabilities.max(axis=1),
+        }
+    )
+    for index, label in enumerate(bundle.label_names):
+        output[f"P({label})"] = probabilities[:, index]
+    return output
+
+
+def pathway_importance(bundle: ModelBundle) -> pd.DataFrame:
+    bundle.model.eval()
+    with torch.no_grad():
+        attention = bundle.model.pathway_attention()
+    records: list[dict[str, float | str]] = []
+    for branch_name, config_key in (
+        ("Gene expression", "gene_pathways"),
+        ("DNA methylation", "dna_pathways"),
+    ):
+        key = "gene_expression" if branch_name == "Gene expression" else "dna_methylation"
+        weights = attention[key].cpu().numpy()
+        pathways = bundle.config[config_key]
+        peak_attention = weights.max(axis=0)
+        entropy = -(
+            weights * np.log(np.clip(weights, 1e-12, None))
+        ).sum(axis=0)
+        for pathway, peak, entropy_value in zip(
+            pathways, peak_attention, entropy, strict=True
+        ):
+            records.append(
+                {
+                    "branch": branch_name,
+                    "pathway": pathway,
+                    "peak_gene_attention": float(peak),
+                    "attention_entropy": float(entropy_value),
+                }
+            )
+    return (
+        pd.DataFrame(records)
+        .sort_values("peak_gene_attention", ascending=False)
+        .reset_index(drop=True)
+    )

requirements-dev.txt

@@ -0,0 +1,2 @@
+-r requirements-local.txt
+pytest>=8,<10

requirements-local.txt

@@ -0,0 +1,2 @@
+-r requirements.txt
+spaces

requirements.txt

@@ -0,0 +1,10 @@
+gradio>=6,<7
+huggingface-hub>=0.27,<2
+numpy>=1.26,<3
+pandas>=2.1,<3
+plotly>=5.20,<7
+pyarrow>=15,<24
+safetensors>=0.4,<1
+scikit-learn>=1.4,<2
+scipy>=1.11,<2
+torch>=2.8,<3

tests/test_core.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import pickle
+
+import numpy as np
+import pandas as pd
+import torch
+
+from biolmnet.artifacts import load_bundle, save_bundle
+from biolmnet.data import (
+    BranchPriors,
+    PreparedWorkspace,
+    attach_embeddings_and_pathways,
+    build_biological_mask,
+    deterministic_gene_embeddings,
+)
+from biolmnet.model import BioMaskedLinear, BioLMNet
+from biolmnet.training import Hyperparameters, predict, train
+
+
+def test_biological_mask_uses_pdi_and_undirected_ppi() -> None:
+    pdi = pd.DataFrame(
+        {"TF": ["A", "C", "outside"], "Target": ["B", "D", "A"]}
+    )
+    ppi = pd.DataFrame(
+        {
+            "protein1": ["A", "X", "B"],
+            "protein2": ["X", "C", "Y"],
+            "combined_score": [950, 950, 710],
+        }
+    )
+    branch = build_biological_mask(["A", "B", "C", "D"], pdi, ppi)
+
+    assert {"B", "D", "X"}.issubset(branch.hidden_genes)
+    x_index = branch.hidden_genes.index("X")
+    assert branch.biological_mask[0, x_index] > 0
+    assert branch.biological_mask[2, x_index] > 0
+    assert branch.pdi_edges == 2
+
+
+def test_masked_linear_disconnects_unlisted_weights() -> None:
+    mask = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+    layer = BioMaskedLinear(mask, bias=False)
+    values = torch.tensor([[2.0, 3.0]])
+    baseline = layer(values).detach().clone()
+    with torch.no_grad():
+        layer.weight[0, 1] = 10_000
+        layer.weight[1, 0] = -10_000
+    changed = layer(values).detach()
+    assert torch.allclose(baseline, changed)
+
+
+def _tiny_workspace() -> PreparedWorkspace:
+    rng = np.random.default_rng(7)
+    sample_count = 48
+    labels = np.repeat(np.array([0, 1]), sample_count // 2)
+    gene_values = rng.normal(size=(sample_count, 4)).astype(np.float32)
+    dna_values = rng.normal(size=(sample_count, 4)).astype(np.float32)
+    gene_values[:, 0] += labels * 1.5
+    dna_values[:, 1] -= labels * 1.2
+    input_genes = ["A", "B", "C", "D"]
+    hidden_genes = ["A", "B", "C"]
+    biological = np.array(
+        [
+            [1.0, 1.0, 0.0],
+            [0.0, 1.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [0.0, 1.0, 0.0],
+        ],
+        dtype=np.float32,
+    )
+    pathway_mapping = pd.DataFrame(
+        {
+            "SYMBOL": ["A", "B", "B", "C"],
+            "PathwayID": ["hsa1", "hsa1", "hsa2", "hsa2"],
+        }
+    )
+    embeddings = deterministic_gene_embeddings(hidden_genes, dimensions=8)
+
+    def branch() -> BranchPriors:
+        value = BranchPriors(
+            input_genes=input_genes.copy(),
+            hidden_genes=hidden_genes.copy(),
+            biological_mask=biological.copy(),
+            pdi_edges=3,
+            ppi_edges=4,
+        )
+        attach_embeddings_and_pathways(
+            value,
+            embeddings,
+            pathway_mapping,
+            precomputed_significant=True,
+        )
+        return value
+
+    return PreparedWorkspace(
+        gene_expression=gene_values,
+        dna_methylation=dna_values,
+        labels=labels,
+        label_names=["control", "case"],
+        gene_branch=branch(),
+        dna_branch=branch(),
+        source_name="unit test",
+    )
+
+
+def test_model_forward_probabilistic_shape() -> None:
+    workspace = _tiny_workspace()
+    gene = workspace.gene_branch
+    dna = workspace.dna_branch
+    model = BioLMNet(
+        torch.from_numpy(gene.biological_mask),
+        torch.from_numpy(dna.biological_mask),
+        torch.from_numpy(gene.embeddings),
+        torch.from_numpy(dna.embeddings),
+        torch.from_numpy(gene.pathway_mask),
+        torch.from_numpy(dna.pathway_mask),
+        n_classes=2,
+        projection_dim=4,
+        fusion_dim=3,
+        dropout=0.0,
+    )
+    logits = model(torch.randn(5, 4), torch.randn(5, 4))
+    assert logits.shape == (5, 2)
+    assert torch.allclose(
+        model.gene_branch.pathway_attention.attention_weights().sum(dim=0),
+        torch.ones(2),
+    )
+
+
+def test_training_artifact_roundtrip(tmp_path) -> None:
+    workspace = _tiny_workspace()
+    result = train(
+        workspace,
+        Hyperparameters(
+            epochs=3,
+            batch_size=8,
+            projection_dim=4,
+            fusion_dim=3,
+            dropout=0.0,
+            early_stopping_patience=3,
+        ),
+    )
+    restored_from_process_boundary = pickle.loads(pickle.dumps(result.bundle))
+    assert restored_from_process_boundary.label_names == ["control", "case"]
+    artifact = save_bundle(result.bundle, tmp_path / "model.zip")
+    restored = load_bundle(artifact)
+    gene_frame = pd.DataFrame(
+        workspace.gene_expression[:5], columns=restored.gene_features
+    )
+    dna_frame = pd.DataFrame(
+        workspace.dna_methylation[:5], columns=restored.dna_features
+    )
+    before = predict(gene_frame, dna_frame, result.bundle)
+    after = predict(gene_frame, dna_frame, restored)
+    probability_columns = [column for column in before if column.startswith("P(")]
+    np.testing.assert_allclose(
+        before[probability_columns].to_numpy(),
+        after[probability_columns].to_numpy(),
+        atol=1e-6,
+    )
+
+
+def test_zerogpu_duration_estimator_is_bounded_and_scales() -> None:
+    from app import estimate_training_duration
+
+    workspace = _tiny_workspace()
+    common = (workspace, 16, 0.001, 0.01, 0.3, 64, 12, 0.2, "Adam", True)
+    short = estimate_training_duration(common[0], 10, *common[1:])
+    long = estimate_training_duration(common[0], 200, *common[1:])
+    assert 30 <= short <= 300
+    assert short <= long <= 300