drizzymedia/SynapseCore-12B

1"""Tunable generation parameters shared by the chat API and the UI.2 3This module is the single source of truth for the generation parameters that4the frontend exposes as sliders/selects. It deliberately avoids heavy imports5(torch, transformers, gradio) so it stays cheap to import from tests and, later,6can be used to drive the frontend controls without pulling in the model.7 8The chat endpoint is a public API: clients can call it directly, bypassing the9UI controls, so the ranges here are enforced server-side rather than trusted10from the request.11"""12 13import json14from dataclasses import dataclass15 16 17@dataclass(frozen=True)18class ParamSpec:19    """Allowed range, default, and UI presentation for one generation parameter.20 21    The numeric range is what the chat API enforces; ``step`` and ``choices``22    describe how the frontend renders the control. ``validate_params`` only uses23    ``minimum``/``maximum``/``is_int``; the rest is consumed by ``ui_config``.24 25    Args:26        minimum: Smallest accepted value (inclusive).27        maximum: Largest accepted value (inclusive).28        default: Value used when the client omits the parameter.29        is_int: Whether the value is coerced to ``int`` (otherwise ``float``).30        step: Slider step for the UI control, or ``None`` for a select.31        choices: Discrete options shown in the UI as a select. The API still32            accepts any value within ``[minimum, maximum]``; these are a UI33            convenience, not an extra constraint.34    """35 36    minimum: float37    maximum: float38    default: float39    is_int: bool = False40    step: float | None = None41    choices: tuple[int, ...] | None = None42 43 44# Single source of truth for the tunable parameters. Both the chat API45# (validate_params) and the UI controls (ui_config, injected into the page)46# derive from this, so the ranges/defaults are defined in exactly one place.47PARAM_SPECS: dict[str, ParamSpec] = {48    "max_new_tokens": ParamSpec(minimum=100, maximum=4000, default=2000, is_int=True, step=10),49    "image_token_budget": ParamSpec(50        minimum=70, maximum=1120, default=280, is_int=True, choices=(70, 140, 280, 560, 1120)51    ),52    "temperature": ParamSpec(minimum=0.0, maximum=2.0, default=1.0, step=0.1),53    "top_p": ParamSpec(minimum=0.0, maximum=1.0, default=0.95, step=0.05),54    "top_k": ParamSpec(minimum=0, maximum=100, default=64, is_int=True, step=1),55    "repetition_penalty": ParamSpec(minimum=1.0, maximum=2.0, default=1.0, step=0.05),56}57 58 59def ui_config() -> dict[str, dict]:60    """Serialize the param specs for the frontend controls.61 62    Returns:63        A JSON-serializable dict, keyed by parameter name, describing each64        control's range, step, default, and discrete choices. The page injects65        this as ``window.PARAM_CONFIG`` so the UI and the API share one66        definition of the parameters.67    """68    return {69        name: {70            "min": spec.minimum,71            "max": spec.maximum,72            "step": spec.step,73            "default": spec.default,74            "choices": list(spec.choices) if spec.choices is not None else None,75        }76        for name, spec in PARAM_SPECS.items()77    }78 79 80# The page applies window.PARAM_CONFIG to its controls; this tag marks where the81# config script is inserted (just before the frontend module loads).82_APP_SCRIPT_TAG = '<script type="module" src="./app.js"></script>'83 84 85def inject_param_config(html: str) -> str:86    """Insert ``window.PARAM_CONFIG`` into the page before the app script.87 88    Lets the served page and the UI-preview stub share one definition of the89    parameters with the chat API. The injected content is fully controlled90    (numbers and parameter names), so no extra escaping is needed.91 92    Args:93        html: The page source containing the app script tag.94 95    Returns:96        The page with a ``<script>`` defining ``window.PARAM_CONFIG`` inserted97        just before the app script.98    """99    blob = f"<script>window.PARAM_CONFIG = {json.dumps(ui_config())};</script>\n    "100    return html.replace(_APP_SCRIPT_TAG, blob + _APP_SCRIPT_TAG)101 102 103def validate_params(values: dict[str, float]) -> dict[str, float]:104    """Coerce and range-check generation parameters against ``PARAM_SPECS``.105 106    Args:107        values: Raw parameter values keyed by name (one per entry in108            ``PARAM_SPECS``).109 110    Returns:111        A new dict with each value coerced to its declared numeric type.112 113    Raises:114        ValueError: If a value is non-numeric or outside its allowed range. The115            message names the parameter and its range so the API is116            self-documenting.117    """118    validated: dict[str, float] = {}119    for name, spec in PARAM_SPECS.items():120        raw = values[name]121        try:122            coerced = int(raw) if spec.is_int else float(raw)123        except (TypeError, ValueError):124            expected = "an integer" if spec.is_int else "a number"125            msg = f"{name} must be {expected}, got {raw!r}"126            raise ValueError(msg) from None127        if not spec.minimum <= coerced <= spec.maximum:128            lo = int(spec.minimum) if spec.is_int else spec.minimum129            hi = int(spec.maximum) if spec.is_int else spec.maximum130            msg = f"{name} must be between {lo} and {hi}, got {coerced}"131            raise ValueError(msg)132        validated[name] = coerced133    return validated134