dinah-0-apache / decision_index_engine.py

Add Dinah-0 Apache: Dinah-0 retrained without non-commercial data

3265ba1 verified 9 days ago

3.06 kB

"""Decision Index engine for Dinah-0 (https://github.com/apolinario/decision-index).
    python -m decision_index run --engine decision_index_engine:DinahEngine \
        --option revision=<commit> --edition 0.2.1 --suite-dir suite-0.2 --out runs/dinah-0-apache
Put this file next to `dinah.py` (both are in the model repository) and on PYTHONPATH.
Kit rules: nothing is truncated and no option is dropped. A request that does not fit in 8,192 tokens,
or a question type the model does not support, raises `Unsupported`.
One call = one request (a state and several questions); all questions of a request go through the
encoder together, 16 at a time. Probabilities are the softmax over the option logits; the confidence
of a `choice` answer comes from the confidence head.
"""
from __future__ import annotations

import sys
from pathlib import Path

from decision_index.engines import Engine, Unsupported

sys.path.insert(0, str(Path(__file__).resolve().parent))
from dinah import Dinah, RequestTooLong  # noqa: E402


class DinahEngine(Engine):
    name = "dinah-0-apache"
    latency = ("In-process request wall time: tokenization plus one encoder pass per request "
               "(all questions of the request batched together); excludes model loading.")

    def __init__(self, repo: str = "Lukitaduarte/dinah-0-apache", revision: str | None = None, device: str | None = None,
                 dtype: str | None = None, batch_size: int = 16, calibrate: str | bool = True, **options):
        super().__init__(**options)
        self.model = Dinah.from_pretrained(repo, device=device, dtype=dtype, revision=revision)
        self.calibrated = str(calibrate).lower() not in ("0", "false", "no", "off")
        if not self.calibrated:   # raw logits and raw confidence head, for comparison with the calibrated run
            self.model.temperature, self.model.platt = None, None
        self.batch_size = int(batch_size)
        self.provenance = {"repo": repo, "revision": revision or "main", "device": self.model.device,
                           "dtype": self.model.dtype, "max_length": self.model.max_length, "calibrated": self.calibrated,
                           "truncation": "none: requests above 8,192 tokens raise Unsupported"}

    def __call__(self, state, questions):
        keys = list(questions)
        batch = []
        for key in keys:
            q = questions[key]
            if q["type"] not in ("choice", "noul"):
                raise Unsupported(f"question type {q['type']}")
            batch.append({"type": q["type"], "state": state, "instructions": q.get("instructions"),
                          "criteria": q.get("criteria")})
        try:
            answers = self.model.predict(batch, batch_size=self.batch_size)
        except RequestTooLong as e:
            raise Unsupported(f"context window: {e}")
        return {"model": "dinah-0-apache", "answers": dict(zip(keys, answers))}, None

    def synchronize(self):
        if self.model.device.startswith("cuda"):
            self.model.torch.cuda.synchronize()

	"""Decision Index engine for Dinah-0 (https://github.com/apolinario/decision-index).

	python -m decision_index run --engine decision_index_engine:DinahEngine \
	--option revision=<commit> --edition 0.2.1 --suite-dir suite-0.2 --out runs/dinah-0-apache

	Put this file next to `dinah.py` (both are in the model repository) and on PYTHONPATH.
	Kit rules: nothing is truncated and no option is dropped. A request that does not fit in 8,192 tokens,
	or a question type the model does not support, raises `Unsupported`.

	One call = one request (a state and several questions); all questions of a request go through the
	encoder together, 16 at a time. Probabilities are the softmax over the option logits; the confidence
	of a `choice` answer comes from the confidence head.
	"""
	from __future__ import annotations

	import sys
	from pathlib import Path

	from decision_index.engines import Engine, Unsupported

	sys.path.insert(0, str(Path(__file__).resolve().parent))
	from dinah import Dinah, RequestTooLong # noqa: E402


	class DinahEngine(Engine):
	name = "dinah-0-apache"
	latency = ("In-process request wall time: tokenization plus one encoder pass per request "
	"(all questions of the request batched together); excludes model loading.")

	def __init__(self, repo: str = "Lukitaduarte/dinah-0-apache", revision: str \| None = None, device: str \| None = None,
	dtype: str \| None = None, batch_size: int = 16, calibrate: str \| bool = True, **options):
	super().__init__(**options)
	self.model = Dinah.from_pretrained(repo, device=device, dtype=dtype, revision=revision)
	self.calibrated = str(calibrate).lower() not in ("0", "false", "no", "off")
	if not self.calibrated: # raw logits and raw confidence head, for comparison with the calibrated run
	self.model.temperature, self.model.platt = None, None
	self.batch_size = int(batch_size)
	self.provenance = {"repo": repo, "revision": revision or "main", "device": self.model.device,
	"dtype": self.model.dtype, "max_length": self.model.max_length, "calibrated": self.calibrated,
	"truncation": "none: requests above 8,192 tokens raise Unsupported"}

	def __call__(self, state, questions):
	keys = list(questions)
	batch = []
	for key in keys:
	q = questions[key]
	if q["type"] not in ("choice", "noul"):
	raise Unsupported(f"question type {q['type']}")
	batch.append({"type": q["type"], "state": state, "instructions": q.get("instructions"),
	"criteria": q.get("criteria")})
	try:
	answers = self.model.predict(batch, batch_size=self.batch_size)
	except RequestTooLong as e:
	raise Unsupported(f"context window: {e}")
	return {"model": "dinah-0-apache", "answers": dict(zip(keys, answers))}, None

	def synchronize(self):
	if self.model.device.startswith("cuda"):
	self.model.torch.cuda.synchronize()