biolmnet/artifacts.py

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", []),
    )