muse_core.py

"""Pure helpers for the Muse Glimmer Space.

This module deliberately has no Torch, Transformers, Gradio, or Spaces dependency so its
behavior can be tested without downloading or loading the 30B checkpoint.
"""

from __future__ import annotations

from dataclasses import dataclass
import secrets
from typing import Any, Callable


MODEL_CONTEXT_TOKENS = 131_072
APP_INPUT_TOKEN_LIMIT = 16_384
DEFAULT_MAX_NEW_TOKENS = 512
MIN_NEW_TOKENS = 32
MAX_NEW_TOKENS = 1_024
DEFAULT_TEMPERATURE = 1.0
DEFAULT_TOP_P = 0.95
DEFAULT_TOP_K = 64
DEFAULT_REPETITION_PENALTY = 1.0
DEFAULT_SEED = 42
VALID_REASONING_STRENGTHS = ("low", "medium", "high", "xhigh")

NATIVE_GREEDY = "Native greedy · checkpoint default"
META_SAMPLING = "Meta recommended sampling"
PRESETS = (NATIVE_GREEDY, META_SAMPLING)


@dataclass(frozen=True)
class ParsedReply:
    reasoning: str
    content: str
    tool_calls: Any = None


def preset_values(name: str) -> tuple[bool, float, float, int]:
    """Return sampling controls for one of the two documented presets."""
    if name == META_SAMPLING:
        return True, DEFAULT_TEMPERATURE, DEFAULT_TOP_P, DEFAULT_TOP_K
    return False, DEFAULT_TEMPERATURE, DEFAULT_TOP_P, DEFAULT_TOP_K


def choose_seed(
    seed: int | float | None,
    randomize: bool,
    randbelow: Callable[[int], int] = secrets.randbelow,
) -> int:
    """Resolve a valid Torch seed without relying on mutable global state."""
    if randomize:
        return int(randbelow(2_147_483_648))
    if seed is None:
        return DEFAULT_SEED
    resolved = int(seed)
    if not 0 <= resolved <= 2_147_483_647:
        raise ValueError("Seed must be between 0 and 2,147,483,647.")
    return resolved


def validate_controls(
    *,
    max_new_tokens: int | float,
    temperature: float,
    top_p: float,
    top_k: int | float,
    repetition_penalty: float,
    reasoning_strength: str,
) -> None:
    tokens = int(max_new_tokens)
    if not MIN_NEW_TOKENS <= tokens <= MAX_NEW_TOKENS:
        raise ValueError(f"Max new tokens must be between {MIN_NEW_TOKENS} and {MAX_NEW_TOKENS}.")
    if not 0.05 <= float(temperature) <= 2.0:
        raise ValueError("Temperature must be between 0.05 and 2.0.")
    if not 0.05 <= float(top_p) <= 1.0:
        raise ValueError("Top-p must be between 0.05 and 1.0.")
    if not 1 <= int(top_k) <= 200:
        raise ValueError("Top-k must be between 1 and 200.")
    if not 0.8 <= float(repetition_penalty) <= 1.3:
        raise ValueError("Repetition penalty must be between 0.8 and 1.3.")
    if reasoning_strength not in VALID_REASONING_STRENGTHS:
        raise ValueError("Unsupported reasoning strength.")


def generation_kwargs(
    *,
    do_sample: bool,
    max_new_tokens: int | float,
    temperature: float,
    top_p: float,
    top_k: int | float,
    repetition_penalty: float,
) -> dict[str, Any]:
    """Build generation arguments while preserving the checkpoint's dual EOS contract."""
    kwargs: dict[str, Any] = {
        "max_new_tokens": int(max_new_tokens),
        "do_sample": bool(do_sample),
        "eos_token_id": [200_001, 200_008],
        "pad_token_id": 200_018,
        "use_cache": True,
        "repetition_penalty": float(repetition_penalty),
    }
    if do_sample:
        kwargs.update(
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
        )
    return kwargs


def estimate_gpu_duration(max_new_tokens: int | float, has_image: bool) -> int:
    """Return a bounded ZeroGPU reservation in seconds.

    This is intentionally conservative for a dense 30B BF16 model. It is a maximum reservation,
    not a claim that each call consumes the full amount.
    """
    tokens = max(MIN_NEW_TOKENS, min(MAX_NEW_TOKENS, int(max_new_tokens)))
    seconds = 55 + int(tokens * 0.13) + (25 if has_image else 0)
    return max(60, min(240, seconds))


def coerce_parsed_reply(message: dict[str, Any] | None) -> ParsedReply:
    """Normalize the fields produced by Transformers' native response parser."""
    message = message or {}

    def text(value: Any) -> str:
        if value is None:
            return ""
        if isinstance(value, str):
            return value.strip()
        if isinstance(value, list):
            chunks: list[str] = []
            for part in value:
                if isinstance(part, str):
                    chunks.append(part)
                elif isinstance(part, dict) and isinstance(part.get("text"), str):
                    chunks.append(part["text"])
            return "\n".join(chunks).strip()
        return str(value).strip()

    return ParsedReply(
        reasoning=text(message.get("reasoning_content")),
        content=text(message.get("content")),
        tool_calls=message.get("tool_calls"),
    )


def protect_reasoning_tags(text: str) -> str:
    """Prevent model-emitted tags from breaking the UI's reasoning region."""
    return text.replace("<think>", "&lt;think&gt;").replace("</think>", "&lt;/think&gt;")


def render_reply(
    reasoning: str,
    content: str,
    *,
    show_reasoning: bool,
    pending: bool = False,
    hit_token_limit: bool = False,
) -> str:
    """Render a parsed reply for Gradio's collapsible reasoning tags."""
    sections: list[str] = []
    if show_reasoning and reasoning.strip():
        sections.append(f"<think>\n{protect_reasoning_tags(reasoning.strip())}\n</think>")
    if content.strip():
        sections.append(protect_reasoning_tags(content.strip()))
    elif pending:
        sections.append("_Generating…_")
    elif hit_token_limit and reasoning.strip():
        sections.append("_The response budget ended before a final-answer region was produced._")
    elif reasoning.strip():
        sections.append("_The model ended without a separate final-answer region._")
    return "\n\n".join(sections).strip()


def friendly_error(error: BaseException) -> str:
    """Map technical failures to messages that do not echo sensitive inputs or paths."""
    name = type(error).__name__.lower()
    message = str(error).lower()
    if "outofmemory" in name or "out of memory" in message:
        return "GPU memory was exhausted. Clear older turns, remove images, or lower the response budget."
    if "timeout" in name or "timed out" in message:
        return "The ZeroGPU allocation timed out. Lower the response budget and try again."
    if isinstance(error, ValueError):
        safe = str(error).strip()
        if safe and len(safe) <= 320:
            safe = safe.replace("/models/muse-glimmer-assistant", "<assistant_mount>")
            safe = safe.replace("/models/muse-glimmer", "<full_mount>")
            if safe.count("/") > 2:
                safe = "A runtime setup error occurred while loading the selected checkpoint."
            return safe
    return "Inference failed. The private Space logs contain the technical error type."