File size: 2,800 Bytes
30e5a24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Load released SimpleP2G prediction heads without an image encoder."""
from __future__ import annotations
import json
from pathlib import Path
from types import SimpleNamespace
import torch
from safetensors.torch import load_file
from src.simpleP2G import SimpleP2G
DATASETS = ("brain", "prad", "her2st", "lung", "ccrcc")
PANELS = ("hmhvg", "hvg")
SETTINGS = ("single-patch", "cross-patch")
def load_pretrained(
repo_id_or_path="king159/SimpleP2G",
*,
dataset="brain",
panel="hmhvg",
setting="single-patch",
device="cpu",
revision=None,
cache_dir=None,
local_files_only=False,
):
"""Return an eval-mode model, ordered gene names and checkpoint metadata.
``repo_id_or_path`` is a Hub model ID or a downloaded repository directory.
Inputs are precomputed image embeddings; no image-encoder weights are loaded.
Predictions follow the order in ``genes`` and use log2(count + 1) units.
"""
dataset, panel = dataset.lower(), panel.lower()
if dataset not in DATASETS or panel not in PANELS or setting not in SETTINGS:
raise ValueError(f"Unsupported model selection: {dataset}/{panel}/{setting}")
folder = f"models/{setting}/{dataset}/{panel}"
local = Path(repo_id_or_path).expanduser()
def resolve(filename):
if local.is_dir():
path = local / folder / filename
if not path.is_file():
raise FileNotFoundError(path)
return path
from huggingface_hub import hf_hub_download
return Path(hf_hub_download(
repo_id=str(repo_id_or_path), filename=f"{folder}/{filename}",
revision=revision, cache_dir=cache_dir, local_files_only=local_files_only,
))
metadata = json.loads(resolve("config.json").read_text())
if metadata.get("format_version") != 1:
raise ValueError("Unsupported checkpoint metadata version")
expected = {"dataset": dataset, "panel": panel, "setting": setting}
if any(metadata.get(key) != value for key, value in expected.items()):
raise ValueError("Checkpoint metadata does not match the requested model")
genes = json.loads(resolve("genes.json").read_text())
config = SimpleNamespace(**metadata["model_config"])
if len(genes) != config.num_genes or len(set(genes)) != len(genes):
raise ValueError("The ordered gene list must contain one unique name per output")
if setting == "single-patch":
model = SimpleP2G(config)
else:
from src.simpleP2G_neigh import SimpleP2GNeigh
model = SimpleP2GNeigh(config)
state = load_file(str(resolve("model.safetensors")), device="cpu")
model.load_state_dict(state, strict=True)
model.to(device=device).eval().requires_grad_(False)
return model, genes, metadata
|