| """Serial inference helpers; no model imports or GPU synchronization when untimed.""" |
| from __future__ import annotations |
|
|
| from contextlib import contextmanager, nullcontext |
| from time import perf_counter |
|
|
| PHASES = ("decode_load", "processor_preprocess", "tokenise_chat_template", |
| "route_adapter_switch", "h2d_copy", "forward", "readout_calibration", "response_build") |
|
|
|
|
| class PhaseTimer: |
| def __init__(self, seconds=None, synchronize=None): |
| self.seconds, self.synchronize = seconds, synchronize |
| if seconds is not None: |
| seconds.update(dict.fromkeys(PHASES, 0.0)) |
|
|
| def phase(self, name): |
| return self._phase(name) if self.seconds is not None else nullcontext() |
|
|
| @contextmanager |
| def _phase(self, name): |
| if self.synchronize is not None: |
| self.synchronize() |
| started = perf_counter() |
| try: |
| yield |
| finally: |
| if self.synchronize is not None: |
| self.synchronize() |
| self.seconds[name] += perf_counter() - started |
|
|
|
|
| class AdapterSwitch: |
| """Keep PEFT's selected adapter fixed; toggle layers only on route transitions. |
| |
| Owns adapter state for a serial reader. Do not mutate adapters externally. |
| """ |
| def __init__(self, model): |
| from open_decisions.image_jev import router |
| self.model = model |
| self.layers = model.base_model |
| model.set_adapter(router.SCREEN_GEOMETRY) |
| self.route = router.SCREEN_GEOMETRY |
|
|
| def __call__(self, route): |
| from open_decisions.image_jev import router |
| if route not in router.ROUTES: |
| raise ValueError("unknown adapter route") |
| if route != self.route: |
| if route == router.SCREEN_GEOMETRY: |
| self.layers.enable_adapter_layers() |
| else: |
| self.layers.disable_adapter_layers() |
| self.route = route |
| return nullcontext() |
|
|
|
|
| def move_to_device(value, device): |
| if device is None or not hasattr(value, "to"): |
| return value |
| if getattr(value, "device", None) == device: |
| return value |
| |
| if getattr(device, "type", str(device).split(":")[0]) == "cuda": |
| if getattr(getattr(value, "device", None), "type", None) == "cpu": |
| value = value.pin_memory() |
| return value.to(device, non_blocking=True) |
| return value.to(device) |
|
|
|
|
| class NativeCache: |
| """Per-processor caches for the immutable letter ids and chat shell.""" |
| def __init__(self, processor): |
| self.processor = processor |
| self.tokenizer = CachedTokenizer(processor.tokenizer) |
| processor.tokenizer = self.tokenizer |
| self.letters = None |
| self.shell = None |
|
|
| def prompt(self, user): |
| from open_decisions.image_jev.vision import IMAGE_SYSTEM |
| processor = self.processor |
| if not hasattr(processor, "apply_chat_template"): |
| return user |
| def render(text): |
| return processor.apply_chat_template( |
| [{"role": "system", "content": IMAGE_SYSTEM}, {"role": "user", "content": text}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| if self.shell is None: |
| marker = "IMAGE_HOPPER_STATIC_SHELL_56" |
| sample = render(marker) |
| parts = sample.split(marker) |
| |
| self.shell = tuple(parts) if len(parts) == 2 and "".join((parts[0], user, parts[1])) == render(user) else False |
| if self.shell is False: |
| return render(user) |
| self.tokenizer.configure(*self.shell) |
| return self.shell[0] + user + self.shell[1] |
|
|
| def letter_ids(self, count): |
| from open_decisions.scoring.prompt import letter_token_ids |
| if self.letters is None: |
| self.letters = letter_token_ids(self.processor.tokenizer) |
| return self.letters[:count] |
|
|
|
|
| class CachedTokenizer: |
| """Cache Qwen chat tokens only at isolated added-special-token boundaries. |
| |
| Unknown kwargs, tokenizer families and shapes use the original tokenizer. |
| The first full encoding also verifies exact token-id equality. |
| """ |
| def __init__(self, tokenizer): |
| self.original = tokenizer |
| self.parts = None |
| self.timer = PhaseTimer() |
|
|
| def __getattr__(self, name): |
| return getattr(self.original, name) |
|
|
| def configure(self, prefix, suffix): |
| if self.parts is not None or not getattr(self.original, "is_fast", False): |
| return |
| if "Qwen" not in type(self.original).__name__ or "<|im_start|>" not in prefix: |
| return |
| boundary = prefix.rfind("<|im_start|>") + len("<|im_start|>") |
| prefix = prefix[:boundary] |
| if not suffix.startswith("<|im_end|>"): |
| return |
| tokens = {str(token): token for token in self.original.added_tokens_decoder.values()} |
| for key in ("<|im_start|>", "<|im_end|>"): |
| token = tokens.get(key) |
| if token is None or not token.special or token.lstrip or token.rstrip: |
| return |
| if self.original.num_special_tokens_to_add(pair=False): |
| return |
| self.parts = (prefix, suffix, self.original.encode(prefix, add_special_tokens=False), |
| self.original.encode(suffix, add_special_tokens=False), False) |
|
|
| def __call__(self, text, **kwargs): |
| with self.timer.phase("tokenise_chat_template"): |
| allowed = {"padding", "return_attention_mask", "return_token_type_ids", "add_special_tokens", "return_tensors"} |
| if (not self.parts or set(kwargs) - allowed or kwargs.get("return_token_type_ids") or |
| not isinstance(text, list) or len(text) != 1 or not isinstance(text[0], str)): |
| return self.original(text, **kwargs) |
| prefix, suffix, before, after, verified = self.parts |
| value = text[0] |
| if not value.startswith(prefix) or not value.endswith(suffix): |
| return self.original(text, **kwargs) |
| middle = value[len(prefix):-len(suffix)] |
| ids = before + self.original.encode(middle, add_special_tokens=False) + after |
| if not verified: |
| reference = self.original(text, **kwargs) |
| expected = reference["input_ids"] |
| if hasattr(expected, "tolist"): |
| expected = expected.tolist() |
| if expected != [ids]: |
| self.parts = False |
| return reference |
| self.parts = (prefix, suffix, before, after, True) |
| return reference |
| return self.original.pad({"input_ids": [ids]}, padding=kwargs.get("padding", False), |
| return_attention_mask=kwargs.get("return_attention_mask"), |
| return_tensors=kwargs.get("return_tensors")) |
|
|
