File size: 2,739 Bytes
63d9fb8 | 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
import mlx.core as mx
import numpy as np
from ark_asr_mlx import ArkASR
from ark_asr_mlx.generation import _next_token
def benchmark(model_path: Path, audio_path: Path) -> dict[str, object]:
mx.reset_peak_memory()
started = time.perf_counter()
asr = ArkASR.from_pretrained(model_path)
load_seconds = time.perf_counter() - started
process_started = time.perf_counter()
inputs = asr.processor.process(audio_path)
processing_seconds = time.perf_counter() - process_started
input_ids = mx.array(inputs.input_ids)
input_features = mx.array(inputs.input_features).astype(
asr.model.audio_encoder.whisper.conv1.weight.dtype
)
generation_started = time.perf_counter()
embeddings = asr.model.prepare_prompt_embeddings(
input_ids,
input_features,
inputs.audio_start,
inputs.audio_count,
)
cache = asr.model.make_cache()
logits = asr.model(input_ids, cache=cache, input_embeddings=embeddings)
mx.eval(logits)
suppressed = asr.processor.suppressed_token_ids(asr.model.config.eos_token_id)
mask = np.zeros(asr.model.config.vocab_size, dtype=np.bool_)
mask[suppressed] = True
suppression_mask = mx.array(mask)
first_token = _next_token(logits, suppression_mask)
first_token_seconds = time.perf_counter() - generation_started
generated: list[int] = []
token = first_token
while len(generated) < 256 and token != asr.model.config.eos_token_id:
generated.append(token)
logits = asr.model(mx.array([[token]], dtype=mx.int32), cache=cache)
mx.eval(logits)
token = _next_token(logits, suppression_mask)
generation_seconds = time.perf_counter() - generation_started
text = asr.processor.tokenizer.decode(generated, skip_special_tokens=True).strip()
return {
"audio": str(audio_path),
"load_seconds": load_seconds,
"processing_seconds": processing_seconds,
"first_token_seconds": first_token_seconds,
"generation_seconds": generation_seconds,
"generation_tokens": len(generated),
"tokens_per_second": len(generated) / generation_seconds,
"peak_memory_gb": mx.get_peak_memory() / 1_000_000_000,
"text": text,
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("audio", type=Path)
parser.add_argument(
"--model",
type=Path,
default=Path("models/ARK-ASR-0.6B-bf16"),
)
args = parser.parse_args()
print(json.dumps(benchmark(args.model, args.audio), indent=2, ensure_ascii=False))
if __name__ == "__main__":
main()
|