File size: 5,024 Bytes
cd02163 | 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 | """High-level API: Decider.from_pretrained(...).decide(request or [requests])."""
import os
import torch
from . import _core as C
class Decider:
"""A loaded Decider 2B model.
decide() takes one Jev/SystemOne request (returns one response) or a list (returns a list in the
same order). Questions that do not fit next to the full state in one pass are split across
passes, each with the whole state; nothing is truncated. Passes from all requests are batched
on the device under a padded-token budget.
"""
def __init__(self, model, name="decider-2b", batch_tokens=32768):
self.model = model
self.name = name
self.batch_tokens = int(batch_tokens)
@classmethod
def from_pretrained(cls, model="Jeetu456/decider-2b", device=None, dtype=None, quantization=None,
abstain_threshold=None, revision=None, token=None, batch_tokens=32768, name="decider-2b"):
"""Load from a Hugging Face repo id or a local folder.
device: "cuda", "mps" or "cpu" (default: cuda if available, else cpu).
quantization: None, "8bit" or "4bit" (bitsandbytes, CUDA only).
abstain_threshold: override the tuned "I don't know" threshold (0 turns it off).
"""
path = model
if not os.path.isdir(path):
from huggingface_hub import snapshot_download
path = snapshot_download(model, revision=revision, token=token)
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
loaded = C.load_release_model(path, device=device, dtype=dtype, quantization=quantization)
if abstain_threshold is not None:
C.ABSTAIN_THRESHOLD = float(abstain_threshold)
return cls(loaded, name=name, batch_tokens=batch_tokens)
@property
def temperature(self):
return self.model.temperature
@property
def device(self):
return self.model.head_device
def _plan(self, request):
C.validate_request(request)
questions = request["questions"]
state = C.request_state_text(request["state"], questions)
rendered = C.items(questions)
passes, state_tokens, sizes = C.plan(self.model.tokenizer, state, rendered, self.model.max_length)
return questions, state, rendered, passes, state_tokens, sizes
@torch.inference_mode()
def _logits(self, requests):
"""Raw logits per request: a list (one per question, in request order) of option-logit lists."""
plans = [self._plan(r) for r in requests]
units = [] # (request index, item indices, estimated length)
for r, (_, _, _, passes, state_tokens, sizes) in enumerate(plans):
for indices in passes:
units.append((r, indices, state_tokens + sum(sizes[i] for i in indices) + 32))
out = [[None] * len(p[2]) for p in plans]
order = sorted(range(len(units)), key=lambda u: units[u][2])
start = 0
while start < len(order):
end, width = start, 0
while end < len(order):
width = max(width, units[order[end]][2])
if end > start and width * (end - start + 1) > self.batch_tokens:
break
end += 1
chosen = [units[u] for u in order[start:end]]
batch = self.model.encode([{"state": plans[r][1], "items": [plans[r][2][i] for i in idx]}
for r, idx, _ in chosen])
logits = self.model(**{k: v.to(self.device) for k, v in batch.items()}).float().cpu()
row = 0
for r, idx, _ in chosen:
for i in idx:
out[r][i] = logits[row, : len(plans[r][2][i]["labels"])]
row += 1
start = end
return plans, out
def logits(self, requests):
"""Raw option logits before temperature: {question id: [logit per option, in request order]}."""
single = isinstance(requests, dict)
plans, out = self._logits([requests] if single else requests)
result = [{item["id"]: values.tolist() for item, values in zip(plan[2], rows)} for plan, rows in zip(plans, out)]
return result[0] if single else result
def decide(self, requests, temperature=None):
"""Jev/SystemOne response(s): {"model", "answers", "usage"} per request."""
single = isinstance(requests, dict)
batch = [requests] if single else list(requests)
plans, out = self._logits(batch)
t = float(temperature or self.model.temperature)
responses = []
for request, (questions, _, rendered, passes, state_tokens, sizes), rows in zip(batch, plans, out):
probabilities = [(values / t).softmax(-1).tolist() for values in rows]
responses.append(C.response(request.get("model") or self.name, questions, probabilities,
len(passes) * state_tokens + sum(sizes)))
return responses[0] if single else responses
|