| """Decision Index engine for FINAL-Bench/Darwin-27B-ZTC.
|
|
|
| Use with the Decision Index kit (github.com/apolinario/decision-index):
|
|
|
| PYTHONPATH=<snapshot of this repo> python -m decision_index pipeline \
|
| --engine ztc_engine:DarwinZTCEngine --option model=FINAL-Bench/Darwin-27B-ZTC --out runs/darwin-27b-ztc
|
|
|
| One request is one row: every question of the row is answered with one forward pass of the
|
| backbone (no generated tokens), and the questions of a row are batched together.
|
| A question whose prompt is longer than max_length tokens is declared unsupported (never truncated).
|
| """
|
|
|
| import math
|
| import sys
|
| import time
|
| from pathlib import Path
|
|
|
| import torch
|
|
|
| from decision_index.engines.base import Engine, Unsupported
|
|
|
|
|
| class DarwinZTCEngine(Engine):
|
| name = "darwin-27b-ztc"
|
| latency = "In-process request wall time (prompt construction, tokenisation and forward passes; excludes model loading)."
|
|
|
| def __init__(self, model="FINAL-Bench/Darwin-27B-ZTC", revision=None, device="cuda", max_length=32768,
|
| batch=8, token_budget=16000, **options):
|
| super().__init__(**options)
|
| path = Path(model)
|
| if not path.exists():
|
| from huggingface_hub import snapshot_download
|
| path = Path(snapshot_download(model, revision=revision))
|
| if str(path) not in sys.path:
|
| sys.path.insert(0, str(path))
|
| from autojev.model import DecisionModel, answer
|
| self._answer = answer
|
| self.model = DecisionModel(checkpoint=path, device=device).eval()
|
| self.max_length, self.batch, self.token_budget = int(max_length), int(batch), int(token_budget)
|
| self.provenance = {"model": model, "revision": revision, "path": str(path), "temperature": self.model.temperature,
|
| "max_length": self.max_length, "batch": self.batch, "token_budget": self.token_budget}
|
|
|
| def runtime(self):
|
| return {"torch": torch.__version__, "cuda": torch.version.cuda,
|
| "device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"}
|
|
|
| def synchronize(self):
|
| if torch.cuda.is_available():
|
| torch.cuda.synchronize()
|
|
|
| @staticmethod
|
| def _est(row):
|
| import json
|
| return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192
|
|
|
| @torch.inference_mode()
|
| def _probs(self, rows):
|
| b = self.model.prepare(rows, max_length=self.max_length)
|
| d = (self.model(b) / self.model.temperature).softmax(-1).cpu().tolist()
|
| return [v[:c] for v, c in zip(d, b.counts)]
|
|
|
| def _one(self, row):
|
| try:
|
| return self._probs([row])[0]
|
| except ValueError as e:
|
| raise Unsupported(str(e))
|
|
|
| def __call__(self, state, questions):
|
| keys = list(questions)
|
| rows = [{"state": state, "question": questions[k]} for k in keys]
|
| order = sorted(range(len(rows)), key=lambda i: self._est(rows[i]))
|
| dists = {}
|
| i = 0
|
| while i < len(order):
|
| group = [order[i]]; tok = self._est(rows[order[i]])
|
| while i + len(group) < len(order) and len(group) < self.batch and tok + self._est(rows[order[i + len(group)]]) <= self.token_budget:
|
| tok += self._est(rows[order[i + len(group)]]); group.append(order[i + len(group)])
|
| i += len(group)
|
| try:
|
| for j, p in zip(group, self._probs([rows[j] for j in group])):
|
| dists[j] = p
|
| except (ValueError, torch.OutOfMemoryError):
|
| torch.cuda.empty_cache()
|
| for j in group:
|
| dists[j] = self._one(rows[j])
|
| answers = {}
|
| for j, k in enumerate(keys):
|
| p = dists[j]
|
| if not all(math.isfinite(x) and x >= 0 for x in p) or sum(p) <= 0:
|
| p = [1.0] * len(p)
|
| answers[k] = self._answer(questions[k], p)
|
| return {"model": self.name, "answers": answers}, None
|
|
|
