| """Letter readout of a stock chat model for the Decision Index. Nothing is generated. |
| |
| The baseline prompt is the "chat layout" of decider-ai 1.6.0 (Apache-2.0, github.com/Mapika/decider), which is how the |
| board entry "Decider chat · Gemma-4-31B" reads the stock model. One row per question: |
| |
| <chat template head>Context:\n<state>\n\nQuestion: <instructions>\nOptions:\n(A) ...\n(B) ...<template tail, thinking off>Answer: ( |
| |
| The answer is read at the final " (": the last hidden state there is multiplied by the output-embedding rows of the |
| labels of the question's options, in float32, and the result goes through a softmax at a temperature set by the number |
| of options. These are the logits before Gemma's final soft cap (30 * tanh(x / 30)); the capped logits the model returns |
| squeeze every confident answer into the same narrow band below 30. Up to 10 options are labelled A..J; above that every |
| option gets its own single-token label (A..Z, then two-letter strings), up to 255. |
| |
| By default the state is read once per request and the request's questions are read side by side after it in one pass |
| (pack_questions); each question sees only the state and itself, so this equals one full row per question up to bf16 |
| rounding. With read_twice the user turn is written twice, as shisa-ai/shisa-de-2 does: |
| |
| Context:\n<state>\n\nQuestion: ...\nOptions: ...\n\nRead the same input again before answering:\nContext:\n<state>\n\nQuestion: ...\nOptions: ... |
| |
| so the second copy of the state is read with the question and its options already in view. Reading twice costs about |
| twice the tokens, so it can be limited: twice_max_tokens reads twice only requests of at most that many tokens read once |
| (state plus all its questions), and reread_below re-reads, written twice, the questions of the other requests whose most |
| likely option has a probability below that value after the first read; the state stays cached between the two reads. A |
| row that written twice would pass the token limit is read once. |
| |
| python -m decision_index run --engine jevlab.reader:Reader --rows sample.jsonl.gz --out runs/reader-sample |
| """ |
| import json |
| import math |
| import string |
|
|
| from decision_index.engines.base import Engine, Unsupported |
|
|
| LETTERS = "ABCDEFGHIJ" |
| MAX_OPTIONS = 255 |
| ANNOTATE_MIN = 8 |
| NOUL_WITHOUT_INSTRUCTIONS = "Which answer fits the context?" |
| READ_AGAIN = "\n\nRead the same input again before answering:\n" |
| QUESTIONS_HEAD = "Questions that will be asked about the context:\n" |
| QUESTIONS_ITEM = "- " |
| LAYOUTS = ("decider", "header") |
| SENTINEL = "@@JEVLAB_USER_CONTENT@@" |
| THINKING_BLOCKS = (("<think>\n", "</think>\n\n"), ("<think>", "</think>\n\n"), ("<|channel>thought\n", "<channel|>")) |
|
|
|
|
| def _txt(v): |
| return v if isinstance(v, str) else json.dumps(v, ensure_ascii=False) |
|
|
|
|
| def _flag(v): |
| return v.strip().lower() in ("1", "true", "yes", "on") if isinstance(v, str) else bool(v) |
|
|
|
|
| def annotate_indices(x, min_len=ANNOTATE_MIN): |
| """Write each element's position into long arrays, so a question naming `records[47]` is a lookup, not a count.""" |
| if isinstance(x, list): |
| if len(x) >= min_len: |
| return [{"_index": i, **annotate_indices(v, min_len)} if isinstance(v, dict) else {"_index": i, "value": annotate_indices(v, min_len)} for i, v in enumerate(x)] |
| return [annotate_indices(v, min_len) for v in x] |
| if isinstance(x, dict): |
| return {k: annotate_indices(v, min_len) for k, v in x.items()} |
| return x |
|
|
|
|
| def render_state(state): |
| return state if isinstance(state, str) else json.dumps(annotate_indices(state), ensure_ascii=False) |
|
|
|
|
| def render_question(q): |
| """-> (question text, option texts, option keys). A noul question is the two options no/yes, keyed False/True.""" |
| kind, crit, raw = q["type"], q.get("criteria"), q.get("instructions", "") |
| if kind == "choice": |
| if isinstance(crit, (list, tuple)): |
| crit = {str(c): None for c in crit} |
| if not isinstance(crit, dict) or not 2 <= len(crit) <= MAX_OPTIONS: |
| raise Unsupported(f"options per choice: a map of 2..{MAX_OPTIONS} options is required") |
| keys = list(crit) |
| return _txt(raw), [k if crit[k] in (None, "") else f"{k}: {_txt(crit[k])}" for k in keys], keys |
| if kind == "noul": |
| crit = crit if isinstance(crit, dict) else {} |
| no, yes = crit.get("false", crit.get(False)), crit.get("true", crit.get(True)) |
| text = NOUL_WITHOUT_INSTRUCTIONS if raw in (None, "") else _txt(raw) |
| return text, ["no" if no in (None, "") else f"no: {_txt(no)}", "yes" if yes in (None, "") else f"yes: {_txt(yes)}"], [False, True] |
| raise Unsupported("Unsupported question type " + str(kind)) |
|
|
|
|
| class Layout: |
| """Token ids of the prompt rows. Needs only the tokenizer.""" |
|
|
| def __init__(self, tok, system=None): |
| if not getattr(tok, "chat_template", None): |
| raise ValueError("the tokenizer has no chat template") |
| self.tok, self.system = tok, system or None |
| kw = {"enable_thinking": False} if "enable_thinking" in tok.chat_template else {} |
| turns = ([{"role": "system", "content": self.system}] if self.system else []) + [{"role": "user", "content": SENTINEL}] |
| head, tail = tok.apply_chat_template(turns, tokenize=False, add_generation_prompt=True, **kw).split(SENTINEL) |
| if self.system and self.system not in head: |
| raise ValueError("the chat template dropped the system turn") |
| for opened, closed in THINKING_BLOCKS: |
| if tail.endswith(opened): |
| tail += closed |
| self.head_text, self.tail_text = head, tail |
| self.head, self.tail = self._enc(head), self._enc(tail) |
| self.answer = self._enc("Answer: (") |
| self.open = self._enc("\n(") |
| names = list(string.ascii_uppercase) + [a + b for a in string.ascii_uppercase for b in string.ascii_uppercase] |
| single = [(n, t[0]) for n, t in ((n, self._enc(n)) for n in names) if len(t) == 1][:MAX_OPTIONS] |
| self.labels, self.label_ids = [n for n, _ in single], [i for _, i in single] |
| if len(self.label_ids) != MAX_OPTIONS or len(set(self.label_ids)) != MAX_OPTIONS: |
| raise ValueError("the tokenizer does not have 255 distinct single-token labels") |
| |
| for n, i in single[: len(LETTERS)]: |
| if self._enc("Answer: (" + n) != self.answer + [i]: |
| raise ValueError(f"the label {n!r} is not one token after 'Answer: ('") |
| if self._enc(tail + "Answer: (") != self.tail + self.answer: |
| raise ValueError("the chat template tail merges with the answer slot") |
|
|
| def _enc(self, text): |
| return self.tok.encode(text, add_special_tokens=False) |
|
|
| @staticmethod |
| def listing(texts): |
| """The header layout's list of every question of the request, placed before the state.""" |
| return QUESTIONS_HEAD + "\n".join(QUESTIONS_ITEM + t for t in texts) + "\n\n" |
|
|
| def context(self, state, texts=None): |
| """Ids up to the end of the state. With `texts` (the request's question texts) the header layout: the list of |
| questions, then the state, each part encoded on its own as wcole3/jev-brain builds it.""" |
| ids = self.head + (self._enc(self.listing(texts)) if texts is not None else []) |
| return ids + self._enc("Context:\n" + render_state(state)) |
|
|
| def again(self, state, texts=None): |
| """Ids that start the second copy of the user turn, up to the end of its state.""" |
| return self._enc(READ_AGAIN + (self.listing(texts) if texts is not None else "") + "Context:\n" + render_state(state)) |
|
|
| def question(self, text, options, again=None): |
| """Ids after the context: the question, its labelled options, the template tail and the answer slot. With `again` |
| (from again()), the user turn is written a second time before the tail.""" |
| ids = self._enc(f"\n\nQuestion: {text}\nOptions:") |
| if len(options) <= len(LETTERS): |
| ids += self._enc("".join(f"\n({LETTERS[j]}) {o}" for j, o in enumerate(options))) |
| else: |
| for j, o in enumerate(options): |
| ids += self.open + [self.label_ids[j]] + self._enc(f") {o}") |
| if again: |
| ids = ids + again + ids |
| return ids + self.tail + self.answer |
|
|
|
|
| class Reader(Engine): |
| name = "jevlab-reader" |
| latency = "In-process request wall time including prompt construction and reuse of the state's KV cache across the questions of one request; excludes model loading." |
|
|
| def __init__(self, model="google/gemma-4-31B-it", revision="842da3794eaa0b77d5f08bae87a17459d91ff475", device=None, dtype="bfloat16", temperature=10.124, temperature_log_slope=-1.633, |
| max_tokens=32768, chunk_tokens=8192, cache_prefix=False, cudnn_sdp=False, pack_questions=True, pack_tokens=4096, read_twice=False, twice_max_tokens=0, reread_below=0.0, reverse_below=0.0, layout="decider", attn_block=512, **options): |
| super().__init__(**options) |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| self.torch = torch |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| |
| |
| |
| if not cudnn_sdp and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): |
| torch.backends.cuda.enable_cudnn_sdp(False) |
| self.tok = AutoTokenizer.from_pretrained(model, revision=revision) |
| self.layout = Layout(self.tok) |
| self.model = AutoModelForCausalLM.from_pretrained(model, revision=revision, dtype=getattr(torch, dtype)).to(self.device).eval() |
| self.decoder = self.model.model |
| head = self.model.get_output_embeddings() |
| if head.bias is not None: |
| raise ValueError("the output layer has a bias; the label readout assumes it has none") |
| self.label_rows = head.weight.detach()[torch.tensor(self.layout.label_ids, device=self.device)].float() |
| self.attn_block = int(attn_block) |
| if self.attn_block > 0: |
| from jevlab.attention import use_blocked_attention |
| use_blocked_attention(self.model, self.attn_block) |
| self.model_id = model |
| |
| |
| self.temperature, self.temperature_log_slope = float(temperature), float(temperature_log_slope) |
| self.limit = int(max_tokens) |
| self.chunk = int(chunk_tokens) |
| self.cache_prefix = _flag(cache_prefix) |
| self.pack, self.pack_tokens = _flag(pack_questions), int(pack_tokens) |
| self.read_twice = _flag(read_twice) |
| self.twice_max_tokens, self.reread_below = int(twice_max_tokens), float(reread_below) |
| self.reverse_below = float(reverse_below) |
| if layout not in LAYOUTS: |
| raise ValueError(f"layout must be one of {LAYOUTS}, not {layout!r}") |
| self.header = layout == "header" |
| cfg = self.model.config.get_text_config() |
| self.layer_types, self.window = set(cfg.layer_types), getattr(cfg, "sliding_window", None) |
| self.cache, self.cached_ids = None, [] |
| self.provenance = { |
| "kind": "inference technique", |
| "repo": model, |
| "revision": revision or getattr(self.model.config, "_commit_hash", None), |
| "device": self.device, |
| "dtype": dtype, |
| "temperature": self.temperature, |
| "temperature_log_slope": self.temperature_log_slope, |
| "max_row_tokens": self.limit, |
| "cudnn_sdp": bool(cudnn_sdp), |
| "pack_questions": self.pack, |
| "read_twice": self.read_twice, |
| "twice_max_tokens": self.twice_max_tokens, |
| "reread_below": self.reread_below, |
| "reverse_below": self.reverse_below, |
| "layout": "header" if self.header else "decider", |
| "attention": f"sdpa in query blocks of {self.attn_block}, each over the span of keys its mask allows" if self.attn_block > 0 else "sdpa", |
| "policy": ("The state is preceded by the list of every question of the request ('Questions that will be asked about the context:'), " |
| "the header layout of wcole3/jev-brain. " if self.header else "") |
| + "One row per question in the model's chat template with thinking off, ending at 'Answer: ('; the last hidden state there times the " |
| "output-embedding rows of the option labels (float32, before the final soft cap) goes through a softmax. Options keep the request's " |
| "order. Rows above the token limit are refused, never cut." |
| + (" The user turn (state, question, options) is written twice, joined by 'Read the same input again before answering:'" |
| + (f", in requests of at most {self.twice_max_tokens} tokens read once" if self.twice_max_tokens > 0 else "") + "." if self.read_twice else "") |
| + (f" Questions of the other requests whose most likely option has probability below {self.reread_below} are read again, written twice " |
| "the same way, after the first read." if self.reread_below > 0 else "") |
| + (f" A question whose most likely option then has probability below {self.reverse_below} is read once more with its options in " |
| "reverse order, and the two distributions are averaged." if self.reverse_below > 0 else "") |
| + (" The state is read once per request; the questions of the request are read side by side after it in one pass, each seeing only " |
| "the state and itself." if self.pack else ""), |
| } |
| if self.pack: |
| self.latency = "In-process request wall time including prompt construction; the state is read once and the request's questions are read together after it; excludes model loading." |
| elif not self.cache_prefix: |
| self.latency = "In-process request wall time including prompt construction; every question row is read in full; excludes model loading." |
|
|
| def runtime(self): |
| torch = self.torch |
| import transformers |
|
|
| info = {"torch": torch.__version__, "transformers": transformers.__version__, "device": self.device} |
| if self.device == "cuda": |
| info.update(cuda=torch.version.cuda, gpu=torch.cuda.get_device_name()) |
| return info |
|
|
| def synchronize(self): |
| if self.device == "cuda": |
| self.torch.cuda.synchronize() |
|
|
| def reset_cache(self): |
| self.cache, self.cached_ids = None, [] |
|
|
| def _forward(self, ids, chunk): |
| """Run `ids`, reusing the cached prefix it shares with the previous row. -> (hidden state at the last position, cached prefix length).""" |
| torch = self.torch |
| from transformers import DynamicCache |
|
|
| p = 0 |
| if self.cache_prefix and self.cache is not None: |
| n = min(len(ids) - 1, len(self.cached_ids)) |
| while p < n and ids[p] == self.cached_ids[p]: |
| p += 1 |
| if p == 0: |
| self.cache = DynamicCache() |
| elif len(self.cached_ids) > p: |
| self.cache.crop(p - len(self.cached_ids)) |
| self.cached_ids = ids[:p] |
| with torch.inference_mode(): |
| for s in range(p, len(ids), chunk): |
| piece = ids[s:s + chunk] |
| out = self.decoder(input_ids=torch.tensor([piece], device=self.device), attention_mask=torch.ones(1, s + len(piece), dtype=torch.long, device=self.device), |
| past_key_values=self.cache, use_cache=True) |
| self.cache = out.past_key_values |
| self.cached_ids = ids[:s + len(piece)] |
| return out.last_hidden_state[0, -1], p |
|
|
| def _pack_masks(self, p, lengths): |
| """Masks for questions laid side by side after a cached prefix of p tokens. Each question token sees the prefix and |
| the earlier tokens of its own question, at the positions it would have in its own row; on sliding-window layers only |
| keys less than `window` positions back, the rule transformers applies to a full row (masking_utils.sliding_window_overlay).""" |
| torch, dev = self.torch, self.device |
| own = torch.repeat_interleave(torch.arange(len(lengths), device=dev), torch.tensor(lengths, device=dev)) |
| qpos = p + torch.cat([torch.arange(m, device=dev) for m in lengths]) |
| kpos = torch.cat([torch.arange(p, device=dev), qpos]) |
| kown = torch.cat([torch.full((p,), -1, device=dev), own]) |
| allowed = (kown[None, :] == -1) | ((kown[None, :] == own[:, None]) & (kpos[None, :] <= qpos[:, None])) |
| masks = {} |
| for t in self.layer_types: |
| if t == "full_attention": |
| masks[t] = allowed[None, None] |
| elif t == "sliding_attention": |
| masks[t] = (allowed & (qpos[:, None] - kpos[None, :] < self.window))[None, None] |
| else: |
| raise ValueError(f"no packed mask for layer type {t!r}") |
| return masks, qpos |
|
|
| def _packed(self, p, pieces): |
| """Hidden states at the last token of each piece, the pieces read side by side after the cached prefix of p tokens.""" |
| torch = self.torch |
| lengths = [len(x) for x in pieces] |
| masks, qpos = self._pack_masks(p, lengths) |
| ids = torch.tensor([[t for x in pieces for t in x]], device=self.device) |
| ends = torch.tensor(lengths, device=self.device).cumsum(0) - 1 |
| try: |
| with torch.inference_mode(): |
| out = self.decoder(input_ids=ids, attention_mask=masks, position_ids=qpos[None], past_key_values=self.cache, use_cache=True) |
| return out.last_hidden_state[0, ends] |
| finally: |
| extra = self.cache.get_seq_length() - p |
| if extra > 0: |
| self.cache.crop(-extra) |
|
|
| def _joint(self, ctx, pieces): |
| """Hidden states at the last token of each piece, the state and all pieces read in one pass: the state causally, |
| each piece after it seeing the state and itself. The state's keys and values stay in the cache afterwards.""" |
| torch, dev = self.torch, self.device |
| from transformers import DynamicCache |
|
|
| p, lengths = len(ctx), [len(x) for x in pieces] |
| own = torch.cat([torch.full((p,), -1, device=dev), torch.repeat_interleave(torch.arange(len(lengths), device=dev), torch.tensor(lengths, device=dev))]) |
| pos = torch.cat([torch.arange(p, device=dev)] + [p + torch.arange(m, device=dev) for m in lengths]) |
| allowed = (pos[None, :] <= pos[:, None]) & ((own[None, :] == -1) | (own[None, :] == own[:, None])) |
| masks = {} |
| for t in self.layer_types: |
| if t == "full_attention": |
| masks[t] = allowed[None, None] |
| elif t == "sliding_attention": |
| masks[t] = (allowed & (pos[:, None] - pos[None, :] < self.window))[None, None] |
| else: |
| raise ValueError(f"no packed mask for layer type {t!r}") |
| ids = torch.tensor([ctx + [t for x in pieces for t in x]], device=dev) |
| ends = p + torch.tensor(lengths, device=dev).cumsum(0) - 1 |
| self.cache, self.cached_ids = DynamicCache(), [] |
| with torch.inference_mode(): |
| out = self.decoder(input_ids=ids, attention_mask=masks, position_ids=pos[None], past_key_values=self.cache, use_cache=True) |
| extra = self.cache.get_seq_length() - p |
| if extra > 0: |
| self.cache.crop(-extra) |
| self.cached_ids = list(ctx) |
| return out.last_hidden_state[0, ends] |
|
|
| def _extend(self, p, piece, chunk): |
| """Hidden state at the last token of one long piece, read in chunks after the cached prefix of p tokens.""" |
| torch = self.torch |
| try: |
| with torch.inference_mode(): |
| for s in range(0, len(piece), chunk): |
| part = piece[s:s + chunk] |
| out = self.decoder(input_ids=torch.tensor([part], device=self.device), attention_mask=torch.ones(1, p + s + len(part), dtype=torch.long, device=self.device), |
| past_key_values=self.cache, use_cache=True) |
| return out.last_hidden_state[0, -1] |
| finally: |
| extra = self.cache.get_seq_length() - p |
| if extra > 0: |
| self.cache.crop(-extra) |
|
|
| def packed_label_logits(self, ctx, pieces, ns, reuse=False): |
| """float32 label logits of every question of a request: the state once, then the questions in packs of about |
| pack_tokens tokens. A question longer than that is read on its own, in chunks. A request of at most pack_tokens |
| tokens in all is read in one pass. With reuse, a state still cached from the first read of the same request is not |
| read again.""" |
| budget = self.pack_tokens |
| while True: |
| try: |
| if not reuse and not self.cache_prefix and len(ctx) + sum(map(len, pieces)) <= budget: |
| hidden = self._joint(ctx, pieces) |
| return [(self.label_rows[:n] @ h.float()).cpu() for h, n in zip(hidden, ns)], 0 |
| if reuse and self.cache is not None and self.cached_ids == ctx and self.cache.get_seq_length() == len(ctx): |
| p0 = len(ctx) |
| else: |
| _, p0 = self._forward(ctx, self.chunk) |
| p = len(ctx) |
| hidden, group, size = {}, [], 0 |
| for i, piece in enumerate(pieces): |
| if len(piece) > budget: |
| hidden[i] = self._extend(p, piece, budget) |
| continue |
| if group and size + len(piece) > budget: |
| hidden.update(zip(group, self._packed(p, [pieces[j] for j in group]))) |
| group, size = [], 0 |
| group.append(i) |
| size += len(piece) |
| if group: |
| hidden.update(zip(group, self._packed(p, [pieces[j] for j in group]))) |
| return [(self.label_rows[:n] @ hidden[i].float()).cpu() for i, n in enumerate(ns)], p0 |
| except self.torch.OutOfMemoryError: |
| self.reset_cache() |
| if self.device == "cuda": |
| self.torch.cuda.empty_cache() |
| if budget <= 256: |
| raise |
| budget //= 2 |
| except Exception: |
| self.reset_cache() |
| raise |
|
|
| def label_logits(self, ids, n): |
| """float32 logits (before the soft cap) of the first n labels at the answer slot of one row.""" |
| torch = self.torch |
| chunk = self.chunk |
| while True: |
| try: |
| hidden, p = self._forward(ids, chunk) |
| return (self.label_rows[:n] @ hidden.float()).cpu(), p |
| except torch.OutOfMemoryError: |
| self.reset_cache() |
| if self.device == "cuda": |
| torch.cuda.empty_cache() |
| if chunk <= 512: |
| raise |
| chunk //= 2 |
| except Exception: |
| self.reset_cache() |
| raise |
|
|
| def _read(self, ctx, pieces, ns, reuse=False): |
| """-> [(float32 label logits, cached prefix tokens)] for the pieces after ctx.""" |
| if self.pack: |
| logits, p0 = self.packed_label_logits(ctx, pieces, ns, reuse) |
| return [(x, p0 if i == 0 else len(ctx)) for i, x in enumerate(logits)] |
| return [self.label_logits(ctx + piece, n) for piece, n in zip(pieces, ns)] |
|
|
| def _probs(self, logits, n): |
| t = max(self.temperature + self.temperature_log_slope * math.log(n), 0.05) |
| return self.torch.softmax(logits.double() / t, -1).tolist() |
|
|
| def __call__(self, state, questions): |
| rendered = [(key, q["type"], *render_question(q)) for key, q in questions.items()] |
| texts = [text for _, _, text, _, _ in rendered] if self.header else None |
| ctx = self.layout.context(state, texts) |
| again = self.layout.again(state, texts) if self.read_twice or self.reread_below > 0 else None |
| rows = [] |
| for key, kind, text, options, keys in rendered: |
| rows.append((key, kind, self.layout.question(text, options), self.layout.question(text, options, again) if again else None, keys)) |
| |
| fits = [again is not None and len(ctx) + len(twice) <= self.limit for _, _, _, twice, _ in rows] |
| once_tokens = len(ctx) + sum(len(once) for _, _, once, _, _ in rows) |
| whole = self.read_twice and all(fits) and (self.twice_max_tokens <= 0 or once_tokens <= self.twice_max_tokens) |
| first = [twice if whole else once for _, _, once, twice, _ in rows] |
| for piece in first: |
| if len(ctx) + len(piece) > self.limit: |
| raise Unsupported(f"maximum context length: a row of {len(ctx) + len(piece)} tokens exceeds the declared limit of {self.limit}") |
| ns = [len(keys) for *_, keys in rows] |
| read = self._read(ctx, first, ns) |
| pieces = list(first) |
| reread = self.reread_below > 0 and not whole |
| unsure = [i for i, ((logits, _), n) in enumerate(zip(read, ns)) if reread and fits[i] and max(self._probs(logits, n)) < self.reread_below] |
| firsts = {} |
| if unsure: |
| again_read = self._read(ctx, [rows[i][3] for i in unsure], [ns[i] for i in unsure], reuse=True) |
| for i, (logits, _) in zip(unsure, again_read): |
| firsts[i] = read[i][0] |
| read[i] = (logits, len(ctx)) |
| pieces[i] = rows[i][3] |
| probs = [self._probs(logits, n) for (logits, _), n in zip(read, ns)] |
| flipped = {} |
| if self.reverse_below > 0: |
| low = [i for i, pr in enumerate(probs) if max(pr) < self.reverse_below] |
| if low: |
| twice = [whole or i in firsts for i in low] |
| rev = [self.layout.question(rendered[i][2], rendered[i][3][::-1], again if t else None) for i, t in zip(low, twice)] |
| for i, piece in zip(low, rev): |
| if len(ctx) + len(piece) > self.limit: |
| raise Unsupported(f"maximum context length: a row of {len(ctx) + len(piece)} tokens exceeds the declared limit of {self.limit}") |
| for i, (logits, _) in zip(low, self._read(ctx, rev, [ns[i] for i in low], reuse=True)): |
| back = self._probs(logits, ns[i])[::-1] |
| flipped[i] = logits |
| probs[i] = [(a + b) / 2 for a, b in zip(probs[i], back)] |
| answers, raw = {}, {} |
| for i, ((key, kind, _, _, keys), (logits, p)) in enumerate(zip(rows, read)): |
| raw[key] = {"label_logits": [round(x, 4) for x in logits.tolist()], "prompt_tokens": len(ctx) + len(pieces[i]), "cached_prefix_tokens": p, |
| "read_twice": whole or i in firsts} |
| if i in firsts: |
| raw[key]["first_label_logits"] = [round(x, 4) for x in firsts[i].tolist()] |
| if i in flipped: |
| raw[key]["reversed_label_logits"] = [round(x, 4) for x in flipped[i].tolist()] |
| probs_i = probs[i] |
| if kind == "noul": |
| answers[key] = {"type": "noul", "noul": probs_i[1]} |
| else: |
| answers[key] = {"type": "choice", "choice": keys[max(range(len(keys)), key=probs_i.__getitem__)], "probabilities": dict(zip(keys, probs_i))} |
| return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw |
|
|
