File size: 1,956 Bytes
4d43445 | 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 | """Small public wrapper around the bundled Basenji Saluki implementation."""
from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Any, Mapping
def load_params(path: str | Path) -> dict[str, Any]:
"""Load parameters and normalize historical scalar multi-head settings."""
with Path(path).open(encoding="utf-8") as handle:
params = json.load(handle)
model_params = dict(params["model"])
num_targets = model_params.get("num_targets", 1)
heads = int(model_params.get("heads", 1))
if not isinstance(num_targets, list):
model_params["num_targets"] = [num_targets] * heads
normalized = dict(params)
normalized["model"] = model_params
return normalized
class SalukiModel:
"""Construct, restore and run one head of the official Saluki network."""
def __init__(self, params: Mapping[str, Any], head: int = 0):
os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
from .basenji.rnann import RnaNN
model_params = dict(params["model"] if "model" in params else params)
num_targets = model_params.get("num_targets", 1)
heads = int(model_params.get("heads", 1))
if not isinstance(num_targets, list):
model_params["num_targets"] = [num_targets] * heads
self.network = RnaNN(model_params)
if not 0 <= head < len(self.network.models):
raise ValueError(
f"head must be in [0, {len(self.network.models) - 1}], got {head}"
)
self.head = head
self.network.model = self.network.models[head]
@property
def keras_model(self):
return self.network.models[self.head]
def restore(self, weight_path: str | Path) -> "SalukiModel":
self.network.restore(str(weight_path), head_i=self.head)
return self
def predict(self, inputs, **kwargs):
return self.keras_model.predict(inputs, **kwargs)
|