| """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 = [] |
| 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 |
|
|
