unsloth/unsloth_cli/commands/inference.py
Nilay f64c3c8aba
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>
2026-06-11 16:09:34 +01:00

76 lines
2.8 KiB
Python

# 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 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."),
prompt: str = typer.Argument(..., help = "Prompt to send to the model."),
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(256, "--max-new-tokens"),
repetition_penalty: float = typer.Option(1.1, "--repetition-penalty"),
system_prompt: str = typer.Option(
"",
"--system-prompt",
help = "Optional system prompt to prepend.",
),
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."""
if not verbose:
configure_quiet_logging()
# 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()