feat: enhance chat CLI with readline history, line editing, and distributed support

mlx_lm/chat.py

@@ -1,6 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+import readline  # noqa: F401  # Enables terminal line editing/history on rank 0.
 
 import mlx.core as mx
 
@@ -16,7 +17,33 @@
 DEFAULT_XTC_THRESHOLD = 0.0
 DEFAULT_SEED = 0
 DEFAULT_MAX_TOKENS = 256
+DEFAULT_RENDER_WINDOW_SIZE = 20
+DEFAULT_REFRESH_RATE = 10
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+DEFAULT_SYSTEM_PROMPT = (
+    "You are a helpful assistant. Your responses are rendered in a terminal with "
+    "Markdown support. Feel free to use Markdown formatting when appropriate: "
+    "**bold**, *italic*, `inline code`, code blocks with syntax highlighting "
+    "(```language), bullet lists, numbered lists, and headers."
+)
+
+
+def broadcast_string(
+    value: str, group: mx.distributed.Group, src: int = 0
+) -> str:
+    """Broadcast a UTF-8 string from src to every rank in group."""
+    if group.size() == 1:
+        return value
+    if group.rank() == src:
+        data = mx.array(value.encode("utf-8"))
+        mx.eval(mx.distributed.all_sum(data.size, group=group))
+        mx.eval(mx.distributed.all_sum(data, group=group))
+        return value
+
+    size = mx.distributed.all_sum(0, group=group).item()
+    data = mx.distributed.all_sum(mx.zeros(size, dtype=mx.uint8), group=group)
+    mx.eval(data)
+    return bytes(data).decode("utf-8")
 
 
 def setup_arg_parser():
@@ -53,8 +80,8 @@ def setup_arg_parser():
     parser.add_argument(
         "--xtc-threshold",
         type=float,
-        default=0.0,
-        help="Thresold the probs of each next token candidate to be sampled by XTC",
+        default=DEFAULT_XTC_THRESHOLD,
+        help="Threshold the probs of each next token candidate to be sampled by XTC",
     )
     parser.add_argument(
         "--seed",
@@ -78,13 +105,31 @@ def setup_arg_parser():
     parser.add_argument(
         "--system-prompt",
         default=None,
-        help="System prompt to be used for the chat template",
+        help="System prompt to be used for the chat template "
+        "(replaces the default Markdown-aware prompt)",
+    )
+    parser.add_argument(
+        "--no-system-prompt",
+        action="store_true",
+        help="Disable the default system prompt entirely",
     )
     parser.add_argument(
         "--pipeline",
         action="store_true",
         help="Use pipelining instead of tensor parallelism",
     )
+    parser.add_argument(
+        "--window-size",
+        type=int,
+        default=DEFAULT_RENDER_WINDOW_SIZE,
+        help="The number of recent rendered lines to keep in the live panel",
+    )
+    parser.add_argument(
+        "--refresh-rate",
+        type=int,
+        default=DEFAULT_REFRESH_RATE,
+        help="The live panel refresh rate during generation",
+    )
     return parser
 
 
@@ -115,7 +160,11 @@ def main():
     with ChatUI(args, rank=rank) as ui:
         prompt_cache = make_prompt_cache(model, args.max_kv_size)
         while True:
-            query = ui.prompt()
+            query = ui.prompt() if rank == 0 else ""
+            query = broadcast_string(query, group).strip()
+
+            if not query:
+                continue
             if query == "q":
                 ui.say_bye()
                 break
@@ -127,7 +176,10 @@ def main():
                 ui.say_help()
                 continue
             messages = []
-            if args.system_prompt is not None:
+            if not args.no_system_prompt:
+                system_content = args.system_prompt or DEFAULT_SYSTEM_PROMPT
+                messages.append({"role": "system", "content": system_content})
+            elif args.system_prompt:
                 messages.append({"role": "system", "content": args.system_prompt})
             messages.append({"role": "user", "content": query})
             prompt = tokenizer.apply_chat_template(

mlx_lm/cli_ui.py

@@ -1,6 +1,7 @@
 # Copyright © 2024 Apple Inc.
 
 import re
+import readline  # noqa: F401  # Enables terminal line editing/history.
 import shutil
 import sys
 from contextlib import contextmanager
@@ -9,6 +10,8 @@
 import mlx.core as mx
 from rich.box import ROUNDED
 from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
 from rich.panel import Panel
 from rich.progress import Progress, ProgressColumn, TextColumn
 from rich.text import Text
@@ -51,6 +54,8 @@ def make_console() -> Console:
         theme=_make_theme(),
         highlight=False,
         color_system="truecolor",
+        force_terminal=True,
+        force_interactive=True,
         width=_terminal_width(),
     )
 
@@ -256,6 +261,10 @@ def __init__(self, args, rank: int = 0):
         self._rank = rank
         self._args = args
         self._console = make_console()
+        self._response_text = ""
+        self._live = None
+        self._window_size = max(getattr(args, "window_size", 20), 1)
+        self._refresh_rate = max(getattr(args, "refresh_rate", 10), 1)
 
     def __enter__(self):
         if self._rank == 0:
@@ -294,16 +303,59 @@ def say_help(self):
         if self._rank == 0:
             print_chat_help(self._console)
 
+    def _display_text(self) -> str:
+        lines = self._response_text.splitlines(keepends=True)
+        if len(lines) > self._window_size:
+            return "".join(lines[-self._window_size :])
+        return self._response_text
+
+    def _ensure_live(self):
+        if self._rank != 0 or self._live is not None:
+            return
+        self._live = Live(
+            Panel(
+                Markdown(""),
+                title="[ui.accent]generating[/ui.accent]",
+                border_style="ui.border",
+                box=ROUNDED,
+            ),
+            console=self._console,
+            refresh_per_second=self._refresh_rate,
+            transient=True,
+        )
+        self._live.start()
+
     def stream_token(self, text: str):
-        rprint(text, flush=True, end="")
+        if self._rank != 0:
+            return
+        self._ensure_live()
+        self._response_text += text
+        self._live.update(
+            Panel(
+                Markdown(self._display_text()),
+                title="[ui.accent]generating[/ui.accent]",
+                border_style="ui.border",
+                box=ROUNDED,
+            )
+        )
 
     def end_turn(self, response):
-        rprint()  # newline after the streamed line
+        if self._live is not None:
+            self._live.stop()
+            self._live = None
         if self._rank != 0 or response is None:
+            self._response_text = ""
             return
+        self._console.print(Markdown(self._response_text))
+        if getattr(response, "finish_reason", None) == "length":
+            self._console.print(
+                f"  [ui.warn]output truncated[/ui.warn] "
+                f"[ui.muted](max tokens: {self._args.max_tokens})[/ui.muted]"
+            )
         self._console.print(
             f"  [ui.muted]{response.generation_tokens} tokens · "
             f"{response.generation_tps:.1f} tok/s · "
             f"prompt {response.prompt_tps:.1f} tok/s · "
             f"peak {response.peak_memory:.2f} GB[/ui.muted]"
         )
+        self._response_text = ""

mlx_lm/sample_utils.py

@@ -141,7 +141,7 @@ def apply_top_k(
     vocab_size = logprobs.shape[-1]
     if not isinstance(top_k, int) or not (0 < top_k < vocab_size):
         raise ValueError(
-            f"`top_k` has to be an integer in the (0, {vocab_size}] interval,"
+            f"`top_k` has to be an integer in the (0, {vocab_size}) interval,"
             f" but is {top_k}."
         )
     mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:]

setup.py

@@ -31,6 +31,7 @@
         "protobuf",
         "pyyaml",
         "jinja2",
+        "rich",
     ],
     packages=[
         "mlx_lm",