"""Python request boundary; all model arithmetic and sampling stay in C++.""" from __future__ import annotations from dataclasses import dataclass import math import threading from typing import Any, Iterator @dataclass(frozen=True) class GenerationChunk: text: str token_ids: list[int] finish_reason: str | None prompt_tokens: int completion_tokens: int spec_proposed: int = 0 spec_accepted: int = 0 @dataclass(frozen=True) class GenerationResult: text: str token_ids: list[int] finish_reason: str prompt_tokens: int completion_tokens: int spec_proposed: int = 0 spec_accepted: int = 0 class Engine: """One loaded native model with serialized generation requests. ``from_pretrained`` imports safetensors directly, without loading a Python model. ``from_weights`` is useful for reference tests and local research. """ def __init__(self, native: Any, tokenizer: Any, config: dict, model_id: str, *, revision: str | None = None) -> None: self._native = native self.tokenizer = tokenizer self.config = dict(config) self.model_id = model_id self.revision = revision self.context_length = config["max_position_embeddings"] self._generation_lock = threading.Lock() self._keep_warm_stop = threading.Event() self._keep_warm_thread: threading.Thread | None = None self._keep_warm_error: BaseException | None = None eos = config.get("eos_token_id", getattr(tokenizer, "eos_token_id", None)) self._eos = [] if eos is None else eos if isinstance(eos, list) else [eos] for token in self._eos: self._validate_token(token) @classmethod def from_pretrained(cls, model: str, *, precision: str = "fp32", revision: str | None = None, trust_remote_code: bool = False, local_files_only: bool = False, cache_dir: str | None = None, threads: int = 1, act_precision: str = "fp32") -> Engine: from .loader import import_model loaded = import_model(model, revision=revision, trust_remote_code=trust_remote_code, local_files_only=local_files_only, cache_dir=cache_dir) return cls.from_weights(loaded.config, loaded.weights, loaded.tokenizer, model_id=model, precision=precision, revision=loaded.revision, threads=threads, act_precision=act_precision) @classmethod def from_weights(cls, config: dict, weights: dict, tokenizer: Any, *, model_id: str = "local", precision: str = "fp32", revision: str | None = None, threads: int = 1, act_precision: str = "fp32") -> Engine: from .loader import validate_config, validate_fwkv_config, validate_surjo_config try: from . import _native except ImportError as exc: raise RuntimeError("CISM native extension is missing. Build with: python -m pip install .") from exc if isinstance(config, dict) and config.get("model_type") == "surjo": normalized = validate_surjo_config(config) return cls(_native.SurjoModel(normalized, weights, precision, threads, act_precision), tokenizer, normalized, model_id, revision=revision) if isinstance(config, dict) and config.get("model_type") == "fwkv": normalized = validate_fwkv_config(config) return cls(_native.FwkvModel(normalized, weights, precision, threads, act_precision), tokenizer, normalized, model_id, revision=revision) normalized = validate_config(config) return cls(_native.Model(normalized, weights, precision, threads, act_precision), tokenizer, normalized, model_id, revision=revision) @property def info(self) -> dict: return {**self._native.info, "model": self.model_id, "revision": self.revision, "context_length": self.context_length, "architecture": self.config["model_type"], "capabilities": ["text", "streaming", "sampling"], "keep_warm": self.keep_warm_active} def lock_pages(self) -> int: """Pin weight pages in physical RAM (no OS paging) until unlock_pages. This cannot pin CPU cache lines; hardware decides those. Returns the number of bytes the OS accepted; a zero return means locking failed. """ with self._generation_lock: return self._native.lock_pages() def unlock_pages(self) -> int: with self._generation_lock: return self._native.unlock_pages() def touch(self) -> int: """Run one immediate cache-residency sweep over all weight bytes.""" with self._generation_lock: return self._native.touch() def scan(self) -> int: """Fast full-weight read for bandwidth probes; returns bytes scanned.""" with self._generation_lock: return self._native.scan() def start_keep_warm(self, interval: float = 2.0) -> None: """Refresh weight cache residency every ``interval`` seconds while idle. Stops automatically during generation and resumes afterward. Cache lines evict naturally once stop_keep_warm is called; there is no way to forbid hardware eviction on x86. """ if type(interval) not in (int, float) or not math.isfinite(interval) or not 0 < interval <= 3600: raise ValueError("interval must be a duration in seconds in (0, 3600]") if self._keep_warm_thread is not None: raise RuntimeError("keep-warm is already running") self._keep_warm_error = None self._keep_warm_stop.clear() def sweep() -> None: while not self._keep_warm_stop.wait(interval): # Never compete with an active generation for cache or lock. if not self._generation_lock.acquire(blocking=False): continue try: self._native.touch() except BaseException as error: self._keep_warm_error = error return finally: self._generation_lock.release() self._keep_warm_thread = threading.Thread(target=sweep, name="cism-keep-warm", daemon=True) self._keep_warm_thread.start() def stop_keep_warm(self) -> None: thread = self._keep_warm_thread if thread is None: return self._keep_warm_stop.set() thread.join() self._keep_warm_thread = None @property def keep_warm_active(self) -> bool: thread = self._keep_warm_thread return thread is not None and thread.is_alive() @property def keep_warm_error(self) -> BaseException | None: return self._keep_warm_error def _validate_token(self, token: int) -> None: if type(token) is not int or not 0 <= token < self.config["vocab_size"]: raise ValueError("Token IDs must be integers within the model vocabulary") def encode(self, prompt: str | list[int]) -> list[int]: if isinstance(prompt, str): tokens = list(self.tokenizer.encode(prompt, add_special_tokens=True)) elif isinstance(prompt, list): tokens = prompt.copy() else: raise TypeError("prompt must be text or a list of token IDs") if not tokens: bos = self.config.get("bos_token_id", getattr(self.tokenizer, "bos_token_id", None)) if bos is None: raise ValueError("Empty prompt requires a BOS token") tokens = [bos] for token in tokens: self._validate_token(token) return tokens def format_chat(self, messages: list[dict]) -> list[int]: if not getattr(self.tokenizer, "chat_template", None): raise ValueError("This tokenizer has no chat template; use /v1/completions") rendered = self.tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True) # Most tokenizers return a bare list of IDs, but one whose __call__ # yields a BatchEncoding ({"input_ids", "attention_mask"}) makes # apply_chat_template return that dict instead. if not isinstance(rendered, list): rendered = rendered["input_ids"] return self.encode(list(rendered)) def logits(self, prompt: str | list[int]): """Return last-position FP32 logits for numerical evaluation.""" with self._generation_lock: return self._native.logits(self.encode(prompt)) def nll(self, prompt: str | list[int], window: int = 128) -> dict: """Teacher-forced mean NLL and perplexity over text or token IDs. Scores in truncated windows of at most ``window`` tokens (each window starts with empty context, mirroring HF sliding-window PPL without overlap). Windows are scored by the native batched pass, so results are deterministic for a given precision. """ if type(window) is not int or window < 2: raise ValueError("window must be an integer >= 2") tokens = self.encode(prompt) if len(tokens) < 2: raise ValueError("scoring requires at least 2 tokens") limit = min(window, self.context_length) if limit < 2: raise ValueError("model context is too small for scoring") total, count = 0.0, 0 step = limit - 1 for start in range(0, len(tokens) - 1, step): # +1 tail token so every window ends with a scored target. chunk = tokens[start:start + limit] scored = len(chunk) - 1 with self._generation_lock: total += self._native.nll(chunk) * scored count += scored mean = total / count try: perplexity = math.exp(mean) except OverflowError: perplexity = float("inf") return {"nll": mean, "perplexity": perplexity, "scored_tokens": count} def stream(self, prompt: str | list[int], *, max_new_tokens: int = 128, temperature: float = 0.0, top_p: float = 1.0, top_k: int = 0, seed: int = 0, stop: list[str] | None = None, cancel: threading.Event | None = None, chunk_size: int = 16, spec_k: int = 0) -> Iterator[GenerationChunk]: if type(chunk_size) is not int or not 1 <= chunk_size <= 1024: raise ValueError("chunk_size must be an integer in [1, 1024]") if type(spec_k) is not int or spec_k != 0 and not 2 <= spec_k <= 16: raise ValueError("spec_k must be 0 (off) or an integer in [2, 16]") if spec_k >= 2 and temperature != 0: # Greedy verification is exact only for temperature=0; sampled # speculation would need rejection sampling and is not implemented. raise ValueError("spec_k requires temperature=0 (greedy) in this version") if type(max_new_tokens) is not int or max_new_tokens < 0: raise ValueError("max_new_tokens must be a nonnegative integer") if type(seed) is not int or not 0 <= seed < 2**64: raise ValueError("seed must be an unsigned 64-bit integer") if type(top_k) is not int or top_k < 0: raise ValueError("top_k must be a nonnegative integer") for name, value in (("temperature", temperature), ("top_p", top_p)): if type(value) not in (float, int) or not math.isfinite(value): raise ValueError(f"{name} must be finite") if temperature < 0 or not 0 < top_p <= 1: raise ValueError("temperature must be nonnegative and top_p in (0, 1]") if stop is not None and (not isinstance(stop, list) or any(not isinstance(s, str) or not s for s in stop)): raise ValueError("stop must be a list of nonempty strings") stops = list(stop or []) tokens = self.encode(prompt) if len(tokens) > self.context_length: raise ValueError("Prompt length exceeds the model context length") # An over-budget generation fills the remaining context window instead # of failing; only a prompt that alone exceeds the context is invalid. max_new_tokens = min(max_new_tokens, self.context_length - len(tokens)) cancel = cancel if cancel is not None else threading.Event() while not self._generation_lock.acquire(timeout=0.05): if cancel.is_set(): return session = None monitor_done = threading.Event() monitor = None monitor_started = False released = False def cleanup() -> None: nonlocal released if released: return try: monitor_done.set() if session is not None: session.cancel() if monitor_started: monitor.join() finally: released = True self._generation_lock.release() try: if cancel.is_set(): return session = self._native.create_session(tokens, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, eos_token_ids=self._eos) def watch_cancel() -> None: # Native prefill releases the GIL, allowing cancellation even # before the first token is produced. No per-token callback. while not monitor_done.wait(0.01): if cancel.is_set(): session.cancel() return monitor = threading.Thread(target=watch_cancel, name="cism-cancel", daemon=True) monitor.start() monitor_started = True generated: list[int] = [] emitted = "" first = True total_proposed, total_accepted = 0, 0 while True: if cancel.is_set(): session.cancel() return # Stop strings require token-at-a-time checks, so sampling never # advances beyond the stopping token. Otherwise amortize bindings. batch = session.next_tokens(1 if first or stops else chunk_size, spec_k=spec_k if not first and not stops else 0) # Duck-typed sessions (tests, third-party backends) predate spec # tallies and report zeros; native sessions provide real counts. proposed, accepted = getattr(session, "spec_stats", (0, 0)) chunk_proposed, chunk_accepted = proposed - total_proposed, accepted - total_accepted total_proposed, total_accepted = proposed, accepted first = False generated.extend(batch) reason = session.finish_reason or None if reason == "cancelled": return text = self.tokenizer.decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False) if reason is None: text = text.rstrip("\ufffd") matches = [text.find(s) for s in stops if s in text] if matches: text = text[:min(matches)] reason = "stop" session.cancel() elif reason is None: hold = 0 for sequence in stops: for size in range(1, min(len(sequence), len(text) + 1)): if text.endswith(sequence[:size]): hold = max(hold, size) if hold: text = text[:-hold] if not text.startswith(emitted): raise ValueError("Tokenizer decoding rewrote emitted text; incremental streaming is unsupported") delta = text[len(emitted):] emitted = text # A client may stop iterating as soon as it sees finish_reason. # Release resources before publishing that terminal event. if reason is not None: cleanup() yield GenerationChunk(delta, batch, reason, len(tokens), len(generated), chunk_proposed, chunk_accepted) if reason is not None: break if not batch: raise RuntimeError("Native decoder made no progress") finally: cleanup() def generate(self, prompt: str | list[int], **kwargs: Any) -> GenerationResult: text: list[str] = [] tokens: list[int] = [] proposed, accepted = 0, 0 last = None for chunk in self.stream(prompt, **kwargs): text.append(chunk.text) tokens.extend(chunk.token_ids) proposed += chunk.spec_proposed accepted += chunk.spec_accepted last = chunk if last is None or last.finish_reason is None: raise RuntimeError("Generation was cancelled") return GenerationResult("".join(text), tokens, last.finish_reason, last.prompt_tokens, last.completion_tokens, proposed, accepted)