| """Shared prefill: encode a state once with the base, then run questions or generation on it. |
| |
| The state is prefilled with the adapter off, which is exactly what the masked adapter computes |
| on state tokens. Each question suffix then runs with the adapter on against that cache, and the |
| cache is cropped back to the state afterwards, so any number of questions, and an escalation to |
| generation, reuse one prefill. |
| """ |
|
|
| import numpy as np |
| import torch |
|
|
| from dmm import lora |
| from dmm.head import option_logits |
| from dmm.template import Template, decision_text |
|
|
|
|
| class StateCache: |
| def __init__(self, model, ids: list[int]): |
| self.model, self.ids = model, list(ids) |
| with torch.no_grad(), lora.disabled(model): |
| out = model.model(input_ids=torch.tensor([ids], device=model.device), use_cache=True) |
| self.kv = out.past_key_values |
|
|
| def __len__(self) -> int: |
| return len(self.ids) |
|
|
| @torch.no_grad() |
| def run_suffix(self, suffix: list[int], adapter: bool = True) -> torch.Tensor: |
| """Final hidden state at the suffix's last token; the cache is cropped back after.""" |
| try: |
| ctx = lora.disabled(self.model) if not adapter else _null() |
| with ctx: |
| out = self.model.model( |
| input_ids=torch.tensor([suffix], device=self.model.device), |
| past_key_values=self.kv, |
| use_cache=True, |
| ) |
| return out.last_hidden_state[0, -1] |
| finally: |
| self.kv.crop(len(self.ids)) |
|
|
|
|
| class _null: |
| def __enter__(self): |
| return None |
|
|
| def __exit__(self, *exc): |
| return False |
|
|
|
|
| class CachedDecisionModel: |
| """decide(items) through the shared-prefill path: one state prefill per distinct state.""" |
|
|
| def __init__(self, model, tokenizer, head, temperatures=None): |
| self.model, self.head, self.tpl = model, head, Template(tokenizer) |
| self.temperatures = temperatures or {} |
| self.prefills = 0 |
|
|
| @torch.no_grad() |
| def option_logits(self, items: list[dict]) -> list[torch.Tensor]: |
| out: list = [None] * len(items) |
| by_state: dict[str, list[int]] = {} |
| for i, it in enumerate(items): |
| by_state.setdefault(it["state"], []).append(i) |
| for _, idx in by_state.items(): |
| cache = StateCache(self.model, self.tpl.state_ids(items[idx[0]])) |
| self.prefills += 1 |
| for i in idx: |
| h = cache.run_suffix(self.tpl.encode(decision_text(items[i]))) |
| out[i] = option_logits(self.head(h[None]), [items[i]])[0].float().cpu() |
| return out |
|
|
| def decide(self, items: list[dict]) -> list[np.ndarray]: |
| out = [] |
| for z, it in zip(self.option_logits(items), items, strict=True): |
| z = z / self.temperatures.get(it["kind"], 1.0) |
| out.append(torch.softmax(z, -1).numpy().astype(np.float64)) |
| return out |
|
|
