Release prep: ensemble-aware CLI/server, clean-install validation, compact Run locally section and local inference doc
5f15b97 verified | """Two-checkpoint ensemble for the Typed Decisions specialist.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from .schema import KIND_TO_ID, parse_questions | |
| from .task_specialist import TaskSpecialistEngine | |
| class EnsembleEngine: | |
| def from_pretrained(cls, pretrained_model_name_or_path: str, device: str | None = None, token: str | None = None, **kwargs) -> "EnsembleEngine": | |
| return cls(pretrained_model_name_or_path, device=device, token=token) | |
| def __init__(self, config_path: str, device: str | None = None, token: str | None = None): | |
| path = config_path | |
| if not os.path.exists(path): | |
| from huggingface_hub import snapshot_download | |
| path = snapshot_download(path, token=token) | |
| if os.path.isdir(path): | |
| candidate = os.path.join(path, "ensemble_config.json") | |
| path = candidate if os.path.exists(candidate) else path | |
| base = os.path.dirname(os.path.abspath(path)) | |
| cfg = json.load(open(path, encoding="utf-8")) | |
| self.config_path = path | |
| self.device = device | |
| def resolve(p: str) -> str: | |
| if os.path.isabs(p) or os.path.exists(p): | |
| return p | |
| return os.path.join(base, p) | |
| self.engines = [TaskSpecialistEngine(resolve(p), device=device) for p in cfg["models"]] | |
| self.weights = cfg.get("weights", [1.0 / len(self.engines)] * len(self.engines)) | |
| self.temperatures = cfg.get("temperature", {}) | |
| self.model_name = cfg.get("name", "pluto-typed-ensemble") | |
| self.max_len = cfg.get("max_len", 512) | |
| def info(self) -> dict: | |
| return { | |
| "name": self.model_name, | |
| "models": [e.info["name"] for e in self.engines], | |
| "weights": self.weights, | |
| "temperatures": self.temperatures, | |
| "device": str(self.engines[0].device) if self.engines else str(self.device), | |
| "params": sum(e.info["params"] for e in self.engines), | |
| } | |
| def _blend(self, per_model: list[list[dict]], items) -> list[dict]: | |
| out = [] | |
| for i, (_, q) in enumerate(items): | |
| ps = [] | |
| for model_results in per_model: | |
| ps.append(model_results[i]["probs"][: q.n_options]) | |
| p = np.zeros(q.n_options, dtype=np.float64) | |
| for w, pi in zip(self.weights, ps): | |
| p += w * pi | |
| p = p / max(p.sum(), 1e-9) | |
| t = self.temperatures.get(q.kind, 1.0) | |
| if t != 1.0: | |
| p = np.power(np.clip(p, 1e-12, 1.0), 1.0 / t) | |
| p = p / p.sum() | |
| top = int(np.argmax(p)) | |
| if q.kind == "choice": | |
| value = q.keys[top] | |
| elif q.kind == "noul": | |
| value = float(p[1]) | |
| else: | |
| value = float(np.dot(np.arange(q.n_options), p)) | |
| out.append( | |
| { | |
| "index": i, | |
| "top": top, | |
| "value": value, | |
| "probs": p.astype(np.float32), | |
| "confidence": float(p[top]), | |
| "accept_probability": float(p[top]), | |
| } | |
| ) | |
| return out | |
| def predict_many(self, items, batch_size: int = 64, max_len: int | None = None) -> list[dict]: | |
| items = list(items) | |
| per_model = [e.predict_many(items, batch_size=batch_size, max_len=max_len) for e in self.engines] | |
| return self._blend(per_model, items) | |
| def predict(self, state, questions, max_len: int | None = None, alpha: float | None = None) -> dict: | |
| qs = parse_questions(questions) if not isinstance(questions, list) else questions | |
| items = [(state, q) for q in qs] | |
| results = self.predict_many(items, max_len=max_len) | |
| answers = {} | |
| escalate = False | |
| for i, q in enumerate(qs): | |
| p = results[i]["probs"][: q.n_options] | |
| top = int(np.argmax(p)) | |
| conf = float(p[top]) | |
| item = { | |
| "type": q.kind, | |
| "confidence": conf, | |
| "accept_probability": conf, | |
| "escalate": conf < 0.5, | |
| "probabilities": {k: float(v) for k, v in zip(q.keys, p)}, | |
| } | |
| if q.kind == "choice": | |
| item["choice"] = q.keys[top] | |
| elif q.kind == "noul": | |
| item["noul"] = float(p[1]) | |
| else: | |
| item["score"] = float(np.dot(np.arange(q.n_options), p)) | |
| item["score_normalized"] = item["score"] / max(1, q.n_options - 1) | |
| answers[q.name] = item | |
| escalate = escalate or item["escalate"] | |
| return {"answers": answers, "escalate": escalate, "routing": {"model": self.model_name}} | |