iamdoubz
40709b257d
resolve_language(Some("EN"), false) was passed through unchanged instead of
being normalized to lowercase "en" - the English-only-model guard only fired
when the requested language *differed* from English, not when it matched with
different casing. Also fixes cargo fmt drift in mod.rs/npu.rs left over from
the M4.2 worktree (stopped mid-verification per session instruction before
running fmt).
548 lines
22 KiB
Rust
548 lines
22 KiB
Rust
//! ONNX transcriber: Whisper ONNX via ONNX Runtime, on either the Intel NPU
|
||
//! (OpenVINO EP, T3.4) or a DX12 GPU (DirectML EP, P2 — the non-Vulkan path for
|
||
//! AMD/Intel). One engine, EP chosen by the requested `BackendId`; same
|
||
//! `Transcriber` trait and `TranscriptSegment` output as the whisper.cpp path,
|
||
//! so callers never branch on engine. The ONNX model artifacts are identical
|
||
//! across EPs (an EP is a device backend, not a different graph).
|
||
//!
|
||
//! Split of work, validated by the T3.4 spike (encoder ~3.6× faster on NPU):
|
||
//! - **encoder** (fixed `[1,80,3000]` shape) runs on the **accelerator** (NPU or
|
||
//! GPU) — the expensive graph, and the fixed shape accelerators want;
|
||
//! - **decoder** (dynamic, autoregressive) runs greedy on the **CPU** EP — a
|
||
//! dynamic KV loop is a poor accelerator fit and the cheap half anyway.
|
||
//!
|
||
//! ponytail: greedy, no-KV-cache decode (re-feeds the full token prefix each
|
||
//! step). Correct and simple; windows are short so the token count is small.
|
||
//! Add `decoder_with_past` only if the decoder shows up in a profile.
|
||
|
||
use super::{AudioWindow, SegmentSink, Transcriber, TrxError};
|
||
use crate::models::{BackendId, TranscriptSegment};
|
||
use ort::execution_providers::{
|
||
DirectMLExecutionProvider, ExecutionProviderDispatch, OpenVINOExecutionProvider,
|
||
};
|
||
use ort::session::Session;
|
||
use ort::value::Tensor;
|
||
use std::collections::HashMap;
|
||
use std::path::Path;
|
||
use std::sync::atomic::{AtomicU64, Ordering};
|
||
use std::sync::Mutex;
|
||
|
||
/// Safety cap on generated tokens per 30 s window (Whisper's own max is 448).
|
||
const MAX_NEW_TOKENS: usize = 224;
|
||
|
||
pub struct OnnxTranscriber {
|
||
encoder: Mutex<Session>,
|
||
decoder: Mutex<Session>,
|
||
enc_input: String,
|
||
enc_output: String,
|
||
dec_ids: String,
|
||
dec_hidden: String,
|
||
dec_logits: String,
|
||
detok: Detok,
|
||
/// `[decoder_start_token_id, ...forced_decoder_ids]` — the fixed prompt
|
||
/// prefix before generation begins.
|
||
initial_tokens: Vec<i64>,
|
||
/// Special tokens start here (`decoder_start_token_id`); anything ≥ this is
|
||
/// masked out of argmax so only text tokens (or `eot`) can be emitted.
|
||
special_floor: i64,
|
||
eot: i64,
|
||
next_id: AtomicU64,
|
||
}
|
||
|
||
impl OnnxTranscriber {
|
||
/// Points `ort` (load-dynamic) at the given on-demand-downloaded runtime DLL
|
||
/// and puts its sibling DLLs on the search path, unless the caller already
|
||
/// set `ORT_DYLIB_PATH` (the test harness does, to target a dev runtime).
|
||
/// Idempotent.
|
||
///
|
||
/// `ORT_DYLIB_PATH` is process-global and `ort` dlopens it once, so a single
|
||
/// process run must not mix EPs from different runtime bundles (OpenVINO vs
|
||
/// DirectML). That's fine: `best()` picks one backend per session.
|
||
// ponytail: one runtime per process. If backend hot-swapping between NPU and
|
||
// DirectML in one run is ever needed, that's a re-init-ort problem to solve then.
|
||
fn ensure_runtime_env(dll: std::path::PathBuf) {
|
||
if std::env::var_os("ORT_DYLIB_PATH").is_some() {
|
||
return;
|
||
}
|
||
std::env::set_var("ORT_DYLIB_PATH", &dll);
|
||
if let Some(dir) = dll.parent() {
|
||
let path = std::env::var_os("PATH").unwrap_or_default();
|
||
let mut dirs = vec![dir.to_path_buf()];
|
||
dirs.extend(std::env::split_paths(&path));
|
||
if let Ok(joined) = std::env::join_paths(dirs) {
|
||
std::env::set_var("PATH", joined);
|
||
}
|
||
}
|
||
}
|
||
|
||
fn find_input(session: &Session, needle: &str) -> Result<String, TrxError> {
|
||
session
|
||
.inputs
|
||
.iter()
|
||
.find(|i| i.name.contains(needle))
|
||
.map(|i| i.name.to_string())
|
||
.ok_or_else(|| TrxError::Load(format!("model missing input '{needle}'")))
|
||
}
|
||
|
||
fn find_output(session: &Session, needle: &str) -> Result<String, TrxError> {
|
||
session
|
||
.outputs
|
||
.iter()
|
||
.find(|o| o.name.contains(needle))
|
||
.map(|o| o.name.to_string())
|
||
.ok_or_else(|| TrxError::Load(format!("model missing output '{needle}'")))
|
||
}
|
||
|
||
/// Full encode→decode over one window's samples, returning the trimmed text.
|
||
fn infer(&self, samples: &[f32]) -> Result<String, TrxError> {
|
||
let mel = super::mel::log_mel_spectrogram(samples);
|
||
let enc_in = Tensor::from_array(([1usize, super::mel::N_MELS, super::mel::N_FRAMES], mel))
|
||
.map_err(|e| TrxError::Inference(e.to_string()))?;
|
||
|
||
// Encoder on the NPU. Copy hidden states out while the session is locked
|
||
// (outputs borrow the session).
|
||
let (hidden_shape, hidden_data) = {
|
||
let mut enc = self
|
||
.encoder
|
||
.lock()
|
||
.map_err(|_| TrxError::Inference("encoder mutex poisoned".into()))?;
|
||
let outputs = enc
|
||
.run(ort::inputs![self.enc_input.as_str() => enc_in.view()])
|
||
.map_err(|e| TrxError::Inference(e.to_string()))?;
|
||
let (shape, data) = outputs[self.enc_output.as_str()]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(|e| TrxError::Inference(e.to_string()))?;
|
||
(
|
||
shape.iter().map(|d| *d as usize).collect::<Vec<_>>(),
|
||
data.to_vec(),
|
||
)
|
||
};
|
||
let hidden = Tensor::from_array((hidden_shape, hidden_data))
|
||
.map_err(|e| TrxError::Inference(e.to_string()))?;
|
||
|
||
// Greedy decode on the CPU.
|
||
let mut tokens = self.initial_tokens.clone();
|
||
for _ in 0..MAX_NEW_TOKENS {
|
||
let ids: Vec<i64> = tokens.clone();
|
||
let ids_tensor = Tensor::from_array(([1usize, ids.len()], ids))
|
||
.map_err(|e| TrxError::Inference(e.to_string()))?;
|
||
|
||
let next = {
|
||
let mut dec = self
|
||
.decoder
|
||
.lock()
|
||
.map_err(|_| TrxError::Inference("decoder mutex poisoned".into()))?;
|
||
let outputs = dec
|
||
.run(ort::inputs![
|
||
self.dec_ids.as_str() => ids_tensor.view(),
|
||
self.dec_hidden.as_str() => hidden.view(),
|
||
])
|
||
.map_err(|e| TrxError::Inference(e.to_string()))?;
|
||
let (shape, data) = outputs[self.dec_logits.as_str()]
|
||
.try_extract_tensor::<f32>()
|
||
.map_err(|e| TrxError::Inference(e.to_string()))?;
|
||
self.argmax_last(shape, data)
|
||
};
|
||
|
||
if next == self.eot {
|
||
break;
|
||
}
|
||
tokens.push(next);
|
||
}
|
||
|
||
Ok(self
|
||
.detok
|
||
.decode(&tokens[self.initial_tokens.len()..])
|
||
.trim()
|
||
.to_string())
|
||
}
|
||
|
||
/// Argmax over the final position's logits, masking special tokens so only
|
||
/// text (or `eot`) can win — greedy, no timestamps.
|
||
fn argmax_last(&self, shape: &[i64], data: &[f32]) -> i64 {
|
||
let vocab = shape[shape.len() - 1] as usize;
|
||
let seq = shape[shape.len() - 2] as usize;
|
||
let row = &data[(seq - 1) * vocab..seq * vocab];
|
||
let mut best = self.eot;
|
||
let mut best_val = f32::MIN;
|
||
for (id, &v) in row.iter().enumerate() {
|
||
let id = id as i64;
|
||
// Mask specials, but never mask eot so decode can always terminate
|
||
// regardless of how the config orders eot vs. the special floor.
|
||
if id >= self.special_floor && id != self.eot {
|
||
continue;
|
||
}
|
||
if v > best_val {
|
||
best_val = v;
|
||
best = id;
|
||
}
|
||
}
|
||
best
|
||
}
|
||
|
||
fn segment(&self, text: String, start_ms: u64, end_ms: u64) -> Option<TranscriptSegment> {
|
||
if text.is_empty() {
|
||
return None;
|
||
}
|
||
Some(TranscriptSegment {
|
||
id: self.next_id.fetch_add(1, Ordering::SeqCst),
|
||
start_ms,
|
||
end_ms: end_ms.max(start_ms),
|
||
speaker: "S1".to_string(), // diarization lands in Phase 4
|
||
text,
|
||
confidence: None, // greedy decode exposes no per-segment score yet
|
||
interim: false,
|
||
})
|
||
}
|
||
}
|
||
|
||
impl Transcriber for OnnxTranscriber {
|
||
/// `model` is the directory holding the ONNX artifacts (see `onnx_models`) —
|
||
/// the same artifacts for every EP. `backend` selects the accelerator:
|
||
/// `Npu` → OpenVINO EP; `Amd`/`Intel` → DirectML EP (the non-Vulkan GPU
|
||
/// path). Any other value is rejected so the dispatcher can fall back to
|
||
/// whisper.cpp rather than us guessing.
|
||
///
|
||
/// `language` (T8.7, FR-TRX-4): the exported ONNX model's
|
||
/// `forced_decoder_ids` bakes its language in at export time — this
|
||
/// engine has no per-inference language selection to apply, unlike
|
||
/// whisper.cpp's `FullParams::set_language`. A non-English, non-auto
|
||
/// request is logged (not silently dropped) so a user picking a language
|
||
/// on the NPU/DirectML path finds out it didn't take, rather than
|
||
/// getting a quietly-wrong transcript; `effective_language`/
|
||
/// `detected_language` fall back to the trait's English-only defaults.
|
||
fn load(model: &Path, backend: BackendId, language: Option<&str>) -> Result<Self, TrxError> {
|
||
if let Some(lang) = language {
|
||
if !lang.eq_ignore_ascii_case("auto") && !lang.eq_ignore_ascii_case("en") {
|
||
tracing::warn!(
|
||
"language '{lang}' requested but the NPU/DirectML engine's ONNX model is \
|
||
English-only (language is fixed at export time); ignoring the request"
|
||
);
|
||
}
|
||
}
|
||
// Pick the runtime bundle + encoder EP for the requested accelerator.
|
||
// error_on_failure makes a failed accelerator registration LOUD (Err)
|
||
// instead of a silent CPU fallback, so the dispatcher can cleanly drop
|
||
// to the whisper.cpp path (T3.4 spike lesson).
|
||
let (runtime_dll, encoder_ep): (std::path::PathBuf, ExecutionProviderDispatch) =
|
||
match backend {
|
||
BackendId::Npu => (
|
||
crate::paths::npu_runtime_dll(),
|
||
OpenVINOExecutionProvider::default()
|
||
.with_device_type("NPU")
|
||
.build()
|
||
.error_on_failure(),
|
||
),
|
||
// ponytail: DirectML device_id 0 = the default DX12 adapter.
|
||
// Map the DXGI adapter index here if a multi-GPU box ever needs
|
||
// to pin a specific AMD/Intel card.
|
||
BackendId::Amd | BackendId::Intel => (
|
||
crate::paths::directml_runtime_dll(),
|
||
DirectMLExecutionProvider::default()
|
||
.with_device_id(0)
|
||
.build()
|
||
.error_on_failure(),
|
||
),
|
||
other => {
|
||
return Err(TrxError::Load(format!(
|
||
"OnnxTranscriber serves NPU/AMD/Intel, got {other:?}"
|
||
)))
|
||
}
|
||
};
|
||
Self::ensure_runtime_env(runtime_dll);
|
||
|
||
let encoder_path = model.join("encoder_model.onnx");
|
||
let decoder_path = model.join("decoder_model.onnx");
|
||
let tokenizer_path = model.join("tokenizer.json");
|
||
let gen_config_path = model.join("generation_config.json");
|
||
|
||
let encoder = Session::builder()
|
||
.map_err(|e| TrxError::Load(e.to_string()))?
|
||
.with_execution_providers([encoder_ep])
|
||
.map_err(|e| TrxError::Load(e.to_string()))?
|
||
.commit_from_file(&encoder_path)
|
||
.map_err(|e| TrxError::Load(e.to_string()))?;
|
||
|
||
// Decoder on the default CPU EP (dynamic shapes; poor NPU fit).
|
||
let decoder = Session::builder()
|
||
.map_err(|e| TrxError::Load(e.to_string()))?
|
||
.commit_from_file(&decoder_path)
|
||
.map_err(|e| TrxError::Load(e.to_string()))?;
|
||
|
||
let enc_input = Self::find_input(&encoder, "input_features")?;
|
||
let enc_output = Self::find_output(&encoder, "last_hidden_state")
|
||
.or_else(|_| Self::find_output(&encoder, "hidden"))?;
|
||
let dec_ids = Self::find_input(&decoder, "input_ids")?;
|
||
let dec_hidden = Self::find_input(&decoder, "encoder_hidden")?;
|
||
let dec_logits = Self::find_output(&decoder, "logits")?;
|
||
|
||
let detok = Detok::from_tokenizer_json(&tokenizer_path)?;
|
||
let gen = GenConfig::from_file(&gen_config_path)?;
|
||
|
||
let mut initial_tokens = vec![gen.decoder_start];
|
||
initial_tokens.extend(gen.forced);
|
||
|
||
Ok(Self {
|
||
encoder: Mutex::new(encoder),
|
||
decoder: Mutex::new(decoder),
|
||
enc_input,
|
||
enc_output,
|
||
dec_ids,
|
||
dec_hidden,
|
||
dec_logits,
|
||
detok,
|
||
initial_tokens,
|
||
special_floor: gen.decoder_start,
|
||
eot: gen.eot,
|
||
next_id: AtomicU64::new(0),
|
||
})
|
||
}
|
||
|
||
fn transcribe_stream(&self, audio: AudioWindow, out: SegmentSink) -> Result<(), TrxError> {
|
||
let text = self.infer(&audio.samples)?;
|
||
let dur_ms = (audio.samples.len() as u64 * 1000) / super::mel::SAMPLE_RATE as u64;
|
||
if let Some(seg) = self.segment(text, audio.offset_ms, audio.offset_ms + dur_ms) {
|
||
let _ = out.send(seg);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn transcribe_file(&self, wav: &Path) -> Result<Vec<TranscriptSegment>, TrxError> {
|
||
let samples =
|
||
crate::audio::read_wav_mono_16k(wav).map_err(|e| TrxError::Load(e.to_string()))?;
|
||
// The encoder is fixed at one 30 s window, so chunk long audio and offset
|
||
// each chunk's segment in time.
|
||
let chunk = super::mel::N_SAMPLES;
|
||
let mut out = Vec::new();
|
||
for (i, part) in samples.chunks(chunk).enumerate() {
|
||
let text = self.infer(part)?;
|
||
let start_ms = (i * chunk) as u64 * 1000 / super::mel::SAMPLE_RATE as u64;
|
||
let dur_ms = (part.len() as u64 * 1000) / super::mel::SAMPLE_RATE as u64;
|
||
if let Some(seg) = self.segment(text, start_ms, start_ms + dur_ms) {
|
||
out.push(seg);
|
||
}
|
||
}
|
||
Ok(out)
|
||
}
|
||
}
|
||
|
||
/// Whisper's generation config: the decoder-start token, end token, and any
|
||
/// forced prefix tokens (e.g. `<|notimestamps|>`).
|
||
struct GenConfig {
|
||
decoder_start: i64,
|
||
eot: i64,
|
||
forced: Vec<i64>,
|
||
}
|
||
|
||
impl GenConfig {
|
||
fn from_file(path: &Path) -> Result<Self, TrxError> {
|
||
let text = std::fs::read_to_string(path)
|
||
.map_err(|e| TrxError::Load(format!("gen config: {e}")))?;
|
||
let v: serde_json::Value =
|
||
serde_json::from_str(&text).map_err(|e| TrxError::Load(format!("gen config: {e}")))?;
|
||
let decoder_start = v["decoder_start_token_id"]
|
||
.as_i64()
|
||
.ok_or_else(|| TrxError::Load("gen config missing decoder_start_token_id".into()))?;
|
||
// eos may be a scalar or an array; take the first if an array.
|
||
let eot = v["eos_token_id"]
|
||
.as_i64()
|
||
.or_else(|| {
|
||
v["eos_token_id"]
|
||
.as_array()
|
||
.and_then(|a| a.first()?.as_i64())
|
||
})
|
||
.unwrap_or(decoder_start - 1);
|
||
// forced_decoder_ids: [[pos, token], ...], applied in position order.
|
||
let mut forced_pairs: Vec<(i64, i64)> = v["forced_decoder_ids"]
|
||
.as_array()
|
||
.map(|arr| {
|
||
arr.iter()
|
||
.filter_map(|p| {
|
||
let p = p.as_array()?;
|
||
Some((p.first()?.as_i64()?, p.get(1)?.as_i64()?))
|
||
})
|
||
.collect()
|
||
})
|
||
.unwrap_or_default();
|
||
forced_pairs.sort_by_key(|(pos, _)| *pos);
|
||
let forced = forced_pairs.into_iter().map(|(_, tok)| tok).collect();
|
||
Ok(Self {
|
||
decoder_start,
|
||
eot,
|
||
forced,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// Minimal Whisper byte-level BPE **detokenizer** (ids → text). We only decode,
|
||
/// so this avoids pulling the heavyweight `tokenizers` crate (and its C deps):
|
||
/// read the vocab from `tokenizer.json`, then invert GPT-2's byte↔unicode map.
|
||
struct Detok {
|
||
id_to_token: HashMap<i64, String>,
|
||
byte_decoder: HashMap<char, u8>,
|
||
}
|
||
|
||
impl Detok {
|
||
fn from_tokenizer_json(path: &Path) -> Result<Self, TrxError> {
|
||
let text =
|
||
std::fs::read_to_string(path).map_err(|e| TrxError::Load(format!("tokenizer: {e}")))?;
|
||
let v: serde_json::Value =
|
||
serde_json::from_str(&text).map_err(|e| TrxError::Load(format!("tokenizer: {e}")))?;
|
||
let vocab = v["model"]["vocab"]
|
||
.as_object()
|
||
.ok_or_else(|| TrxError::Load("tokenizer.json missing model.vocab".into()))?;
|
||
let id_to_token = vocab
|
||
.iter()
|
||
.filter_map(|(tok, id)| Some((id.as_i64()?, tok.clone())))
|
||
.collect();
|
||
Ok(Self {
|
||
id_to_token,
|
||
byte_decoder: byte_decoder(),
|
||
})
|
||
}
|
||
|
||
fn decode(&self, ids: &[i64]) -> String {
|
||
let mut bytes = Vec::new();
|
||
for id in ids {
|
||
if let Some(tok) = self.id_to_token.get(id) {
|
||
for ch in tok.chars() {
|
||
if let Some(b) = self.byte_decoder.get(&ch) {
|
||
bytes.push(*b);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
String::from_utf8_lossy(&bytes).into_owned()
|
||
}
|
||
}
|
||
|
||
/// GPT-2 / Whisper `bytes_to_unicode`, inverted to map each printable proxy
|
||
/// char back to its raw byte.
|
||
fn byte_decoder() -> HashMap<char, u8> {
|
||
let mut bs: Vec<u32> = Vec::new();
|
||
for b in b'!'..=b'~' {
|
||
bs.push(b as u32);
|
||
}
|
||
for b in 0xA1u32..=0xAC {
|
||
bs.push(b);
|
||
}
|
||
for b in 0xAEu32..=0xFF {
|
||
bs.push(b);
|
||
}
|
||
let mut cs = bs.clone();
|
||
let mut n = 0u32;
|
||
for b in 0u32..256 {
|
||
if !bs.contains(&b) {
|
||
bs.push(b);
|
||
cs.push(256 + n);
|
||
n += 1;
|
||
}
|
||
}
|
||
bs.iter()
|
||
.zip(cs.iter())
|
||
.filter_map(|(&b, &c)| Some((char::from_u32(c)?, b as u8)))
|
||
.collect()
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn byte_decoder_maps_gpt2_space_proxy() {
|
||
// 'Ġ' (U+0120 = 256+32) is GPT-2's proxy for a space byte (0x20).
|
||
let d = byte_decoder();
|
||
assert_eq!(d.get(&'Ġ'), Some(&b' '));
|
||
// Printable ASCII maps to itself.
|
||
assert_eq!(d.get(&'A'), Some(&b'A'));
|
||
}
|
||
|
||
#[test]
|
||
fn byte_decoder_covers_all_256_bytes() {
|
||
let d = byte_decoder();
|
||
let mut seen = [false; 256];
|
||
for &b in d.values() {
|
||
seen[b as usize] = true;
|
||
}
|
||
assert!(
|
||
seen.iter().all(|&s| s),
|
||
"byte decoder must be a bijection over 0..256"
|
||
);
|
||
}
|
||
|
||
/// Real NPU inference, opt-in (needs the runtime + a downloaded model + a
|
||
/// WAV, none of which exist in CI). Run locally after the T3.4 spike setup:
|
||
/// ORT_DYLIB_PATH, PATH, WA_NPU_MODEL_DIR, WA_NPU_TEST_WAV[, WA_NPU_EXPECT]
|
||
#[test]
|
||
#[ignore = "requires NPU runtime + model + wav; run manually"]
|
||
fn npu_transcribes_speech() {
|
||
let Some(model_dir) = std::env::var_os("WA_NPU_MODEL_DIR") else {
|
||
eprintln!("skip: set WA_NPU_MODEL_DIR");
|
||
return;
|
||
};
|
||
let wav = std::env::var("WA_NPU_TEST_WAV").expect("set WA_NPU_TEST_WAV");
|
||
let t0 = std::time::Instant::now();
|
||
let t = OnnxTranscriber::load(Path::new(&model_dir), BackendId::Npu, None)
|
||
.expect("load NPU transcriber");
|
||
let load_ms = t0.elapsed().as_millis();
|
||
let t1 = std::time::Instant::now();
|
||
let segs = t.transcribe_file(Path::new(&wav)).expect("transcribe");
|
||
let infer_ms = t1.elapsed().as_millis();
|
||
let text = segs
|
||
.iter()
|
||
.map(|s| s.text.as_str())
|
||
.collect::<Vec<_>>()
|
||
.join(" ")
|
||
.to_lowercase();
|
||
eprintln!("[spike] backend=Npu load={load_ms}ms infer={infer_ms}ms text={text:?}");
|
||
assert!(!text.trim().is_empty(), "transcript was empty");
|
||
if let Ok(expect) = std::env::var("WA_NPU_EXPECT") {
|
||
assert!(
|
||
text.contains(&expect.to_lowercase()),
|
||
"transcript {text:?} missing expected {expect:?}"
|
||
);
|
||
}
|
||
}
|
||
|
||
/// Real DirectML GPU inference (P2), opt-in. Point `ORT_DYLIB_PATH` at a
|
||
/// DirectML-EP `onnxruntime.dll` (or stage `runtime\directml\`), reuse the
|
||
/// same ONNX model dir as the NPU spike, and run on the Intel/AMD GPU:
|
||
/// ORT_DYLIB_PATH=…directml\onnxruntime.dll WA_DML_MODEL_DIR=… \
|
||
/// WA_DML_TEST_WAV=… cargo test --release --features npu \
|
||
/// directml_transcribes -- --ignored --nocapture
|
||
#[test]
|
||
#[ignore = "requires a DirectML runtime + model + wav; run manually on a GPU"]
|
||
fn directml_transcribes_speech() {
|
||
let Some(model_dir) = std::env::var_os("WA_DML_MODEL_DIR") else {
|
||
eprintln!("skip: set WA_DML_MODEL_DIR");
|
||
return;
|
||
};
|
||
let wav = std::env::var("WA_DML_TEST_WAV").expect("set WA_DML_TEST_WAV");
|
||
// Intel and AMD both take the DirectML EP; either BackendId exercises it.
|
||
let backend = match std::env::var("WA_DML_BACKEND").as_deref() {
|
||
Ok("amd") => BackendId::Amd,
|
||
_ => BackendId::Intel,
|
||
};
|
||
let t0 = std::time::Instant::now();
|
||
let t = OnnxTranscriber::load(Path::new(&model_dir), backend, None).expect("load DirectML");
|
||
let load_ms = t0.elapsed().as_millis();
|
||
let t1 = std::time::Instant::now();
|
||
let segs = t.transcribe_file(Path::new(&wav)).expect("transcribe");
|
||
let infer_ms = t1.elapsed().as_millis();
|
||
let text = segs
|
||
.iter()
|
||
.map(|s| s.text.as_str())
|
||
.collect::<Vec<_>>()
|
||
.join(" ")
|
||
.to_lowercase();
|
||
eprintln!("[spike] backend={backend:?} load={load_ms}ms infer={infer_ms}ms text={text:?}");
|
||
assert!(!text.trim().is_empty(), "transcript was empty");
|
||
if let Ok(expect) = std::env::var("WA_DML_EXPECT") {
|
||
assert!(
|
||
text.contains(&expect.to_lowercase()),
|
||
"transcript {text:?} missing expected {expect:?}"
|
||
);
|
||
}
|
||
}
|
||
}
|