"""Python entrypoint of chat.""" import dataclasses from typing import Any, Dict, List, Optional, Union # noqa: UP035 from prompt_toolkit import prompt as get_prompt from prompt_toolkit.key_binding import KeyBindings from mlc_llm.json_ffi import JSONFFIEngine from mlc_llm.protocol import openai_api_protocol from mlc_llm.serve.config import EngineConfig from mlc_llm.serve.engine import MLCEngine from mlc_llm.serve.engine_base import _query_engine_metrics from mlc_llm.support import argparse from mlc_llm.support.config import ConfigOverrideBase def _print_help_str(): help_str = """You can use the following special commands: /help print the special commands /exit quit the cli /stats print out stats of last request (token/sec) /metrics print out full engine metrics /reset restart a fresh chat /set [overrides] override settings in the generation config. For example, `/set temperature=0.5;top_p=0.8;seed=23;max_tokens=100;stop=str1,str2` Note: Separate stop words in the `stop` option with commas (,). Multi-line input: Use escape+enter to start a new line. """ print(help_str) def _set_up_key_bindings(): kb = KeyBindings() @kb.add("escape", "enter") def _(event): event.current_buffer.insert_text("\n") @kb.add("enter") def _(event): event.current_buffer.validate_and_handle() return kb @dataclasses.dataclass class ChatCompletionOverride(ConfigOverrideBase): """Flags for overriding chat completions.""" temperature: Optional[float] = None top_p: Optional[float] = None frequency_penalty: Optional[float] = None presence_penalty: Optional[float] = None max_tokens: Optional[int] = None seed: Optional[int] = None stop: Optional[Union[str, List[str]]] = None # noqa: UP006 @staticmethod def from_str(source: str) -> "ChatCompletionOverride": """Parse model config override values from a string.""" parser = argparse.ArgumentParser(description="chat completion override values") parser.add_argument("--temperature", type=float, default=None) parser.add_argument("--top_p", type=float, default=None) parser.add_argument("--frequency_penalty", type=float, default=None) parser.add_argument("--presence_penalty", type=float, default=None) parser.add_argument("--max_tokens", type=int, default=None) parser.add_argument("--seed", type=int, default=None) parser.add_argument("--stop", type=str, default=None) results = parser.parse_args([f"--{i}" for i in source.split(";") if i]) return ChatCompletionOverride( temperature=results.temperature, top_p=results.top_p, frequency_penalty=results.frequency_penalty, presence_penalty=results.presence_penalty, max_tokens=results.max_tokens, seed=results.seed, stop=results.stop.split(",") if results.stop is not None else None, ) @dataclasses.dataclass class ModelConfigOverride(ConfigOverrideBase): """Flags for overriding model config.""" context_window_size: Optional[int] = None sliding_window_size: Optional[int] = None prefill_chunk_size: Optional[int] = None attention_sink_size: Optional[int] = None tensor_parallel_shards: Optional[int] = None pipeline_parallel_stages: Optional[int] = None opt: Optional[str] = None @staticmethod def from_str(source: str) -> "ModelConfigOverride": """Parse model config override values from a string.""" parser = argparse.ArgumentParser(description="model config override values") parser.add_argument("--tensor_parallel_shards", type=int, default=None) parser.add_argument("--pipeline_parallel_stages", type=int, default=None) parser.add_argument("--opt", type=str, default=None) parser.add_argument("--context_window_size", type=int, default=None) parser.add_argument("--sliding_window_size", type=int, default=None) parser.add_argument("--prefill_chunk_size", type=int, default=None) parser.add_argument("--attention_sink_size", type=int, default=None) results = parser.parse_args([f"--{i}" for i in source.split(";") if i]) return ModelConfigOverride( tensor_parallel_shards=results.tensor_parallel_shards, pipeline_parallel_stages=results.pipeline_parallel_stages, opt=results.opt, context_window_size=results.context_window_size, sliding_window_size=results.sliding_window_size, prefill_chunk_size=results.prefill_chunk_size, attention_sink_size=results.attention_sink_size, ) class ChatState: """Simple helper class to manage chat state. Chat state wraps around a engine instance and exposes the minimum set of tools to perform interactive chat. It provides support for mlc_llm chat. It also can be used to do interactive debugging with different engine instance. Examples -------- .. code:: python from openai import OpenAI from mlc_llm import MLCEngine from mlc_llm.serve import PopenServer from mlc_llm.interface.chat import ChatState def chat_with_engine(model): # hookup with MLCEngine ChatState(MLCEngine(model)).chat() def chat_with_server(model): # hookup with AsyncMLCEngine backed api server with PopenServer(model) as server: ChatState( OpenAI(base_url=server.openai_v1_base_url, api_key="None") ).chat() """ history: List[Dict[str, Any]] # noqa: UP006 history_begin: int # kwargs passed to completions overrides: ChatCompletionOverride # Underlying engine engine: Union[JSONFFIEngine, MLCEngine] last_finished_request_usage: Optional[openai_api_protocol.CompletionUsage] def __init__(self, engine: Union[JSONFFIEngine, MLCEngine]): self.engine = engine self.history = [] self.history_window_begin = 0 self.overrides = ChatCompletionOverride() # model is mainly used for compact reasons self.model = "chat_model" self.last_finished_request_usage = None def slide_history(self): """Slide history to fit into context window""" history_window_size = len(self.history) - self.history_window_begin assert history_window_size % 2 == 0 self.history_window_begin += ((history_window_size + 3) // 4) * 2 def process_system_prompts(self): """Process system prompts""" # TODO(mlc-team): possibly leverage debug option # pass a simple prompt to warm up for _ in self.engine.chat.completions.create( messages=[{"role": "user", "content": ""}], max_tokens=1, model=self.model, stream=True, ): pass def generate(self, prompt: str): """Run one generation with the prompt. Parameters ---------- prompt: str The input prompt """ self.history.append({"role": "user", "content": prompt}) output_text = "" finish_reason_length = False messages = self.history[self.history_window_begin :] for response in self.engine.chat.completions.create( messages=messages, model=self.model, stream=True, stream_options={"include_usage": True}, **dataclasses.asdict(self.overrides), ): if response.usage is not None: self.last_finished_request_usage = response.usage continue for choice in response.choices: assert choice.delta.role == "assistant" if isinstance(choice.delta.content, str): output_text += choice.delta.content print(choice.delta.content, end="", flush=True) if choice.finish_reason == "length": finish_reason_length = True if finish_reason_length: print(" [output truncated due to context length limit...]") # print additional \n when generation ends print() # record the history self.history.append({"role": "assistant", "content": output_text}) if finish_reason_length: self.slide_history() def stats(self): """Print statistics of the prefill and decode speed.""" def get_stats_text(): """Get text""" if self.last_finished_request_usage is None: return "N/A" last_finished_request = self.last_finished_request_usage.extra if last_finished_request is None: return "N/A" prefill_speed = last_finished_request.get("prefill_tokens_per_s", None) decode_speed = last_finished_request.get("decode_tokens_per_s", None) prefill_speed = f"{prefill_speed:.1f}" if prefill_speed is not None else "N/A" decode_speed = f"{decode_speed:.1f}" if decode_speed is not None else "N/A" return f"prefill: {prefill_speed} tok/s, decode: {decode_speed} tok/s" print(get_stats_text(), flush=True) def metrics(self): """Print metrics as prometheus text""" print(_query_engine_metrics(self.engine).prometheus_text(), flush=True) def reset(self): """Reset the chat history""" self.history = [] self.history_window_begin = 0 def chat(self): """Start an interactive chat session.""" _print_help_str() self.process_system_prompts() # Multi-line input support: set escape+enter as start a new line kb = _set_up_key_bindings() while True: try: prompt = get_prompt( ">>> ", key_bindings=kb, multiline=True, ) except (KeyboardInterrupt, EOFError): break if prompt[:4] == "/set": overrides = ChatCompletionOverride.from_str(prompt.split()[1]) for key, value in dataclasses.asdict(overrides).items(): if value is not None: setattr(self.overrides, key, value) elif prompt[:6] == "/stats": self.stats() elif prompt[:8] == "/metrics": self.metrics() elif prompt[:6] == "/reset": self.reset() elif prompt[:5] == "/exit": break elif prompt[:5] == "/help": _print_help_str() else: self.generate(prompt) def chat( model: str, device: str, model_lib: Optional[str], overrides: ModelConfigOverride, ): """Chat cli entry""" # By default we use JSONFFIEngine engine = JSONFFIEngine( model, device, model_lib=model_lib, mode="interactive", engine_config=EngineConfig( max_single_sequence_length=overrides.context_window_size, prefill_chunk_size=overrides.prefill_chunk_size, sliding_window_size=overrides.sliding_window_size, attention_sink_size=overrides.attention_sink_size, tensor_parallel_shards=overrides.tensor_parallel_shards, pipeline_parallel_stages=overrides.pipeline_parallel_stages, opt=overrides.opt, ), ) try: ChatState(engine).chat() finally: engine.terminate()