File size: 4,979 Bytes
2778af7 d9a201b 2778af7 5f15b97 2778af7 | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """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:
@classmethod
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)
@property
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}}
|