rlx_whisper/
bench_fixture.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Downloaded bench clips with paired reference transcripts (`just fetch-whisper-bench`).
17
18use anyhow::{Context, Result, bail, ensure};
19use std::path::{Path, PathBuf};
20
21/// Default cache dir: `<repo>/.cache/whisper-bench`.
22pub fn bench_cache_dir() -> PathBuf {
23    std::env::var("RLX_WHISPER_BENCH_DIR")
24        .map(PathBuf::from)
25        .unwrap_or_else(|_| default_bench_cache_dir())
26}
27
28fn default_bench_cache_dir() -> PathBuf {
29    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../.cache/whisper-bench")
30}
31
32/// OpenAI Whisper test clip (16 kHz mono). Override with `RLX_WHISPER_WAV`.
33pub fn jfk_wav_path() -> PathBuf {
34    std::env::var("RLX_WHISPER_WAV")
35        .map(PathBuf::from)
36        .unwrap_or_else(|_| bench_cache_dir().join("jfk_16k.wav"))
37}
38
39/// Greedy reference transcript for [`jfk_wav_path`]. Override with `RLX_WHISPER_REFERENCE`.
40pub fn jfk_reference_path() -> PathBuf {
41    std::env::var("RLX_WHISPER_REFERENCE")
42        .map(PathBuf::from)
43        .unwrap_or_else(|_| bench_cache_dir().join("jfk.reference.txt"))
44}
45
46/// Embedded fallback when the reference file has not been fetched yet.
47pub const JFK_REFERENCE: &str = " And so my fellow Americans ask not what your country can do for you ask what you can do for your country.";
48
49/// Load the JFK reference string (file if present, else embedded constant).
50pub fn load_jfk_reference() -> Result<String> {
51    let path = jfk_reference_path();
52    if path.is_file() {
53        std::fs::read_to_string(&path)
54            .with_context(|| format!("read reference {path:?}"))
55            .map(|s| s.trim_end_matches(['\r', '\n']).to_string())
56    } else {
57        Ok(JFK_REFERENCE.to_string())
58    }
59}
60
61/// Ensure wav + reference exist; useful from tests before running ASR.
62pub fn ensure_jfk_fixture() -> Result<(PathBuf, String)> {
63    let wav = jfk_wav_path();
64    ensure!(
65        wav.is_file(),
66        "missing JFK wav at {wav:?}; run `just fetch-whisper-bench`"
67    );
68    let reference = load_jfk_reference()?;
69    Ok((wav, reference))
70}
71
72/// Collapse whitespace for stable one-to-one comparison (Whisper may emit a leading space).
73pub fn normalize_transcript(text: &str) -> String {
74    text.split_whitespace().collect::<Vec<_>>().join(" ")
75}
76
77/// Exact match after whitespace normalization.
78pub fn transcripts_match(got: &str, want: &str) -> bool {
79    normalize_transcript(got) == normalize_transcript(want)
80}
81
82/// Assert greedy decode matches the paired reference transcript.
83pub fn assert_transcript_matches_reference(got: &str, want: &str) {
84    assert!(
85        transcripts_match(got, want),
86        "transcript mismatch.\n  got:  {got:?}\n  want: {want:?}\n  \
87         (normalized got:  {:?})\n  (normalized want: {:?})",
88        normalize_transcript(got),
89        normalize_transcript(want),
90    );
91}
92
93/// Read manifest path if present.
94pub fn manifest_path() -> PathBuf {
95    bench_cache_dir().join("manifest.json")
96}
97
98/// Return `(wav, reference)` for clip id `jfk` only (extensible via manifest later).
99pub fn clip_paths(id: &str) -> Result<(PathBuf, PathBuf)> {
100    if id == "jfk" {
101        return Ok((jfk_wav_path(), jfk_reference_path()));
102    }
103    bail!("unknown whisper bench clip {id:?}")
104}
105
106/// Load reference text for a bench clip.
107pub fn load_reference_for(id: &str) -> Result<String> {
108    match id {
109        "jfk" => load_jfk_reference(),
110        other => bail!("unknown whisper bench clip {other:?}"),
111    }
112}
113
114/// Check fixture files on disk (for preflight / CLI).
115pub fn fixture_status(dir: &Path) -> Result<()> {
116    let wav = dir.join("jfk_16k.wav");
117    let reference = dir.join("jfk.reference.txt");
118    ensure!(wav.is_file(), "missing {wav:?}");
119    ensure!(reference.is_file(), "missing {reference:?}");
120    let _ = load_jfk_reference()?;
121    Ok(())
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn normalize_collapses_whitespace() {
130        assert_eq!(
131            normalize_transcript("  And  so my fellow Americans  "),
132            "And so my fellow Americans"
133        );
134    }
135
136    #[test]
137    fn embedded_jfk_reference_matches_openai_whisper() {
138        assert!(JFK_REFERENCE.starts_with(' '));
139        assert!(transcripts_match(
140            JFK_REFERENCE,
141            "And so my fellow Americans ask not what your country can do for you ask what you can do for your country."
142        ));
143    }
144}
rlx_whisper/bench_fixture.rs

rlx_whisper/
bench_fixture.rs
