chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,311 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user