mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
Studio: add unsloth chat CLI command (#6170)
* Studio: add `unsloth chat` CLI command Interactive chat REPL on the shared Studio backend: trained-model picker when no model is given, /think and /compare toggles (adapter toggle on CUDA, side-by-side base-model load on MLX), markdown streaming, and connect-if-running Studio server mode so models stay warm across sessions and are shared with the UI. * fix settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix error handling and compare base precision * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix chat CLI backend imports and GGUF drafter loading * Hide split thinking tags in chat CLI streams --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 <samleejackson0@gmail.com>
This commit is contained in:
Nilay
GitHub
parent
3733e0b274
commit
f64c3c8aba
6 changed files with 1208 additions and 40 deletions
|
|
@ -26,6 +26,7 @@ classifiers = [
|
|||
]
|
||||
dependencies = [
|
||||
"typer",
|
||||
"rich",
|
||||
"pydantic",
|
||||
"pyyaml",
|
||||
"nest-asyncio",
|
||||
|
|
|
|||
|
|
@ -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.")
|
||||
|
|
|
|||
417
unsloth_cli/_inference.py
Normal file
417
unsloth_cli/_inference.py
Normal file
|
|
@ -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>"
|
||||
_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", 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 <think> 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
|
||||
340
unsloth_cli/commands/chat.py
Normal file
340
unsloth_cli/commands/chat.py
Normal file
|
|
@ -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 <model>`.",
|
||||
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 <think> 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")
|
||||
|
|
@ -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 <think> 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()
|
||||
|
|
|
|||
401
unsloth_cli/tests/test_inference_chat.py
Normal file
401
unsloth_cli/tests/test_inference_chat.py
Normal file
|
|
@ -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 = "<think>reasoning</think>answer"
|
||||
assert visible_text(text, show_thinking = True) == text
|
||||
|
||||
|
||||
def test_visible_text_strips_closed_think_block():
|
||||
text = "<think>step 1\nstep 2</think>The answer is 42."
|
||||
assert visible_text(text, show_thinking = False) == "The answer is 42."
|
||||
|
||||
|
||||
def test_visible_text_holds_unclosed_think():
|
||||
# An open <think> is held back so partial reasoning never leaks mid-stream.
|
||||
assert visible_text("<think>still thinking", show_thinking = False) == ""
|
||||
assert visible_text("done.<think>more thinking", show_thinking = False) == "done."
|
||||
|
||||
|
||||
def test_visible_text_holds_partial_think_prefix():
|
||||
# Streams are cumulative, so the opening tag can arrive as "<", "<thi",
|
||||
# then "<think>". Hold possible tag prefixes until they are disambiguated.
|
||||
assert visible_text("<", show_thinking = False) == ""
|
||||
assert visible_text("<thi", show_thinking = False) == ""
|
||||
assert visible_text("done.<thi", show_thinking = False) == "done."
|
||||
assert visible_text("2 < 3", show_thinking = False) == "2 < 3"
|
||||
|
||||
|
||||
def _option(command_fn, name):
|
||||
return inspect.signature(command_fn).parameters[name].default
|
||||
|
||||
|
||||
def test_inference_think_defaults_off():
|
||||
from unsloth_cli.commands.inference import inference
|
||||
|
||||
opt = _option(inference, "think")
|
||||
assert getattr(opt, "default", None) is False
|
||||
# typer stores a flag/--no-flag pair as one combined decl.
|
||||
assert "--think/--no-think" in (getattr(opt, "param_decls", None) or [])
|
||||
|
||||
|
||||
def test_chat_command_is_registered_with_options():
|
||||
params = inspect.signature(chatmod.chat).parameters
|
||||
assert "model" in params
|
||||
|
||||
think = _option(chatmod.chat, "think")
|
||||
assert "--think/--no-think" in (getattr(think, "param_decls", None) or [])
|
||||
|
||||
compare = _option(chatmod.chat, "compare")
|
||||
assert "--compare/--no-compare" in (getattr(compare, "param_decls", None) or [])
|
||||
|
||||
verbose = _option(chatmod.chat, "verbose")
|
||||
assert {"--verbose", "-v"} <= set(getattr(verbose, "param_decls", None) or [])
|
||||
|
||||
|
||||
class _FakeBackend:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def generate_chat_response(self, **kwargs):
|
||||
self.calls.append(("plain", None, kwargs))
|
||||
return iter(["hi"])
|
||||
|
||||
def generate_with_adapter_control(self, *, use_adapter, **kwargs):
|
||||
self.calls.append(("adapter", use_adapter, kwargs))
|
||||
return iter(["hi"])
|
||||
|
||||
|
||||
_STREAM_KWARGS = dict(
|
||||
system_prompt = "",
|
||||
temperature = 0.7,
|
||||
top_p = 0.9,
|
||||
top_k = 40,
|
||||
max_new_tokens = 8,
|
||||
repetition_penalty = 1.1,
|
||||
enable_thinking = False,
|
||||
)
|
||||
|
||||
|
||||
def test_chatbackend_routes_compare_to_adapter_control():
|
||||
fake = _FakeBackend()
|
||||
backend = ChatBackend("unsloth", fake)
|
||||
|
||||
list(backend.stream([{"role": "user", "content": "x"}], use_adapter = False, **_STREAM_KWARGS))
|
||||
list(backend.stream([{"role": "user", "content": "x"}], use_adapter = True, **_STREAM_KWARGS))
|
||||
|
||||
assert [(path, flag) for path, flag, _ in fake.calls] == [
|
||||
("adapter", False),
|
||||
("adapter", True),
|
||||
]
|
||||
|
||||
|
||||
def test_chatbackend_normal_path_skips_adapter_control():
|
||||
fake = _FakeBackend()
|
||||
backend = ChatBackend("unsloth", fake)
|
||||
|
||||
list(backend.stream([{"role": "user", "content": "x"}], **_STREAM_KWARGS))
|
||||
|
||||
assert fake.calls[0][0] == "plain"
|
||||
|
||||
|
||||
def test_collect_stream_returns_last_cumulative_think_stripped():
|
||||
stream = iter(["<think>r</think>hel", "<think>r</think>hello"])
|
||||
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"}
|
||||
Loading…
Add table
Add a link
Reference in a new issue