mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-09-18 10:26:58 +08:00
51 lines
1.1 KiB
Python
51 lines
1.1 KiB
Python
# Copyright © 2024 Apple Inc.
|
|
|
|
"""
|
|
An example of a multi-turn chat with prompt caching.
|
|
"""
|
|
|
|
from mlx_lm import generate, load
|
|
from mlx_lm.models.cache import make_prompt_cache
|
|
|
|
model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
|
|
|
|
# Make the initial prompt cache for the model
|
|
prompt_cache = make_prompt_cache(model)
|
|
|
|
# User turn
|
|
prompt = "Hi my name is <Name>."
|
|
messages = [{"role": "user", "content": prompt}]
|
|
prompt = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True
|
|
)
|
|
|
|
# Assistant response
|
|
response = generate(
|
|
model,
|
|
tokenizer,
|
|
prompt=prompt,
|
|
verbose=True,
|
|
max_tokens=1024,
|
|
temp=0.0,
|
|
prompt_cache=prompt_cache,
|
|
)
|
|
messages.append({"role": "assistant", "content": response})
|
|
|
|
# User turn
|
|
prompt = "What's my name?"
|
|
messages = [{"role": "user", "content": prompt}]
|
|
prompt = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True
|
|
)
|
|
|
|
# Assistant response
|
|
response = generate(
|
|
model,
|
|
tokenizer,
|
|
prompt=prompt,
|
|
verbose=True,
|
|
max_tokens=1024,
|
|
temp=0.0,
|
|
prompt_cache=prompt_cache,
|
|
)
|