| """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}") |
|
|
