| import mlx.core as mx |
|
|
| from ark_asr_mlx.config import ArkASRConfig |
| from ark_asr_mlx.model import ArkASRModel |
|
|
|
|
| def _tiny_config() -> ArkASRConfig: |
| return ArkASRConfig.from_dict( |
| { |
| "hidden_size": 8, |
| "intermediate_size": 16, |
| "num_attention_heads": 2, |
| "num_hidden_layers": 2, |
| "num_key_value_heads": 1, |
| "rms_norm_eps": 1e-6, |
| "rope_theta": 10_000, |
| "vocab_size": 20, |
| "audio_token_id": 10, |
| "eos_token_id": 2, |
| "pad_token_id": 0, |
| "merge_factor": 2, |
| "whisper_config": { |
| "d_model": 8, |
| "encoder_attention_heads": 2, |
| "encoder_ffn_dim": 16, |
| "encoder_layers": 1, |
| "num_mel_bins": 4, |
| }, |
| } |
| ) |
|
|
|
|
| def test_prompt_injection_and_kv_cache() -> None: |
| model = ArkASRModel(_tiny_config()) |
| input_ids = mx.array([[1, 10, 10, 3, 4]]) |
| input_features = mx.zeros((1, 4, 8)) |
| embeddings = model.prepare_prompt_embeddings( |
| input_ids, |
| input_features, |
| audio_start=1, |
| audio_count=2, |
| ) |
| assert embeddings.shape == (1, 5, 8) |
|
|
| cache = model.make_cache() |
| prompt_logits = model(input_ids, cache=cache, input_embeddings=embeddings) |
| next_logits = model(mx.array([[5]]), cache=cache) |
| mx.eval(prompt_logits, next_logits) |
|
|
| assert prompt_logits.shape == (1, 5, 20) |
| assert next_logits.shape == (1, 1, 20) |
| assert all(layer_cache.offset == 6 for layer_cache in cache) |
|
|
