diff --git a/pyproject.toml b/pyproject.toml index 79c14c907..9f6465d3a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ classifiers = [ ] dependencies = [ "typer", + "rich", "pydantic", "pyyaml", "nest-asyncio", diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index e1047a651..08d7e17e8 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -10,6 +10,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference +from unsloth_cli.commands.chat import chat from unsloth_cli.commands.export import export, list_checkpoints from unsloth_cli.commands.studio import ( run as studio_run, @@ -72,6 +73,7 @@ def main( app.command()(train) app.command()(inference) +app.command()(chat) app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py new file mode 100644 index 000000000..5e734ee2a --- /dev/null +++ b/unsloth_cli/_inference.py @@ -0,0 +1,417 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Model loading and streaming shared by `inference` and `chat`.""" + +import os +import re +import sys +from pathlib import Path +from typing import Optional + +import typer + +_THINK_OPEN = "" +_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL) + + +def ensure_studio_backend_path() -> None: + backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend") + if backend_dir not in sys.path: + sys.path.insert(0, backend_dir) + + +def configure_quiet_logging() -> None: + import logging + + import structlog + + # The CLI never configures structlog, so without this every backend INFO + # line prints. LOG_LEVEL is exported so the worker subprocess inherits it. + level_name = os.environ.setdefault("LOG_LEVEL", "WARNING").upper() + level = getattr(logging, level_name, logging.WARNING) + structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level)) + os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + + +def visible_text(text: str, show_thinking: bool) -> str: + if show_thinking: + return text + text = _THINK_BLOCK.sub("", text) + # Hold back an unclosed trailing so reasoning never leaks mid-stream. + open_idx = text.find(_THINK_OPEN) + if open_idx != -1: + text = text[:open_idx] + max_prefix = min(len(text), len(_THINK_OPEN) - 1) + for size in range(max_prefix, 0, -1): + if _THINK_OPEN.startswith(text[-size:]): + return text[:-size] + return text + + +def stream_to_stdout(stream, show_thinking: bool) -> str: + # Backends yield the full text-so-far on each step (llama.cpp ends with a + # metadata dict, skipped); print the growing tail, return the raw text. + raw = "" + shown = "" + for chunk in stream: + if not isinstance(chunk, str): + continue + raw = chunk + rendered = visible_text(chunk, show_thinking) + delta = rendered[len(shown) :] + if delta: + sys.stdout.write(delta) + sys.stdout.flush() + shown = rendered + sys.stdout.write("\n") + sys.stdout.flush() + return raw + + +def stream_markdown(stream, show_thinking: bool, *, console) -> str: + from rich.live import Live + from rich.markdown import Markdown + from rich.text import Text + + raw = "" + with Live(console = console, refresh_per_second = 12, vertical_overflow = "visible") as live: + for chunk in stream: + if not isinstance(chunk, str): + continue + raw = chunk + visible = visible_text(chunk, show_thinking) + live.update(Markdown(visible) if visible.strip() else Text("")) + return raw + + +def collect_stream(stream, show_thinking: bool) -> str: + raw = "" + for chunk in stream: + if isinstance(chunk, str): + raw = chunk + return visible_text(raw, show_thinking) + + +def render_columns( + left_label: str, + left_text: str, + right_label: str, + right_text: str, + *, + console = None, +) -> None: + from rich import box + from rich.console import Console + from rich.table import Table + + table = Table(box = box.MINIMAL, expand = True, padding = (0, 1), pad_edge = False) + table.add_column(left_label, header_style = "bold yellow", ratio = 1, overflow = "fold") + table.add_column(right_label, header_style = "bold magenta", ratio = 1, overflow = "fold") + table.add_row(left_text or "", right_text or "") + (console or Console()).print(table) + + +class ChatBackend: + """Uniform stream()/close() over the llama-server and Unsloth backends.""" + + def __init__(self, kind: str, backend) -> None: + self._kind = kind # "gguf" | "unsloth" + self._backend = backend + + def stream( + self, + messages: list, + *, + system_prompt: str, + temperature: float, + top_p: float, + top_k: int, + max_new_tokens: int, + repetition_penalty: float, + enable_thinking: bool, + use_adapter: Optional[bool] = None, + ): + if self._kind == "gguf": + # llama-server takes the system prompt as the first message. + msgs = list(messages) + if system_prompt: + msgs = [{"role": "system", "content": system_prompt}, *msgs] + return self._backend.generate_chat_completion( + messages = msgs, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = enable_thinking, + ) + gen_kwargs = dict( + messages = messages, + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = enable_thinking, + ) + if use_adapter is not None: + return self._backend.generate_with_adapter_control( + use_adapter = use_adapter, **gen_kwargs + ) + return self._backend.generate_chat_response(**gen_kwargs) + + def close(self) -> None: + # Shut the worker down directly: the graceful unload_model waits for + # an ack that compare mode can swallow, hanging exit for minutes. + try: + if self._kind == "gguf": + self._backend.unload_model() + else: + self._backend._shutdown_subprocess(timeout = 2.0) + except Exception: + pass + + +def resolve_model_config(model: str, *, hf_token: Optional[str]): + ensure_studio_backend_path() + from utils.models import ModelConfig + + model_config = ModelConfig.from_identifier(model_id = model, hf_token = hf_token) + if not model_config: + typer.echo("Could not resolve model config", err = True) + raise typer.Exit(code = 1) + return model_config + + +def _load_gguf_backend(model_config, *, hf_token, max_seq_length): + ensure_studio_backend_path() + from core.inference.llama_cpp import LlamaCppBackend + + llama_backend = LlamaCppBackend() + common = dict( + hf_variant = model_config.gguf_variant, + model_identifier = model_config.identifier, + is_vision = model_config.is_vision, + n_ctx = max_seq_length, + ) + if model_config.gguf_hf_repo: + loaded = llama_backend.load_model( + hf_repo = model_config.gguf_hf_repo, hf_token = hf_token, **common + ) + else: + loaded = llama_backend.load_model( + gguf_path = model_config.gguf_file, + mmproj_path = model_config.gguf_mmproj_file, + mtp_draft_path = model_config.gguf_mtp_file, + **common, + ) + if not loaded: + typer.echo("Model load failed", err = True) + raise typer.Exit(code = 1) + return ChatBackend("gguf", llama_backend) + + +def load_chat_backend( + model: str, + *, + hf_token: Optional[str], + max_seq_length: int, + load_in_4bit: bool, + model_config = None, + fresh_backend: bool = False, +): + """Load `model` in-process: GGUF via llama-server, else the orchestrator. + + fresh_backend uses a private orchestrator so a second model (compare's + base column) can run alongside the main one. + """ + if model_config is None: + model_config = resolve_model_config(model, hf_token = hf_token) + + typer.echo(f"Loading {model}", err = True) + + if model_config.is_gguf: + return _load_gguf_backend(model_config, hf_token = hf_token, max_seq_length = max_seq_length) + + if fresh_backend: + ensure_studio_backend_path() + from core.inference import InferenceOrchestrator + backend = InferenceOrchestrator() + else: + ensure_studio_backend_path() + from core.inference import get_inference_backend + backend = get_inference_backend() + if not backend.load_model( + config = model_config, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + hf_token = hf_token, + ): + typer.echo("Model load failed", err = True) + raise typer.Exit(code = 1) + return ChatBackend("unsloth", backend) + + +def find_studio_server(timeout: float = 0.4) -> Optional[str]: + import urllib.request + base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") + try: + with urllib.request.urlopen(f"{base}/api/health", timeout = timeout): + return base + except Exception: + return None + + +def _studio_token() -> Optional[str]: + """Self-issue a JWT: the CLI runs as the same OS user as the server, so it + signs with the same stored secret the server validates against.""" + try: + import studio.backend.core # noqa: F401 puts studio/backend on sys.path + + from studio.backend.auth import storage + from studio.backend.auth.authentication import create_access_token + + row = storage.get_connection().execute("SELECT username FROM auth_user LIMIT 1").fetchone() + return create_access_token(row[0], desktop = True) if row else None + except Exception: + return None + + +class HttpChatBackend: + """Chat against a running Studio server over its OpenAI-compatible API. + + close() leaves the model loaded on purpose — the next session (or the + UI) starts instantly. + """ + + def __init__(self, base_url: str, token: str) -> None: + self._base = base_url + self._token = token + + def _request( + self, + method: str, + path: str, + payload = None, + timeout = None, + ): + import json + import urllib.request + + request = urllib.request.Request( + self._base + path, + data = None if payload is None else json.dumps(payload).encode(), + headers = { + "Authorization": f"Bearer {self._token}", + "Content-Type": "application/json", + }, + method = method, + ) + return urllib.request.urlopen(request, timeout = timeout) + + def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None: + typer.echo(f"Loading {model} on the Studio server", err = True) + try: + self._request( + "POST", + "/api/inference/load", + { + "model_path": model, + "hf_token": hf_token, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + }, + ).close() + except Exception as exc: + typer.echo(f"Model load failed: {exc}", err = True) + raise typer.Exit(code = 1) + + def stream( + self, + messages: list, + *, + system_prompt: str, + temperature: float, + top_p: float, + top_k: int, + max_new_tokens: int, + repetition_penalty: float, + enable_thinking: bool, + use_adapter: Optional[bool] = None, + ): + import json + + msgs = list(messages) + if system_prompt: + msgs = [{"role": "system", "content": system_prompt}, *msgs] + resp = self._request( + "POST", + "/v1/chat/completions", + { + "model": "default", + "messages": msgs, + "stream": True, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "max_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + "enable_thinking": enable_thinking, + }, + ) + + def cumulative(): + # Accumulate SSE deltas into the full-text-so-far convention the + # stream helpers expect. + text = "" + with resp: + for raw_line in resp: + line = raw_line.decode("utf-8", "replace").strip() + if not line.startswith("data:"): + continue + data = line[len("data:") :].strip() + if data == "[DONE]": + break + try: + parsed = json.loads(data) + except ValueError: + continue + if "error" in parsed: + raise RuntimeError( + f"Server error: {parsed['error'].get('message', 'Unknown server error')}" + ) + try: + delta = parsed["choices"][0]["delta"].get("content") + except (KeyError, IndexError): + continue + if not delta: + continue + text += delta + # An emoji can arrive split across two deltas as lone + # surrogate halves: hold back a trailing half, merge pairs. + visible = text + if "\ud800" <= visible[-1] <= "\udbff": + visible = visible[:-1] + yield visible.encode("utf-16", "surrogatepass").decode("utf-16", "replace") + + return cumulative() + + def close(self) -> None: + pass + + +def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit): + """Backend on a running Studio server, or None (caller loads locally).""" + base_url = find_studio_server() + if not base_url: + return None + token = _studio_token() + if not token: + return None + backend = HttpChatBackend(base_url, token) + backend.ensure_loaded( + model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit + ) + return backend diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py new file mode 100644 index 000000000..c62916bc7 --- /dev/null +++ b/unsloth_cli/commands/chat.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from typing import Optional + +import typer +from rich.console import Console + +from unsloth_cli._inference import ( + collect_stream, + configure_quiet_logging, + connect_studio_server, + ensure_studio_backend_path, + load_chat_backend, + render_columns, + resolve_model_config, + stream_markdown, + visible_text, +) + +_HELP = ( + "Commands: /exit (quit), /reset (clear history), " + "/think (toggle reasoning), /compare (base vs tuned), /help" +) + + +def _you_prompt(colors: bool) -> str: + # The prompt must go through input(), not a separate print — readline + # redraws erase anything they didn't draw, eating the label. GNU readline + # wants colors wrapped in \001/\002; libedit (macOS) prints those + # literally, so it gets raw ANSI. + try: + import readline + except ImportError: + return "\n\x1b[1;36mYou: \x1b[0m" if colors else "\nYou: " + libedit = ( + "libedit" in (readline.__doc__ or "") or getattr(readline, "backend", "") == "editline" + ) + if not colors: + return "\nYou: " + if libedit: + return "\n\x1b[1;36mYou: \x1b[0m" + return "\n\001\x1b[1;36m\002You: \001\x1b[0m\002" + + +def _compare_blocked_reason(model_config) -> Optional[str]: + if model_config.is_gguf: + return ( + "GGUF models can't toggle adapters — load a LoRA fine-tune " + "(transformers backend) to compare base vs tuned." + ) + if not model_config.is_lora: + return ( + "this isn't a LoRA adapter — compare turns the adapter off for the " + "'base' column, so there's nothing to compare against." + ) + return None + + +def _get_base_load_in_4bit(model_config) -> bool: + """Determine load_in_4bit for base model based on tuned adapter precision.""" + if not model_config.is_lora or not model_config.path: + # Fallback to default if not a LoRA or no path + return True + + try: + import json + from pathlib import Path + + adapter_cfg_path = Path(model_config.path) / "adapter_config.json" + if not adapter_cfg_path.exists(): + return True + + with open(adapter_cfg_path) as f: + adapter_cfg = json.load(f) + + training_method = adapter_cfg.get("unsloth_training_method") + if training_method == "lora": + return False + elif training_method == "qlora": + return True + elif not training_method: + # Fallback: check base model name for -bnb-4bit suffix + if model_config.base_model and "-bnb-4bit" not in model_config.base_model.lower(): + return False + return True + return True + except Exception: + return True + + +def _compare_needs_second_model() -> bool: + # MLX can't toggle the adapter off, so compare loads the base separately. + # detect_hardware() would print into the chat (and import torch), so + # probe its MLX condition quietly: Apple Silicon with mlx installed. + try: + from studio.backend.utils.hardware import hardware as hw + + if hw.DEVICE is not None: + return hw.DEVICE == hw.DeviceType.MLX + if not hw.is_apple_silicon(): + return False + import mlx.core # noqa: F401 + + return True + except Exception: + return False + + +def _pick_trained_model(console) -> str: + ensure_studio_backend_path() + from utils.models import scan_trained_models + + trained = scan_trained_models() + if not trained: + typer.echo( + "No trained models found in your outputs folder. " + "Pass a model id or path: `unsloth chat `.", + err = True, + ) + raise typer.Exit(code = 1) + + console.print("Your trained models (newest first):", style = "bold") + for i, (display_name, _, model_type) in enumerate(trained, 1): + console.print(f" {i}. {display_name} ({model_type})", markup = False) + + while True: + try: + raw = input(f"Chat with [1-{len(trained)}, Enter = 1]: ").strip() + except (EOFError, KeyboardInterrupt): + raise typer.Exit(code = 1) + if not raw: + return trained[0][1] + if raw.isdigit() and 1 <= int(raw) <= len(trained): + return trained[int(raw) - 1][1] + console.print(f"Pick a number between 1 and {len(trained)}.", style = "yellow") + + +def chat( + model: Optional[str] = typer.Argument( + None, help = "HF model id or local path. Omit to pick one of your trained models." + ), + hf_token: Optional[str] = typer.Option( + None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed." + ), + temperature: float = typer.Option(0.7, "--temperature"), + top_p: float = typer.Option(0.9, "--top-p"), + top_k: int = typer.Option(40, "--top-k"), + max_new_tokens: int = typer.Option(512, "--max-new-tokens"), + repetition_penalty: float = typer.Option(1.1, "--repetition-penalty"), + system_prompt: str = typer.Option( + "", "--system-prompt", help = "Optional system prompt for the conversation." + ), + max_seq_length: int = typer.Option(4096, "--max-seq-length"), + load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"), + think: bool = typer.Option( + False, + "--think/--no-think", + help = "Start with the model's reasoning shown. Toggle live with /think.", + ), + compare: bool = typer.Option( + False, + "--compare/--no-compare", + help = "Answer each prompt twice — base vs fine-tuned — side by side. " + "Needs a LoRA adapter. Toggle live with /compare.", + ), + verbose: bool = typer.Option( + False, "--verbose", "-v", help = "Show backend and llama-server logs." + ), + no_server: bool = typer.Option( + False, + "--no-server", + help = "Load the model in-process even if a Studio server is running.", + ), +): + """Start an interactive chat with a model (loads once, stays warm).""" + if not verbose: + configure_quiet_logging() + + console = Console() + err = Console(stderr = True) + + if model is None: + model = _pick_trained_model(console) + + # Resolve first so --compare can be rejected before the slow load. + model_config = resolve_model_config(model, hf_token = hf_token) + compare_blocked = _compare_blocked_reason(model_config) + if compare and compare_blocked: + err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False) + raise typer.Exit(code = 1) + + load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit) + + # Prefer a running Studio server: instant starts, model shared with the UI. + chat_backend = None if no_server else connect_studio_server(model, **load_opts) + server_mode = chat_backend is not None + if server_mode: + console.print( + "(Studio server connected — model stays warm after /exit)", + style = "bright_black", + ) + else: + chat_backend = load_chat_backend(model, model_config = model_config, **load_opts) + + name = model_config.display_name or model + show_thinking = think + compare_mode = compare + messages = [] + + # Compare's base column: server mode keeps the tuned model remote and + # loads the base locally; local MLX (no adapter toggle) does the same; + # local CUDA just toggles the adapter on the one loaded model. + dual_compare = compare_blocked is None and (server_mode or _compare_needs_second_model()) + base_backend = None + + def load_base_for_compare(): + nonlocal base_backend + if base_backend is not None: + return True + base_id = model_config.base_model + if not base_id: + console.print( + "(compare unavailable: this adapter doesn't record its base model)", + style = "yellow", + ) + return False + console.print( + f"(loading base model {base_id} for compare — keeps two models in memory)", + style = "bright_black", + markup = False, + ) + try: + # Use the same precision as the tuned model for fair comparison + base_load_opts = dict(load_opts) # Copy original options + base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config) + base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts) + except Exception as exc: + err.print(f"(base model load failed: {exc})", style = "red", markup = False) + return False + return True + + if compare and dual_compare and not load_base_for_compare(): + raise typer.Exit(code = 1) + + def generate(backend = None, use_adapter = None): + # Reads messages and show_thinking live, so /reset and /think apply. + return (backend or chat_backend).stream( + messages, + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = show_thinking, + use_adapter = use_adapter, + ) + + console.print() + console.print(f"Chatting with {name}", style = "bold green", markup = False) + console.print(_HELP, style = "bright_black") + + # legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage. + you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows) + assistant_label = "[bold magenta]Assistant:[/bold magenta]" + + try: + while True: + try: + user = input(you_prompt).strip() + except (EOFError, KeyboardInterrupt): + console.print() + break + + if not user: + continue + if user in ("/exit", "/quit"): + break + if user == "/reset": + messages = [] + console.print("(history cleared)", style = "bright_black") + continue + if user == "/think": + show_thinking = not show_thinking + state = "on" if show_thinking else "off" + console.print(f"(thinking {state})", style = "bright_black") + continue + if user == "/compare": + if compare_blocked: + console.print(f"(compare unavailable: {compare_blocked})", style = "yellow") + continue + if not compare_mode and dual_compare and not load_base_for_compare(): + continue + compare_mode = not compare_mode + state = "on" if compare_mode else "off" + console.print(f"(compare {state})", style = "bright_black") + continue + if user in ("/help", "/?"): + console.print(_HELP, style = "bright_black") + continue + + messages.append({"role": "user", "content": user}) + + try: + if compare_mode: + console.print("(comparing base vs tuned…)", style = "bright_black") + if dual_compare: + base_text = collect_stream(generate(backend = base_backend), show_thinking) + tuned_text = collect_stream(generate(), show_thinking) + else: + base_text = collect_stream(generate(use_adapter = False), show_thinking) + tuned_text = collect_stream(generate(use_adapter = True), show_thinking) + console.print() + render_columns( + "base", base_text, f"{name} (tuned)", tuned_text, console = console + ) + # History continues as the tuned model; base is just the reference. + answer = tuned_text + else: + console.print(assistant_label) + answer = stream_markdown(generate(), show_thinking, console = console) + except KeyboardInterrupt: + # Ctrl-C aborts this answer only; drop the unanswered turn. + console.print("\n(interrupted)", style = "bright_black") + messages.pop() + continue + except Exception as exc: + err.print(f"\n(error: {exc})", style = "red", markup = False) + messages.pop() + continue + + messages.append( + {"role": "assistant", "content": visible_text(answer, show_thinking = False)} + ) + finally: + chat_backend.close() + if base_backend is not None: + base_backend.close() + err.print("\nBye.", style = "bright_black") diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py index 18de68c4d..5dbc32c7d 100644 --- a/unsloth_cli/commands/inference.py +++ b/unsloth_cli/commands/inference.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import sys from typing import Optional import typer +from unsloth_cli._inference import ( + configure_quiet_logging, + connect_studio_server, + load_chat_backend, + stream_to_stdout, +) + def inference( model: str = typer.Argument(..., help = "HF model id or local path."), @@ -25,45 +31,46 @@ def inference( ), max_seq_length: int = typer.Option(2048, "--max-seq-length"), load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"), + think: bool = typer.Option( + False, + "--think/--no-think", + help = "Show the model's reasoning. Off by default so reasoning " + "models answer directly instead of spending the token budget thinking.", + ), + verbose: bool = typer.Option( + False, + "--verbose", + "-v", + help = "Show backend and llama-server logs (otherwise only the answer).", + ), + no_server: bool = typer.Option( + False, + "--no-server", + help = "Load the model in-process even if a Studio server is running.", + ), ): """Run a single inference using the specified model.""" - from studio.backend.core import ModelConfig, get_inference_backend + if not verbose: + configure_quiet_logging() - inference_backend = get_inference_backend() - model_config = ModelConfig.from_ui_selection( - dropdown_value = model, search_value = None, hf_token = hf_token, is_lora = False - ) - if not model_config: - typer.echo("Could not resolve model config", err = True) - raise typer.Exit(code = 1) - - if not inference_backend.load_model( - config = model_config, - max_seq_length = max_seq_length, - load_in_4bit = load_in_4bit, - hf_token = hf_token, - ): - typer.echo("Model load failed", err = True) - raise typer.Exit(code = 1) - - messages = [{"role": "user", "content": prompt}] - stream = inference_backend.generate_chat_response( - messages = messages, - system_prompt = system_prompt, - temperature = temperature, - top_p = top_p, - top_k = top_k, - max_new_tokens = max_new_tokens, - repetition_penalty = repetition_penalty, - ) - - typer.echo("Assistant:", nl = True) - previous = "" - for chunk in stream: - delta = chunk[len(previous) :] - if delta: - sys.stdout.write(delta) - sys.stdout.flush() - previous = chunk - sys.stdout.write("\n") - sys.stdout.flush() + # A running Studio server keeps the model warm between runs, which is + # exactly what a one-shot command wants. + load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit) + chat_backend = None if no_server else connect_studio_server(model, **load_opts) + if chat_backend is None: + chat_backend = load_chat_backend(model, **load_opts) + try: + stream = chat_backend.stream( + [{"role": "user", "content": prompt}], + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = think, + ) + typer.echo("Assistant:") + stream_to_stdout(stream, show_thinking = think) + finally: + chat_backend.close() diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py new file mode 100644 index 000000000..629cb4640 --- /dev/null +++ b/unsloth_cli/tests/test_inference_chat.py @@ -0,0 +1,401 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the `unsloth chat` / `unsloth inference` CLI — fakes only, no model loads.""" + +from __future__ import annotations + +import inspect +import sys +import types +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +import typer +from rich.console import Console +from typer.testing import CliRunner + +import unsloth_cli.commands.chat as chatmod +from unsloth_cli._inference import ( + ChatBackend, + HttpChatBackend, + collect_stream, + render_columns, + visible_text, +) + + +class _FakeConfig: + is_gguf = False + is_lora = True + display_name = "fake-model" + base_model = "fake/base" + path = None + + +def _chat_app(): + cli = typer.Typer() + cli.command()(chatmod.chat) + return cli + + +def test_visible_text_passthrough_when_shown(): + text = "reasoninganswer" + assert visible_text(text, show_thinking = True) == text + + +def test_visible_text_strips_closed_think_block(): + text = "step 1\nstep 2The answer is 42." + assert visible_text(text, show_thinking = False) == "The answer is 42." + + +def test_visible_text_holds_unclosed_think(): + # An open is held back so partial reasoning never leaks mid-stream. + assert visible_text("still thinking", show_thinking = False) == "" + assert visible_text("done.more thinking", show_thinking = False) == "done." + + +def test_visible_text_holds_partial_think_prefix(): + # Streams are cumulative, so the opening tag can arrive as "<", "". Hold possible tag prefixes until they are disambiguated. + assert visible_text("<", show_thinking = False) == "" + assert visible_text("rhel", "rhello"]) + assert collect_stream(stream, show_thinking = False) == "hello" + + +def test_render_columns_emits_both_answers_with_separator(capsys): + render_columns("base", "alpha", "tuned", "beta") + out = capsys.readouterr().out + assert "base" in out and "tuned" in out + assert "alpha" in out and "beta" in out + assert "│" in out + + +def test_you_prompt_matches_readline_backend(monkeypatch): + gnu = types.ModuleType("readline") + gnu.__doc__ = "Importing this module enables command line editing using GNU readline." + monkeypatch.setitem(sys.modules, "readline", gnu) + prompt = chatmod._you_prompt(colors = True) + assert "You: " in prompt and "\001" in prompt + + libedit = types.ModuleType("readline") + libedit.__doc__ = "Importing this module enables command line editing using libedit readline." + monkeypatch.setitem(sys.modules, "readline", libedit) + assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m" + assert chatmod._you_prompt(colors = False) == "\nYou: " + + # Windows: no readline module at all; the console's own line editing + # handles backspace, so plain ANSI color (no markers) is safe. + monkeypatch.setitem(sys.modules, "readline", None) + assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m" + assert chatmod._you_prompt(colors = False) == "\nYou: " + + +def test_chat_registered_on_app(): + from unsloth_cli import app + + # cmd.name is None until typer resolves it from the callback name. + names = {(cmd.name or cmd.callback.__name__) for cmd in app.registered_commands} + assert "chat" in names + + +def test_chat_exits_cleanly_on_slash_exit(monkeypatch): + closed = [] + + class _FakeChatBackend: + def stream(self, *a, **k): + return iter(["hello"]) + + def close(self): + closed.append(True) + + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None) + + runner = CliRunner() + for args in (["fake-model"], ["fake-model", "--compare"]): + closed.clear() + result = runner.invoke(_chat_app(), args, input = "hi\n/exit\n") + assert result.exit_code == 0, result.output + assert closed == [True] + assert "Bye." in result.output + # The prompt must go through input() (readline-safe), not a print. + assert "You: " in result.output + assert "You: You:" not in result.output + + +def test_pick_trained_model_lists_and_selects(monkeypatch): + fake_models = types.ModuleType("utils.models") + fake_models.scan_trained_models = lambda: [ + ("run-new", "outputs/run-new", "lora"), + ("run-old", "outputs/run-old", "merged"), + ] + monkeypatch.setitem(sys.modules, "utils.models", fake_models) + + monkeypatch.setattr("builtins.input", lambda prompt = "": "2") + assert chatmod._pick_trained_model(Console()) == "outputs/run-old" + + monkeypatch.setattr("builtins.input", lambda prompt = "": "") + assert chatmod._pick_trained_model(Console()) == "outputs/run-new" + + +def test_chat_no_arg_chats_with_picked_trained_model(monkeypatch): + class _FakeChatBackend: + def stream(self, *a, **k): + return iter(["hello"]) + + def close(self): + pass + + resolved = [] + monkeypatch.setattr(chatmod, "_pick_trained_model", lambda console: "outputs/run-42") + monkeypatch.setattr( + chatmod, + "resolve_model_config", + lambda model, **k: (resolved.append(model), _FakeConfig())[1], + ) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None) + + result = CliRunner().invoke(_chat_app(), [], input = "/exit\n") + assert result.exit_code == 0, result.output + assert resolved == ["outputs/run-42"] + + +def test_find_studio_server_none_when_not_running(monkeypatch): + import urllib.request + + from unsloth_cli import _inference + + def refuse(*a, **k): + raise OSError("connection refused") + + monkeypatch.setattr(urllib.request, "urlopen", refuse) + assert _inference.find_studio_server() is None + + +class _FakeSSEResponse: + def __init__(self, lines): + self._lines = lines + + def __iter__(self): + return iter(self._lines) + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def test_http_backend_streams_cumulative_text(monkeypatch): + backend = HttpChatBackend("http://localhost:8888", "token") + response = _FakeSSEResponse( + [ + b'data: {"choices":[{"delta":{"content":"He"}}]}\n', + b"\n", + b'data: {"choices":[{"delta":{"content":"llo"}}]}\n', + b"data: [DONE]\n", + ] + ) + monkeypatch.setattr(backend, "_request", lambda *a, **k: response) + + out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS)) + assert out == ["He", "Hello"] + + +def test_http_backend_merges_emoji_split_across_deltas(monkeypatch): + backend = HttpChatBackend("http://localhost:8888", "token") + response = _FakeSSEResponse( + [ + b'data: {"choices":[{"delta":{"content":"hi "}}]}\n', + b'data: {"choices":[{"delta":{"content":"\\ud83d"}}]}\n', + b'data: {"choices":[{"delta":{"content":"\\ude0a"}}]}\n', + b"data: [DONE]\n", + ] + ) + monkeypatch.setattr(backend, "_request", lambda *a, **k: response) + + out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS)) + # The lone high surrogate is held back, then merged with its other half. + assert out == ["hi ", "hi ", "hi 😊"] + + +def test_chat_prefers_running_studio_server(monkeypatch): + closed = [] + + class _FakeHttpBackend: + def stream(self, *a, **k): + return iter(["hello"]) + + def close(self): + closed.append("http") + + local_loads = [] + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend()) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: local_loads.append(1)) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + + result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n") + + assert result.exit_code == 0, result.output + assert local_loads == [] + assert "stays warm" in result.output + assert closed == ["http"] + + +def test_chat_server_mode_compare_loads_base_locally(monkeypatch): + streamed, closed, base_loads = [], [], [] + + class _FakeHttpBackend: + def stream(self, *a, **k): + streamed.append("tuned") + return iter(["tuned-answer"]) + + def close(self): + closed.append("http") + + class _FakeBaseBackend: + def stream(self, *a, **k): + streamed.append("base") + return iter(["base-answer"]) + + def close(self): + closed.append("base") + + def fake_local_load(model, **kwargs): + base_loads.append((model, kwargs.get("fresh_backend", False))) + return _FakeBaseBackend() + + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend()) + monkeypatch.setattr(chatmod, "load_chat_backend", fake_local_load) + + result = CliRunner().invoke(_chat_app(), ["tuned-run"], input = "/compare\nhi\n/exit\n") + + assert result.exit_code == 0, result.output + assert "(compare on)" in result.output + # Only the base model loaded locally, on its own private backend. + assert base_loads == [("fake/base", True)] + assert streamed == ["base", "tuned"] + assert set(closed) == {"http", "base"} + + +def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch): + loads, streamed, closed = [], [], [] + + class _FakeLocalBackend: + def __init__(self, role): + self.role = role + + def stream(self, *a, **k): + streamed.append((self.role, k.get("use_adapter"))) + return iter([f"{self.role}-answer"]) + + def close(self): + closed.append(self.role) + + def fake_load(model, **kwargs): + fresh = kwargs.get("fresh_backend", False) + loads.append((model, fresh)) + return _FakeLocalBackend("base" if fresh else "tuned") + + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "load_chat_backend", fake_load) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: True) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None) + + result = CliRunner().invoke(_chat_app(), ["tuned-run", "--compare"], input = "hi\n/exit\n") + + assert result.exit_code == 0, result.output + assert loads == [("tuned-run", False), ("fake/base", True)] + # Both models answered the turn, via plain generation (no adapter toggle). + assert ("base", None) in streamed and ("tuned", None) in streamed + assert set(closed) == {"tuned", "base"}