File size: 2,689 Bytes
04558eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | #!/usr/bin/env python3
"""Benchmark inference latency for LoRA vs DoRA vs base model."""
import time
from mlx_lm import load, generate
SYS = "Convert the dictated text into the exact syntax it represents. Output only the result."
# A few representative samples at different lengths
samples = [
{"label": "short", "input": "dash dash verbose", "expected": "--verbose"},
{"label": "medium", "input": "export all caps API underscore KEY equals quote my dash key dash one two three quote", "expected": 'export API_KEY="my-key-123"'},
{"label": "long", "input": "git add dash A and and git commit dash M quote fix typo quote and and git push", "expected": 'git add -A && git commit -m "fix typo" && git push'},
]
configs = [
{"label": "Qwen 0.5B (base)", "model": "mlx-community/Qwen2.5-0.5B-Instruct-4bit", "adapter": None},
{"label": "Qwen 0.5B + LoRA", "model": "mlx-community/Qwen2.5-0.5B-Instruct-4bit", "adapter": "/Users/arach/dev/talkie/datasets/finetune/adapters/qwen-0.5b-lora"},
{"label": "Qwen 0.5B + DoRA", "model": "mlx-community/Qwen2.5-0.5B-Instruct-4bit", "adapter": "/Users/arach/dev/talkie/datasets/finetune/adapters/qwen-0.5b-dora"},
]
WARMUP = 2
RUNS = 10
for cfg in configs:
print(f"\n{'='*60}")
print(f" {cfg['label']}")
print(f"{'='*60}")
kwargs = {"adapter_path": cfg["adapter"]} if cfg["adapter"] else {}
model, tokenizer = load(cfg["model"], **kwargs)
for sample in samples:
messages = [
{"role": "system", "content": SYS},
{"role": "user", "content": sample["input"]},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
# Warmup
for _ in range(WARMUP):
generate(model, tokenizer, prompt=prompt, max_tokens=60, verbose=False)
# Timed runs
times = []
outputs = []
for _ in range(RUNS):
t0 = time.perf_counter()
out = generate(model, tokenizer, prompt=prompt, max_tokens=60, verbose=False)
t1 = time.perf_counter()
times.append((t1 - t0) * 1000) # ms
outputs.append(out.strip())
avg = sum(times) / len(times)
mn = min(times)
mx = max(times)
last_out = outputs[-1]
match = "✓" if last_out == sample["expected"] else "✗"
print(f"\n [{sample['label']}] \"{sample['input'][:50]}{'...' if len(sample['input'])>50 else ''}\"")
print(f" avg: {avg:.1f}ms min: {mn:.1f}ms max: {mx:.1f}ms ({RUNS} runs)")
print(f" output: {last_out} {match}")
del model, tokenizer
print(f"\n{'='*60}")
print(" Done")
print(f"{'='*60}")
|