mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 09:21:18 +08:00
42672f5446
* add an option to apply the tokenizer chat template * fix the option to apply the tokenizer chat template * better error messages for chat template issues * apply the chat template by default when possible * nit in comment' * rebase --------- Co-authored-by: Awni Hannun <awni@apple.com>
113 lines
3.2 KiB
Python
113 lines
3.2 KiB
Python
import argparse
|
|
import time
|
|
|
|
import mlx.core as mx
|
|
|
|
from .utils import generate_step, load
|
|
|
|
DEFAULT_MODEL_PATH = "mlx_model"
|
|
DEFAULT_PROMPT = "hello"
|
|
DEFAULT_MAX_TOKENS = 100
|
|
DEFAULT_TEMP = 0.6
|
|
DEFAULT_SEED = 0
|
|
|
|
|
|
def setup_arg_parser():
|
|
"""Set up and return the argument parser."""
|
|
parser = argparse.ArgumentParser(description="LLM inference script")
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="mlx_model",
|
|
help="The path to the local model directory or Hugging Face repo.",
|
|
)
|
|
parser.add_argument(
|
|
"--trust-remote-code",
|
|
action="store_true",
|
|
help="Enable trusting remote code for tokenizer",
|
|
)
|
|
parser.add_argument(
|
|
"--eos-token",
|
|
type=str,
|
|
default=None,
|
|
help="End of sequence token for tokenizer",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt", default=DEFAULT_PROMPT, help="Message to be processed by the model"
|
|
)
|
|
parser.add_argument(
|
|
"--max-tokens",
|
|
"-m",
|
|
type=int,
|
|
default=DEFAULT_MAX_TOKENS,
|
|
help="Maximum number of tokens to generate",
|
|
)
|
|
parser.add_argument(
|
|
"--temp", type=float, default=DEFAULT_TEMP, help="Sampling temperature"
|
|
)
|
|
parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
|
|
parser.add_argument(
|
|
"--ignore-chat-template",
|
|
action="store_true",
|
|
help="Use the raw prompt without the tokenizer's chat template.",
|
|
)
|
|
return parser
|
|
|
|
|
|
def main(args):
|
|
mx.random.seed(args.seed)
|
|
|
|
# Building tokenizer_config
|
|
tokenizer_config = {"trust_remote_code": True if args.trust_remote_code else None}
|
|
if args.eos_token is not None:
|
|
tokenizer_config["eos_token"] = args.eos_token
|
|
|
|
model, tokenizer = load(args.model, tokenizer_config=tokenizer_config)
|
|
|
|
if not args.ignore_chat_template and (
|
|
hasattr(tokenizer, "apply_chat_template")
|
|
and tokenizer.chat_template is not None
|
|
):
|
|
messages = [{"role": "user", "content": args.prompt}]
|
|
prompt = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True
|
|
)
|
|
else:
|
|
prompt = args.prompt
|
|
|
|
print("=" * 10)
|
|
print("Prompt:", prompt)
|
|
prompt = tokenizer.encode(prompt)
|
|
prompt = mx.array(prompt)
|
|
tic = time.time()
|
|
tokens = []
|
|
skip = 0
|
|
for token, n in zip(
|
|
generate_step(prompt, model, args.temp), range(args.max_tokens)
|
|
):
|
|
if token == tokenizer.eos_token_id:
|
|
break
|
|
if n == 0:
|
|
prompt_time = time.time() - tic
|
|
tic = time.time()
|
|
tokens.append(token.item())
|
|
s = tokenizer.decode(tokens)
|
|
print(s[skip:], end="", flush=True)
|
|
skip = len(s)
|
|
print(tokenizer.decode(tokens)[skip:], flush=True)
|
|
gen_time = time.time() - tic
|
|
print("=" * 10)
|
|
if len(tokens) == 0:
|
|
print("No tokens generated for this prompt")
|
|
return
|
|
prompt_tps = prompt.size / prompt_time
|
|
gen_tps = (len(tokens) - 1) / gen_time
|
|
print(f"Prompt: {prompt_tps:.3f} tokens-per-sec")
|
|
print(f"Generation: {gen_tps:.3f} tokens-per-sec")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = setup_arg_parser()
|
|
args = parser.parse_args()
|
|
main(args)
|