src/cism/engine.py

"""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)