from __future__ import annotations import os import shlex from typing import Any import chainlit as cl from hydra.core.global_hydra import GlobalHydra from omegaconf import DictConfig from mini_transformer.inference import run_inference from mini_transformer.model_loader import ( MODEL_NAME_ENV, compose_model_config, ensure_models_root, list_model_names, ) JOB_NAME = "chainlit_app" MODEL_SESSION_KEY = "selected_model" GENERATION_SESSION_KEY = "generation_overrides" HAS_ACTION_PROMPT = hasattr(cl, "AskActionMessage") and hasattr(cl, "Action") def _discover_action_fields() -> set[str]: if not HAS_ACTION_PROMPT: return set() action_cls = cl.Action for attr in ("model_fields", "__fields__"): fields = getattr(action_cls, attr, None) if isinstance(fields, dict): return set(fields.keys()) return set() ACTION_FIELDS = _discover_action_fields() GENERATION_FIELD_ORDER = [ "max_new_tokens", "temperature", "top_k", "top_p", "do_sample", "presence_penalty", "frequency_penalty", "no_repeat_ngram", "min_steps_before_eos", ] OPTIONAL_KEYS = {"top_k", "top_p", "no_repeat_ngram"} BOOL_KEYS = {"do_sample"} def _parse_bool(value: str) -> bool: candidates = {"true", "1", "yes", "y", "on"} anti_candidates = {"false", "0", "no", "n", "off"} lower = value.lower() if lower in candidates: return True if lower in anti_candidates: return False raise ValueError("expected a boolean (true/false)") def _parse_optional(value: str) -> str | None: lower = value.lower() if lower in {"none", "null", "default"}: return None return value def _parse_generation_value(key: str, raw: str) -> Any: parsed = _parse_optional(raw) if key in OPTIONAL_KEYS else raw if parsed is None: return None if key in BOOL_KEYS: return _parse_bool(parsed) if key in {"max_new_tokens", "top_k", "no_repeat_ngram", "min_steps_before_eos"}: return int(parsed) return float(parsed) def _get_generation_overrides() -> dict[str, Any]: overrides: dict[str, Any] | None = cl.user_session.get(GENERATION_SESSION_KEY) return dict(overrides) if overrides else {} def _set_generation_overrides(overrides: dict[str, Any]) -> None: cl.user_session.set(GENERATION_SESSION_KEY, overrides) def _format_generation_overrides(overrides: dict[str, Any]) -> str: if not overrides: return " (none)" lines: list[str] = [] for key in GENERATION_FIELD_ORDER: if key not in overrides: continue value = overrides[key] if isinstance(value, bool): display = "true" if value else "false" else: display = "None" if value is None else str(value) lines.append(f"- {key} = {display}") return "\n".join(lines) if lines else " (none)" def _render_generation_help(overrides: dict[str, Any]) -> str: allowed = ", ".join(GENERATION_FIELD_ORDER) summary = _format_generation_overrides(overrides) return ( "Current generation overrides:\n" f"{summary}\n\n" "Usage: /config key=value [key=value ...]\n" "Set optional fields to 'none' to restore default values.\n" "Run `/config reset` to clear all overrides.\n" f"Allowed keys: {allowed}" ) class ModelOption: """Lightweight container describing a trained model configuration.""" def __init__(self, name: str) -> None: self.name = name @property def job_name(self) -> str: return f"{JOB_NAME}_{self.name}" def discover_model_options() -> list[ModelOption]: ensure_models_root() return [ModelOption(name) for name in list_model_names()] def restore_model_from_session() -> ModelOption | None: data: dict[str, str] | None = cl.user_session.get(MODEL_SESSION_KEY) if not data: env_model = os.environ.get(MODEL_NAME_ENV) if env_model and env_model in list_model_names(): return ModelOption(name=env_model) return None name = data.get("name") if not name: return None if name not in list_model_names(): return None return ModelOption(name=name) def persist_model_selection(option: ModelOption) -> None: cl.user_session.set(MODEL_SESSION_KEY, {"name": option.name}) def extract_action_value(result: Any) -> str | None: if result is None: return None if isinstance(result, str): return result if isinstance(result, dict): payload = result.get("payload") if isinstance(payload, dict): selected = payload.get("model") or payload.get("value") or payload.get("name") if selected: return str(selected) for key in ("value", "name", "action"): value = result.get(key) if value: return str(value) return None value = getattr(result, "value", None) if value: return str(value) name = getattr(result, "name", None) if name: return str(name) payload = getattr(result, "payload", None) if isinstance(payload, dict): selected = payload.get("model") or payload.get("value") or payload.get("name") if selected: return str(selected) return None async def prompt_model_selection(header: str | None = None) -> ModelOption | None: options = discover_model_options() if not options: await cl.Message( content=( "No trained models found in `trained_models/`. " "Place a model folder with configs/config_inference.yaml there and restart the chat." ) ).send() return None if not HAS_ACTION_PROMPT: chosen = options[0] await cl.Message( content=f"Defaulting to `{chosen.name}`. Action prompts are unavailable." ).send() return chosen fields = ACTION_FIELDS or {"name", "value", "label", "payload"} actions = [] for option in options: kwargs: dict[str, Any] = {"name": option.name} if "label" in fields: kwargs["label"] = option.name if "value" in fields: kwargs["value"] = option.name if not fields or "payload" in fields: kwargs["payload"] = {"model": option.name} actions.append(cl.Action(**kwargs)) message = header or "Select a trained model to use for inference:" result = await cl.AskActionMessage(content=message, actions=actions).send() selected_name = extract_action_value(result) if not selected_name: selected_name = options[0].name return next((opt for opt in options if opt.name == selected_name), options[0]) def build_config(user_input: str, model: ModelOption) -> DictConfig: if GlobalHydra.instance().is_initialized(): GlobalHydra.instance().clear() cfg = compose_model_config(model.name, job_name=model.job_name) overrides = _get_generation_overrides() if overrides: for key, value in overrides.items(): if key in cfg.generation: cfg.generation[key] = value cfg.input_text = user_input return cfg async def ensure_model_selected(header: str | None = None) -> ModelOption | None: model = restore_model_from_session() if model: return model model = await prompt_model_selection(header=header) if model: persist_model_selection(model) return model async def handle_model_command(content: str) -> None: parts = content.split(maxsplit=1) options = discover_model_options() if not options: await cl.Message( content="No trained models available to select. Add a model under `trained_models/` and try again." ).send() return if len(parts) == 1: model = await prompt_model_selection(header="Pick a new model:") else: requested = parts[1].strip() model = next((opt for opt in options if opt.name == requested), None) if not model: available = ", ".join(opt.name for opt in options) await cl.Message( content=f"Model `{requested}` not found. Available models: {available}." ).send() return if model: persist_model_selection(model) await cl.Message( content=f"Switched to `{model.name}`. Provide text to generate a continuation." ).send() async def handle_generation_command(content: str) -> None: tokens = shlex.split(content) overrides = _get_generation_overrides() if len(tokens) == 1: await cl.Message(content=_render_generation_help(overrides)).send() return first = tokens[1].lower() if first == "reset": _set_generation_overrides({}) await cl.Message( content="Generation overrides cleared. Using model defaults for all parameters." ).send() return updates: dict[str, Any] = {} errors: list[str] = [] for token in tokens[1:]: if "=" not in token: errors.append(f"Missing '=' in '{token}'. Use key=value format.") continue key, raw_value = token.split("=", 1) normalized_key = key.strip().lower() if normalized_key not in GENERATION_FIELD_ORDER: errors.append(f"Unknown parameter '{key}'.") continue try: value = _parse_generation_value(normalized_key, raw_value.strip()) except ValueError as exc: errors.append(f"Invalid value for '{key}': {exc}") continue updates[normalized_key] = value if updates: overrides.update(updates) _set_generation_overrides(overrides) parts: list[str] = [] if updates: parts.append("Updated generation overrides:\n" + _format_generation_overrides(overrides)) if errors: formatted_errors = "\n".join(f"- {msg}" for msg in errors) parts.append(f"Issues detected:\n{formatted_errors}") if not parts: parts.append(_render_generation_help(overrides)) await cl.Message(content="\n\n".join(parts)).send() @cl.on_chat_start async def on_chat_start() -> None: model = await ensure_model_selected() if not model: return await cl.Message( content=( f"Using trained model `{model.name}`. Send text for generation, or type `/model` to switch models." " Use `/config` to tweak generation settings." ) ).send() @cl.on_message async def on_message(message: cl.Message) -> None: user_input = message.content.strip() if user_input.lower().startswith("/model"): await handle_model_command(user_input) return if user_input.lower().startswith("/config"): await handle_generation_command(user_input) return if not user_input: await cl.Message(content="Please provide some text for the model.").send() return model = await ensure_model_selected(header="Select a model before requesting generations:") if not model: return cfg = build_config(user_input, model) try: outputs = run_inference(cfg) except Exception as exc: # pragma: no cover await cl.Message(content=f"Inference error: {exc}").send() return response = outputs[0] if outputs else "No output generated." await cl.Message(content=response).send()