New UI for chat and lora (ml-explore#1344)

Author: nastya236

mlx_lm/chat.py

@@ -4,6 +4,7 @@
 
 import mlx.core as mx
 
+from .cli_ui import ChatUI
 from .generate import stream_generate
 from .models.cache import make_prompt_cache
 from .sample_utils import make_sampler
@@ -96,10 +97,6 @@ def main():
     pipeline_group = group if args.pipeline else None
     tensor_group = group if not args.pipeline else None
 
-    def rprint(*args, **kwargs):
-        if rank == 0:
-            print(*args, **kwargs)
-
     mx.random.seed(args.seed)
 
     if group.size() > 1:
@@ -115,51 +112,48 @@ def rprint(*args, **kwargs):
             },
         )
 
-    def print_help():
-        rprint("The command list:")
-        rprint("- 'q' to exit")
-        rprint("- 'r' to reset the chat")
-        rprint("- 'h' to display these commands")
-
-    rprint(f"[INFO] Starting chat session with {args.model}.")
-    print_help()
-    prompt_cache = make_prompt_cache(model, args.max_kv_size)
-    while True:
-        query = input(">> " if rank == 0 else "")
-        if query == "q":
-            break
-        if query == "r":
-            prompt_cache = make_prompt_cache(model, args.max_kv_size)
-            continue
-        if query == "h":
-            print_help()
-            continue
-        messages = []
-        if args.system_prompt is not None:
-            messages.append({"role": "system", "content": args.system_prompt})
-        messages.append({"role": "user", "content": query})
-        prompt = tokenizer.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-        )
-        for response in stream_generate(
-            model,
-            tokenizer,
-            prompt,
-            max_tokens=args.max_tokens,
-            sampler=make_sampler(
-                args.temp,
-                args.top_p,
-                xtc_threshold=args.xtc_threshold,
-                xtc_probability=args.xtc_probability,
-                xtc_special_tokens=(
-                    tokenizer.encode("\n") + list(tokenizer.eos_token_ids)
+    with ChatUI(args, rank=rank) as ui:
+        prompt_cache = make_prompt_cache(model, args.max_kv_size)
+        while True:
+            query = ui.prompt()
+            if query == "q":
+                ui.say_bye()
+                break
+            if query == "r":
+                prompt_cache = make_prompt_cache(model, args.max_kv_size)
+                ui.say_reset()
+                continue
+            if query == "h":
+                ui.say_help()
+                continue
+            messages = []
+            if args.system_prompt is not None:
+                messages.append({"role": "system", "content": args.system_prompt})
+            messages.append({"role": "user", "content": query})
+            prompt = tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+            )
+            last_response = None
+            for response in stream_generate(
+                model,
+                tokenizer,
+                prompt,
+                max_tokens=args.max_tokens,
+                sampler=make_sampler(
+                    args.temp,
+                    args.top_p,
+                    xtc_threshold=args.xtc_threshold,
+                    xtc_probability=args.xtc_probability,
+                    xtc_special_tokens=(
+                        tokenizer.encode("\n") + list(tokenizer.eos_token_ids)
+                    ),
                 ),
-            ),
-            prompt_cache=prompt_cache,
-        ):
-            rprint(response.text, flush=True, end="")
-        rprint()
+                prompt_cache=prompt_cache,
+            ):
+                ui.stream_token(response.text)
+                last_response = response
+            ui.end_turn(last_response)
 
 
 if __name__ == "__main__":