File size: 2,991 Bytes
189e0ad e697bd2 189e0ad e697bd2 189e0ad | 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 model weights and PCA components."""
import json
from pathlib import Path
def shared_resources(directory, resources_directory):
if resources_directory is not None:
return Path(resources_directory)
for parent in directory.parents:
resources = parent / "public-resources"
if resources.is_dir():
return resources
raise FileNotFoundError(
"Shared public-resources directory is missing; supply resources_directory explicitly"
)
def load_model(model_directory, adata=None, resources_directory=None):
directory = Path(model_directory).resolve()
if (directory / "encoder.ckpt").is_file():
from _native_scimilarity import load_scimilarity
return load_scimilarity(directory)
config = json.loads((directory / "config.json").read_text())
if config.get("model_type") == "bert":
from transformers import BertForMaskedLM
model, info = BertForMaskedLM.from_pretrained(
directory, local_files_only=True, output_loading_info=True,
attn_implementation="eager",
)
if any(info[key] for key in (
"missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs",
)):
raise ValueError(f"Geneformer strict loading failed: {info}")
model.eval()
return model
family = config.get("family")
if family == "PCA":
import numpy as np
components = np.load(directory / "components.npy", allow_pickle=False)
if list(components.shape) != config["shape"] or str(components.dtype) != config["dtype"]:
raise ValueError("PCA components disagree with the checkpoint config")
if components.dtype.kind != "f" or not np.isfinite(components).all():
raise ValueError("PCA components are not finite numeric data")
return components
from safetensors.torch import load_file
weights = load_file(str(directory / "model.safetensors"), device="cpu")
if family == "SSL":
from _native_ssl import native_ssl_class
resources = shared_resources(directory, resources_directory)
model = native_ssl_class(resources)(**config["constructor"])
model.load_state_dict(weights, strict=True)
model.eval()
return model
if family == "scVI":
import scvi
if adata is None:
raise ValueError("scVI requires AnnData with raw counts in X and the exact gene order")
if list(adata.var_names) != config["gene_names"]:
raise ValueError("AnnData gene order does not match the checkpoint config")
scvi.model.SCVI.setup_anndata(adata, **config["setup_anndata"])
model = scvi.model.SCVI(adata, **config["constructor"])
model.module.load_state_dict(weights, strict=True)
model.module.eval()
model.is_trained = True
model.to_device("cpu")
return model
raise ValueError(f"Unsupported portable model family: {family}")
|