"""
Frox AI Morph 1.1 — Inference Engine
The single-process serving engine used for Colab / edge / local dev.
(Production multi-user serving uses vLLM — see backend architecture doc.
 This engine still matters: it's what actually runs training-time eval,
 the demo CLI, and any deployment without a GPU cluster.)

Improvements over Morph 1.0:
  - Streaming is now the core primitive (`generate_stream`); `generate()`
    is a thin wrapper that consumes the stream. 1.0 had two separate,
    diverging implementations of the sampling logic — one for `generate()`
    and a different, buggier one for `stream()` (it re-decoded the full
    sequence every single token via `tokenizer.decode(generated_ids)`
    instead of yielding deltas, and it never applied `top_p`/`top_k`/
    `repetition_penalty`, only temperature). Fixed here.
  - Session-based KV cache: `chat()` keeps a MorphSessionCache alive
    across calls so a multi-turn conversation doesn't re-run the full
    prefill on every turn — only the new user message is processed.
  - Speculative decoding: optional small draft model proposes several
    tokens, the main model verifies them in a single forward pass.
  - Quantized loading now actually quantizes (see inference/quantize),
    instead of the 1.0 stub that logged a message and loaded fp16 anyway.
  - Every generation call reports tokens/sec, matching what the training
    loop already reports, so speed regressions are visible immediately.
"""
from __future__ import annotations

import gc
import time
from pathlib import Path
from typing import Dict, Generator, List, Optional, Tuple, Union

import torch
import torch.nn.functional as F

from config.model_config import MorphConfig
from multimodal.fusion.morph_multimodal import MorphMultimodalModel
from tokenizer.morph_tokenizer import (
    apply_chat_template, build_morph_tokenizer,
    format_thinking,
)
from inference.cache.kv_cache import MorphKVCache, MorphSessionCache
from inference.quantize.quantize import quantize_4bit, quantize_8bit, print_quantization_report
from multimodal.fusion.generation_pipeline import detect_vram_state, VRAMState, free_memory


class MorphInferenceEngine:
    """
    Frox Morph 1.1 inference engine.

    Usage:
        engine = MorphInferenceEngine.from_pretrained("./frox-morph-1-1")
        response = engine.generate([{"role": "user", "content": "Hello!"}])

        # Streaming
        for chunk in engine.generate_stream(messages):
            print(chunk, end="", flush=True)

        # Multi-turn with persistent KV cache (no re-prefill per turn)
        engine.chat("session-123", "What's the capital of France?")
        engine.chat("session-123", "What's its population?")   # reuses cache
    """

    def __init__(
        self,
        model: MorphMultimodalModel,
        tokenizer,
        config: MorphConfig,
        device: Optional[torch.device] = None,
        dtype: torch.dtype = torch.float16,
    ):
        self.config = config
        self.tokenizer = tokenizer
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = dtype

        self.model = model.to(self.device, dtype=self.dtype)
        self.model.eval()

        self.vram_state = detect_vram_state()

        # Session cache for multi-turn conversations without re-prefill
        self.sessions = MorphSessionCache(
            num_layers=config.text.num_hidden_layers,
            num_kv_heads=config.text.num_key_value_heads,
            head_dim=config.text.head_dim,
            dtype=self.dtype,
            device=self.device,
            max_seq_len_per_session=config.text.max_position_embeddings,
        )

        # Optional speculative decoding draft model
        self.draft_model: Optional[torch.nn.Module] = None
        if config.inference.use_speculative and config.inference.draft_model_path:
            self._load_draft_model(config.inference.draft_model_path)

        print(f"✓ Morph 1.1 Inference Engine ready")
        print(f"  Device: {self.device} | dtype: {self.dtype} | VRAM tier: {self.vram_state}")
        params = model.param_count()
        print(f"  Params: {params['total_billions']}B "
              f"(LM: {params['lm_billions']}B, Vision: {params['vision_billions']}B)")

    # ── Construction ──────────────────────────────────────────────

    @classmethod
    def from_pretrained(
        cls,
        path: str,
        device: Optional[str] = None,
        dtype: str = "float16",
        quantization: Optional[str] = None,   # None | "4bit" | "8bit"
    ) -> "MorphInferenceEngine":
        p = Path(path)
        dev = torch.device(device) if device else torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        torch_dtype = getattr(torch, dtype)

        print(f"📂 Loading Morph 1.1 from {path}...")
        model = MorphMultimodalModel.from_saved(str(p), device="cpu")  # load to CPU first

        model = cls._apply_quantization(model, quantization, torch_dtype)

        tokenizer_path = p / "tokenizer"
        if tokenizer_path.exists():
            tokenizer = build_morph_tokenizer(tokenizer_path=str(tokenizer_path))
        else:
            print("  ⚠ No saved tokenizer found — building a fresh one (vocab won't match!)")
            tokenizer = build_morph_tokenizer()

        return cls(model=model, tokenizer=tokenizer, config=model.config,
                   device=dev, dtype=torch_dtype)

    @staticmethod
    def _apply_quantization(model, quantization: Optional[str], compute_dtype: torch.dtype):
        if quantization == "4bit":
            model = quantize_4bit(model, compute_dtype=compute_dtype)
        elif quantization == "8bit":
            model = quantize_8bit(model)
        print_quantization_report(model, dtype_label=quantization or "float16")
        return model

    def _load_draft_model(self, path: str):
        """Load a small draft model for speculative decoding."""
        try:
            from model.architecture.morph_model import MorphForCausalLM
            self.draft_model = MorphForCausalLM.from_saved(path, device=str(self.device))
            self.draft_model.to(self.device, dtype=self.dtype).eval()
            print(f"  ✓ Draft model loaded for speculative decoding: {path}")
        except Exception as e:
            print(f"  ⚠ Draft model load failed ({e}) — speculative decoding disabled")
            self.draft_model = None

    def _free_vram(self):
        free_memory()

    # ── Core streaming generation ─────────────────────────────────

    @torch.no_grad()
    def generate_stream(
        self,
        messages: List[Dict[str, str]],
        system_prompt: Optional[str] = None,
        max_new_tokens: int = 1024,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 50,
        repetition_penalty: float = 1.1,
        do_sample: bool = True,
        pixel_values: Optional[torch.Tensor] = None,
    ) -> Generator[str, None, None]:
        """
        Stream response text incrementally (yields new text deltas, not
        the full accumulated string — fixes the 1.0 bug where every
        yielded chunk re-decoded from scratch).
        """
        prompt = apply_chat_template(
            messages, self.tokenizer, add_generation_prompt=True,
            system_prompt=system_prompt,
        )
        input_ids = self.tokenizer.encode(
            prompt, return_tensors="pt", add_special_tokens=False,
        ).to(self.device)

        if pixel_values is not None:
            pixel_values = pixel_values.to(self.device, dtype=self.dtype)

        past_key_values = None
        current_ids = input_ids
        generated_ids: List[int] = []
        prev_text = ""
        t0 = time.perf_counter()

        for step in range(max_new_tokens):
            with torch.amp.autocast("cuda", dtype=self.dtype, enabled=self.device.type == "cuda"):
                forward_kwargs = dict(
                    input_ids=current_ids,
                    past_key_values=past_key_values,
                    use_cache=True,
                )
                if step == 0 and pixel_values is not None:
                    forward_kwargs["pixel_values"] = pixel_values
                outputs = self.model(**forward_kwargs)

            logits = outputs.logits[:, -1, :]
            past_key_values = outputs.past_key_values

            next_token = self._sample(
                logits, generated_ids, temperature, top_p, top_k,
                repetition_penalty, do_sample,
            )
            token_id = next_token.item()
            generated_ids.append(token_id)
            current_ids = next_token

            full_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
            delta = full_text[len(prev_text):]
            prev_text = full_text
            if delta:
                yield delta

            if token_id == self.tokenizer.eos_token_id:
                break

        elapsed = time.perf_counter() - t0
        tok_s = len(generated_ids) / max(elapsed, 1e-6)
        print(f"  [{len(generated_ids)} tokens in {elapsed:.2f}s, {tok_s:.1f} tok/s]")

    def _sample(
        self,
        logits: torch.Tensor,
        generated_ids: List[int],
        temperature: float,
        top_p: float,
        top_k: int,
        repetition_penalty: float,
        do_sample: bool,
    ) -> torch.Tensor:
        """Shared sampling logic — used by both generate_stream and speculative decoding."""
        logits = logits.clone()

        if repetition_penalty != 1.0 and generated_ids:
            for tid in set(generated_ids):
                if logits[0, tid] < 0:
                    logits[0, tid] *= repetition_penalty
                else:
                    logits[0, tid] /= repetition_penalty

        if not do_sample:
            return logits.argmax(dim=-1, keepdim=True)

        if temperature != 1.0:
            logits = logits / max(temperature, 1e-5)

        if top_k > 0:
            top_k_vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < top_k_vals[:, -1:]] = float("-inf")

        if top_p < 1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            remove = cum_probs - F.softmax(sorted_logits, dim=-1) > top_p
            sorted_logits[remove] = float("-inf")
            logits = torch.zeros_like(logits).scatter_(1, sorted_idx, sorted_logits)

        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)

    # ── Non-streaming convenience wrapper ─────────────────────────

    def generate(
        self,
        messages: List[Dict[str, str]],
        system_prompt: Optional[str] = None,
        max_new_tokens: int = 1024,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 50,
        repetition_penalty: float = 1.1,
        do_sample: bool = True,
        pixel_values: Optional[torch.Tensor] = None,
    ) -> str:
        """Generate a full response (collects the stream)."""
        chunks = list(self.generate_stream(
            messages, system_prompt, max_new_tokens, temperature, top_p,
            top_k, repetition_penalty, do_sample, pixel_values,
        ))
        return "".join(chunks).strip()

    # ── Session-based chat (persistent KV cache across turns) ────

    @torch.no_grad()
    def chat(
        self,
        session_id: str,
        user_message: str,
        system_prompt: Optional[str] = None,
        max_new_tokens: int = 1024,
        temperature: float = 0.7,
        top_p: float = 0.9,
        stream_callback=None,
    ) -> str:
        """
        Multi-turn chat with a persistent per-session KV cache.
        Only the NEW user message is tokenized and prefilled on each
        call — prior turns are already resident in the session cache,
        so a 20-turn conversation's 5th reply doesn't re-process turns
        1-4 from scratch the way a stateless `generate()` call would.

        Bridging note: the model's attention layers internally
        concatenate past+new KV via `torch.cat` and return the full
        tensor each call. We store only the newly-computed slice back
        into the paged session cache (the "delta"), since the cache
        already holds everything before this turn.
        """
        cache = self.sessions.get_or_create(session_id)
        is_first_turn = cache.seq_len == 0

        if is_first_turn:
            prompt = apply_chat_template(
                [{"role": "user", "content": user_message}],
                self.tokenizer, add_generation_prompt=True,
                system_prompt=system_prompt,
            )
        else:
            # Only encode the new turn — the system prompt + prior turns
            # are already baked into the cached KV states.
            prompt = (
                f"<|start_header_id|>user<|end_header_id|>\n{user_message}<|eot_id|>"
                f"<|start_header_id|>assistant<|end_header_id|>\n"
            )

        new_ids = self.tokenizer.encode(prompt, return_tensors="pt",
                                        add_special_tokens=False).to(self.device)

        # Build past_key_values list from the session cache. On a brand new
        # session's very first step there's nothing cached yet, so we pass
        # None (matches a fresh forward pass); every step after that reuses
        # whatever the previous step returned.
        has_prior_context = cache.seq_len > 0
        if has_prior_context:
            past_kv = [cache.get(i) for i in range(self.config.text.num_hidden_layers)]
        else:
            past_kv = None

        generated_ids: List[int] = []
        prev_text = ""
        current_ids = new_ids
        t0 = time.perf_counter()

        for step in range(max_new_tokens):
            with torch.amp.autocast("cuda", dtype=self.dtype, enabled=self.device.type == "cuda"):
                outputs = self.model(
                    input_ids=current_ids,
                    past_key_values=past_kv,
                    use_cache=True,
                )

            logits = outputs.logits[:, -1, :]
            new_past_kv = outputs.past_key_values

            # Persist only the delta (newly computed tokens) into the paged cache.
            # IMPORTANT: cache.update() writes at offset `cache.seq_len` but does
            # NOT advance it — that's step()'s job. All layers must be written
            # at the SAME offset (they process the same tokens in lockstep), so
            # we call step() exactly once, after every layer has been updated.
            delta_len = current_ids.shape[1]
            for layer_idx, (k_full, v_full) in enumerate(new_past_kv):
                k_new = k_full[:, :, -delta_len:, :]
                v_new = v_full[:, :, -delta_len:, :]
                cache.update(layer_idx, k_new, v_new)
            cache.step(delta_len)

            past_kv = new_past_kv

            next_token = self._sample(
                logits, generated_ids, temperature, top_p, top_k=50,
                repetition_penalty=1.1, do_sample=True,
            )
            token_id = int(next_token.item())
            generated_ids.append(token_id)
            current_ids = next_token

            if stream_callback is not None:
                full_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
                delta = full_text[len(prev_text):]
                prev_text = full_text
                if delta:
                    stream_callback(delta)

            if token_id == self.tokenizer.eos_token_id:
                break

        elapsed = time.perf_counter() - t0
        response = self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
        print(f"  [session={session_id[:8]}... | {len(generated_ids)} tok in {elapsed:.2f}s | "
              f"cache={cache.seq_len} tok, {cache.memory_mb():.1f}MB]")
        return response

    def reset_session(self, session_id: str):
        self.sessions.reset_session(session_id)

    def end_session(self, session_id: str):
        self.sessions.delete_session(session_id)

    # ── Speculative decoding ───────────────────────────────────────

    @torch.no_grad()
    def generate_speculative(
        self,
        messages: List[Dict[str, str]],
        system_prompt: Optional[str] = None,
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        k: Optional[int] = None,
    ) -> Tuple[str, Dict]:
        """
        Speculative decoding (Leviathan et al. / Chen et al. algorithm).
        The small draft model proposes `k` tokens greedily; the main
        model verifies all `k` in a single forward pass and accepts a
        prefix via rejection sampling, guaranteeing the same output
        distribution as sampling from the main model alone — just
        fewer expensive forward passes through it.

        Falls back to standard `generate()` if no draft model is loaded.
        """
        if self.draft_model is None:
            text = self.generate(messages, system_prompt, max_new_tokens, temperature)
            return text, {"speculative": False, "acceptance_rate": None}

        k = k or self.config.inference.speculative_k
        prompt = apply_chat_template(messages, self.tokenizer, add_generation_prompt=True,
                                     system_prompt=system_prompt)
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt",
                                          add_special_tokens=False).to(self.device)

        generated: List[int] = []
        total_proposed, total_accepted = 0, 0
        t0 = time.perf_counter()

        current_ids = input_ids
        main_past, draft_past = None, None

        while len(generated) < max_new_tokens:
            # 1. Draft model proposes k tokens
            draft_tokens = []
            draft_ids = current_ids
            dpast = draft_past
            for _ in range(k):
                out = self.draft_model(input_ids=draft_ids, past_key_values=dpast, use_cache=True)
                logits = out.logits[:, -1, :] / max(temperature, 1e-5)
                probs = F.softmax(logits, dim=-1)
                tok = torch.multinomial(probs, 1)
                draft_tokens.append((tok, probs))
                draft_ids = tok
                dpast = out.past_key_values

            proposed_ids = torch.cat([current_ids] + [t for t, _ in draft_tokens], dim=1)

            # Length of the valid cache *before* this round (needed below to
            # truncate away any rejected draft tokens' KV entries).
            prior_cache_len = main_past[0][0].shape[2] if main_past is not None else 0

            # 2. Main model verifies all k+1 positions in one forward pass
            main_out = self.model(input_ids=proposed_ids, past_key_values=main_past, use_cache=True)
            main_logits = main_out.logits[:, -(k + 1):, :] / max(temperature, 1e-5)
            main_probs = F.softmax(main_logits, dim=-1)

            # 3. Accept/reject each proposed token (standard spec-decoding test)
            accepted = 0
            for i, (draft_tok, draft_prob) in enumerate(draft_tokens):
                # BUGFIX: this used to live at the end of the loop body, after
                # the reject branch's `break` — so a rejected token never got
                # counted, and total_proposed only ever counted acceptances.
                # acceptance_rate (total_accepted/total_proposed) was
                # therefore always ~100% regardless of real performance.
                total_proposed += 1

                tok_id = draft_tok.item()
                p_main = main_probs[0, i, tok_id].item()
                p_draft = draft_prob[0, tok_id].item()
                accept_prob = min(1.0, p_main / max(p_draft, 1e-10))

                if torch.rand(1).item() < accept_prob:
                    generated.append(tok_id)
                    accepted += 1
                    total_accepted += 1
                else:
                    # Reject — resample from the residual distribution
                    residual = (main_probs[0, i] - draft_prob[0]).clamp(min=0)
                    residual = residual / residual.sum().clamp(min=1e-10)
                    resampled = torch.multinomial(residual, 1).item()
                    generated.append(resampled)
                    break

            # If all k accepted, sample one more "bonus" token from the main model
            if accepted == k:
                bonus_probs = main_probs[0, k]
                bonus_tok = torch.multinomial(bonus_probs, 1).item()
                generated.append(bonus_tok)

            current_ids = torch.tensor([[generated[-1]]], device=self.device)

            # BUGFIX: main_out.past_key_values contains KV entries for the
            # seed token + ALL k drafted tokens, but on an early rejection
            # only `accepted` of those k draft tokens are actually part of
            # the real sequence (the rest were proposals that got thrown
            # out). Keeping the full cache left stale/phantom key-value
            # entries in place for tokens that never happened, which both
            # corrupts future attention (the model attends to keys for
            # rejected tokens) and desyncs RoPE position ids (computed from
            # cache length) from the true sequence length. Also, the
            # resampled replacement token has no cache entry yet — it gets
            # one on the next round's forward pass, same as the "bonus
            # token" case. Truncate to seed(1) + accepted confirmed drafts.
            valid_len = prior_cache_len + 1 + accepted
            main_past = tuple(
                (k_full[:, :, :valid_len, :], v_full[:, :, :valid_len, :])
                for k_full, v_full in main_out.past_key_values
            )
            draft_past = None  # simplification: draft cache rebuilt next round

            if generated and generated[-1] == self.tokenizer.eos_token_id:
                break

        elapsed = time.perf_counter() - t0
        text = self.tokenizer.decode(generated, skip_special_tokens=True).strip()
        acceptance_rate = total_accepted / max(total_proposed, 1)

        print(f"  [speculative: {len(generated)} tok in {elapsed:.2f}s | "
              f"acceptance={acceptance_rate:.1%}]")

        return text, {"speculative": True, "acceptance_rate": round(acceptance_rate, 3),
                      "tokens": len(generated), "elapsed_s": round(elapsed, 2)}

    # ── Image understanding ───────────────────────────────────────

    def understand_image(self, image, question: str, max_new_tokens: int = 512) -> str:
        """Answer a question about an image."""
        from PIL import Image as PILImage
        if isinstance(image, str):
            image = PILImage.open(image).convert("RGB")

        pixel_values = self.model.vision_module.preprocess_image(image, device=self.device)
        messages = [{"role": "user", "content": f"<|image|>\n{question}"}]
        return self.generate(messages=messages, pixel_values=pixel_values,
                             max_new_tokens=max_new_tokens)

    # ── Tool / generation-request parsing ──────────────────────────

    def parse_tool_calls(self, response: str) -> List[Dict]:
        """Parse <|tool_call|>{...}<|/tool_call|> blocks from model output."""
        import re, json
        tool_calls = []
        for match in re.findall(r"<\|tool_call\|>(.*?)<\|/tool_call\|>", response, re.DOTALL):
            try:
                tool_calls.append(json.loads(match.strip()))
            except json.JSONDecodeError:
                pass
        return tool_calls

    def parse_generation_requests(self, response: str) -> List[Dict]:
        """Parse <|gen_image|>/<|gen_video|>/<|gen_3d|> blocks from model output."""
        import re
        requests = []
        for gen_type in ("image", "video", "3d"):
            for match in re.findall(rf"<\|gen_{gen_type}\|>(.*?)<\|/gen\|>", response, re.DOTALL):
                requests.append({"type": gen_type, "prompt": match.strip()})
        return requests

    def parse_thinking(self, response: str) -> Tuple[Optional[str], str]:
        """Split <|think|>...<|/think|> reasoning from the visible answer."""
        import re
        match = re.search(r"<\|think\|>(.*?)<\|/think\|>", response, re.DOTALL)
        if not match:
            return None, response
        thinking = match.group(1).strip()
        answer = response[:match.start()] + response[match.end():]
        return thinking, answer.strip()

    # ── Batch inference ───────────────────────────────────────────

    @torch.no_grad()
    def batch_generate(
        self, prompts: List[str], max_new_tokens: int = 512, temperature: float = 0.7,
    ) -> List[str]:
        """Generate responses for multiple independent prompts at once."""
        encodings = self.tokenizer(
            prompts, return_tensors="pt", padding=True, truncation=True, max_length=4096,
        ).to(self.device)

        output_ids = self.model.language_model.generate(
            input_ids=encodings["input_ids"],
            attention_mask=encodings["attention_mask"],
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True,
            eos_token_id=self.tokenizer.eos_token_id,
            pad_token_id=self.tokenizer.pad_token_id,
        )

        responses = []
        for out in output_ids:
            new_tokens = out[encodings["input_ids"].shape[1]:]
            responses.append(self.tokenizer.decode(new_tokens, skip_special_tokens=True))
        return responses

    # ── Embeddings (for memory / RAG tools) ────────────────────────

    @torch.no_grad()
    def embed(self, text: str, max_length: int = 512) -> List[float]:
        """
        Produce an embedding vector using Morph's own hidden states —
        no external embedding API. Mean-pools the last transformer
        layer's hidden states over non-padding tokens, then L2-normalizes
        so downstream cosine-similarity search is a plain dot product.

        This is deliberately simple (no dedicated embedding head/training
        objective) rather than a from-scratch contrastive embedding
        model — good enough for the in-repo memory/RAG tools, and the
        production backend architecture's Qdrant-based RAG pipeline can
        swap in a dedicated embedding model later without changing the
        tool interface.
        """
        tokens = self.tokenizer(
            text, return_tensors="pt", truncation=True, max_length=max_length,
        ).to(self.device)

        with torch.amp.autocast("cuda", dtype=self.dtype, enabled=self.device.type == "cuda"):
            outputs = self.model.language_model(
                input_ids=tokens["input_ids"],
                attention_mask=tokens.get("attention_mask"),
                output_hidden_states=True,
                use_cache=False,
            )

        last_hidden = outputs.hidden_states[-1].float()          # [1, S, H]
        mask = tokens.get("attention_mask")
        if mask is not None:
            mask = mask.unsqueeze(-1).float()                    # [1, S, 1]
            pooled = (last_hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-6)
        else:
            pooled = last_hidden.mean(dim=1)

        pooled = torch.nn.functional.normalize(pooled, p=2, dim=-1)
        return pooled.squeeze(0).cpu().tolist()

    def embed_batch(self, texts: List[str], max_length: int = 512) -> List[List[float]]:
        """Convenience loop over embed() — fine for tool-scale batches (docs/memories, not bulk indexing)."""
        return [self.embed(t, max_length=max_length) for t in texts]

    # ── Diagnostics ──────────────────────────────────────────────

    def stats(self) -> Dict:
        return {
            "device": str(self.device),
            "dtype": str(self.dtype),
            "vram_tier": self.vram_state,
            "vram_free_gb": (torch.cuda.mem_get_info()[0] / 1e9) if torch.cuda.is_available() else None,
            "sessions": self.sessions.stats(),
            "speculative_decoding": self.draft_model is not None,
            "params": self.model.param_count(),
        }