|
| 1 | +//! On-disk cache for Whisper transcription output. |
| 2 | +//! |
| 3 | +//! Whisper is the slowest stage of `dpub convert --transcribe`. The |
| 4 | +//! output is deterministic given the audio bytes, the model bytes, |
| 5 | +//! the language code, and our serialisation schema — so we hash those, |
| 6 | +//! key a JSON file by the result, and skip re-running Whisper when |
| 7 | +//! the same combination has been seen before. |
| 8 | +//! |
| 9 | +//! Cache layout: |
| 10 | +//! - Directory: `~/.cache/dpub/transcripts/` (Unix), `%LOCALAPPDATA%\dpub\transcripts\` (Windows). |
| 11 | +//! - Filename: `<combined_hash>.json` where `combined_hash` derives from |
| 12 | +//! `(audio_sha256, model_sha256, language, schema_version)`. |
| 13 | +//! - Format: JSON envelope with diagnostic metadata + the |
| 14 | +//! `Vec<Segment>` payload. |
| 15 | +//! |
| 16 | +//! The cache is purely an optimisation: read failures fall back to a |
| 17 | +//! fresh transcription, write failures are logged and ignored. |
| 18 | +//! `DPUB_NO_TRANSCRIPT_CACHE=1` disables both reads and writes. |
| 19 | +
|
| 20 | +use std::fs; |
| 21 | +use std::io::{Read, Write}; |
| 22 | +use std::path::{Path, PathBuf}; |
| 23 | + |
| 24 | +use serde::{Deserialize, Serialize}; |
| 25 | +use sha2::{Digest, Sha256}; |
| 26 | + |
| 27 | +use dpub_whisper::{Segment, TranscribeOptions, Transcriber}; |
| 28 | + |
| 29 | +/// Bumped whenever the on-disk JSON shape changes. Old cache files |
| 30 | +/// hash to a different key after a bump and will simply be ignored |
| 31 | +/// (and overwritten on the next miss). No deletion needed. |
| 32 | +const SCHEMA_VERSION: u32 = 1; |
| 33 | + |
| 34 | +/// Disk cache wrapper around `dpub_whisper::Transcriber`. Keeps the |
| 35 | +/// model loaded and its hash memoised across all calls in one run. |
| 36 | +pub(crate) struct CachedTranscriber { |
| 37 | + inner: Transcriber, |
| 38 | + model_sha: String, |
| 39 | + language: String, |
| 40 | + cache_dir: PathBuf, |
| 41 | + cache_enabled: bool, |
| 42 | +} |
| 43 | + |
| 44 | +impl CachedTranscriber { |
| 45 | + pub(crate) fn new(opts: &TranscribeOptions) -> crate::Result<Self> { |
| 46 | + let inner = Transcriber::new(opts)?; |
| 47 | + let model_sha = hash_file(&opts.model_path).unwrap_or_else(|e| { |
| 48 | + // Hashing failure isn't fatal — it just disables the |
| 49 | + // cache for this run. Log it so the user knows why they |
| 50 | + // didn't get a speedup. |
| 51 | + tracing::warn!( |
| 52 | + "transcript cache: model hash failed ({e}); cache disabled this run" |
| 53 | + ); |
| 54 | + String::new() |
| 55 | + }); |
| 56 | + let cache_enabled = !model_sha.is_empty() |
| 57 | + && std::env::var_os("DPUB_NO_TRANSCRIPT_CACHE").is_none(); |
| 58 | + let cache_dir = transcripts_cache_dir(); |
| 59 | + if cache_enabled { |
| 60 | + // Create the dir lazily; ignore failures (we'll log on first write). |
| 61 | + let _ = fs::create_dir_all(&cache_dir); |
| 62 | + } |
| 63 | + Ok(Self { |
| 64 | + inner, |
| 65 | + model_sha, |
| 66 | + language: opts.language.clone(), |
| 67 | + cache_dir, |
| 68 | + cache_enabled, |
| 69 | + }) |
| 70 | + } |
| 71 | + |
| 72 | + pub(crate) fn transcribe(&self, audio_path: &Path) -> crate::Result<Vec<Segment>> { |
| 73 | + if !self.cache_enabled { |
| 74 | + return Ok(self.inner.transcribe(audio_path)?); |
| 75 | + } |
| 76 | + let audio_sha = match hash_file(audio_path) { |
| 77 | + Ok(s) => s, |
| 78 | + Err(e) => { |
| 79 | + tracing::warn!( |
| 80 | + "transcript cache: audio hash failed for {} ({e}); transcribing without cache", |
| 81 | + audio_path.display() |
| 82 | + ); |
| 83 | + return Ok(self.inner.transcribe(audio_path)?); |
| 84 | + } |
| 85 | + }; |
| 86 | + let key = combined_key(&audio_sha, &self.model_sha, &self.language); |
| 87 | + let cache_path = self.cache_dir.join(format!("{key}.json")); |
| 88 | + |
| 89 | + if let Some(segments) = read_cached(&cache_path) { |
| 90 | + tracing::info!( |
| 91 | + "transcript cache: hit for {} ({} segments)", |
| 92 | + audio_path.display(), |
| 93 | + segments.len() |
| 94 | + ); |
| 95 | + return Ok(segments); |
| 96 | + } |
| 97 | + |
| 98 | + let segments = self.inner.transcribe(audio_path)?; |
| 99 | + let envelope = Envelope { |
| 100 | + schema_version: SCHEMA_VERSION, |
| 101 | + audio_sha256: audio_sha, |
| 102 | + model_sha256: self.model_sha.clone(), |
| 103 | + language: self.language.clone(), |
| 104 | + dpub_whisper_version: env!("CARGO_PKG_VERSION").to_owned(), |
| 105 | + segments: segments.clone(), |
| 106 | + }; |
| 107 | + if let Err(e) = write_cached(&cache_path, &envelope) { |
| 108 | + tracing::warn!( |
| 109 | + "transcript cache: write failed for {} ({e}); transcript will be re-computed next time", |
| 110 | + cache_path.display() |
| 111 | + ); |
| 112 | + } else { |
| 113 | + tracing::debug!( |
| 114 | + "transcript cache: stored {} ({} segments)", |
| 115 | + cache_path.display(), |
| 116 | + envelope.segments.len() |
| 117 | + ); |
| 118 | + } |
| 119 | + Ok(segments) |
| 120 | + } |
| 121 | +} |
| 122 | + |
| 123 | +/// JSON envelope written to disk. The metadata fields duplicate the |
| 124 | +/// inputs that already feed into the cache key — they're for `jq` |
| 125 | +/// debugging, not lookup. |
| 126 | +#[derive(Debug, Serialize, Deserialize)] |
| 127 | +struct Envelope { |
| 128 | + schema_version: u32, |
| 129 | + audio_sha256: String, |
| 130 | + model_sha256: String, |
| 131 | + language: String, |
| 132 | + dpub_whisper_version: String, |
| 133 | + segments: Vec<Segment>, |
| 134 | +} |
| 135 | + |
| 136 | +/// Look up the cache file. Returns `Some(segments)` on a clean hit. |
| 137 | +/// Any error (missing file, corrupt JSON, schema mismatch) yields |
| 138 | +/// `None`; missing files are silent, real errors log a warning. |
| 139 | +fn read_cached(path: &Path) -> Option<Vec<Segment>> { |
| 140 | + let bytes = match fs::read(path) { |
| 141 | + Ok(b) => b, |
| 142 | + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return None, |
| 143 | + Err(e) => { |
| 144 | + tracing::warn!("transcript cache: read failed for {}: {e}", path.display()); |
| 145 | + return None; |
| 146 | + } |
| 147 | + }; |
| 148 | + let env: Envelope = match serde_json::from_slice(&bytes) { |
| 149 | + Ok(e) => e, |
| 150 | + Err(e) => { |
| 151 | + tracing::warn!( |
| 152 | + "transcript cache: ignoring malformed entry {}: {e}", |
| 153 | + path.display() |
| 154 | + ); |
| 155 | + return None; |
| 156 | + } |
| 157 | + }; |
| 158 | + if env.schema_version != SCHEMA_VERSION { |
| 159 | + return None; |
| 160 | + } |
| 161 | + Some(env.segments) |
| 162 | +} |
| 163 | + |
| 164 | +/// Atomically write the cache entry (`.partial` then rename). Same |
| 165 | +/// pattern as the model downloader in `dpub-cli/src/setup.rs`. |
| 166 | +fn write_cached(path: &Path, envelope: &Envelope) -> std::io::Result<()> { |
| 167 | + if let Some(parent) = path.parent() { |
| 168 | + fs::create_dir_all(parent)?; |
| 169 | + } |
| 170 | + let partial = path.with_extension("json.partial"); |
| 171 | + let json = serde_json::to_vec(envelope).map_err(std::io::Error::other)?; |
| 172 | + { |
| 173 | + let mut f = fs::File::create(&partial)?; |
| 174 | + f.write_all(&json)?; |
| 175 | + f.sync_data()?; |
| 176 | + } |
| 177 | + fs::rename(&partial, path)?; |
| 178 | + Ok(()) |
| 179 | +} |
| 180 | + |
| 181 | +/// Stream-hash a file's bytes with SHA-256. Mirrors the helper used |
| 182 | +/// for `dpub setup --whisper-model …` model verification but lives |
| 183 | +/// here to avoid a cross-crate dependency for ~15 lines. |
| 184 | +fn hash_file(path: &Path) -> std::io::Result<String> { |
| 185 | + let mut file = fs::File::open(path)?; |
| 186 | + let mut hasher = Sha256::new(); |
| 187 | + let mut buf = vec![0u8; 64 * 1024]; |
| 188 | + loop { |
| 189 | + let n = file.read(&mut buf)?; |
| 190 | + if n == 0 { |
| 191 | + break; |
| 192 | + } |
| 193 | + hasher.update(&buf[..n]); |
| 194 | + } |
| 195 | + Ok(hex(hasher.finalize().as_slice())) |
| 196 | +} |
| 197 | + |
| 198 | +fn hex(bytes: &[u8]) -> String { |
| 199 | + use std::fmt::Write; |
| 200 | + let mut s = String::with_capacity(bytes.len() * 2); |
| 201 | + for b in bytes { |
| 202 | + let _ = write!(&mut s, "{b:02x}"); |
| 203 | + } |
| 204 | + s |
| 205 | +} |
| 206 | + |
| 207 | +/// Combined cache key: `sha256(audio_sha || model_sha || lang || schema_version)`, |
| 208 | +/// truncated to 32 hex chars. Truncation is fine: SHA-256 has no |
| 209 | +/// adversary here, only the normal birthday-bound risk, which at 128 |
| 210 | +/// bits of entropy is ~2^64 inputs before a collision is even |
| 211 | +/// plausible. Real-world cache will have a few thousand entries max. |
| 212 | +fn combined_key(audio_sha: &str, model_sha: &str, language: &str) -> String { |
| 213 | + let mut hasher = Sha256::new(); |
| 214 | + hasher.update(audio_sha.as_bytes()); |
| 215 | + hasher.update(b"\0"); |
| 216 | + hasher.update(model_sha.as_bytes()); |
| 217 | + hasher.update(b"\0"); |
| 218 | + hasher.update(language.as_bytes()); |
| 219 | + hasher.update(b"\0"); |
| 220 | + hasher.update(SCHEMA_VERSION.to_le_bytes()); |
| 221 | + let hex = hex(hasher.finalize().as_slice()); |
| 222 | + hex[..32].to_owned() |
| 223 | +} |
| 224 | + |
| 225 | +/// Return the platform-appropriate transcripts cache directory. |
| 226 | +/// Mirrors the layout of `~/.cache/dpub/models/` in `dpub-cli/setup.rs`. |
| 227 | +fn transcripts_cache_dir() -> PathBuf { |
| 228 | + if cfg!(target_os = "windows") { |
| 229 | + let base = std::env::var_os("LOCALAPPDATA") |
| 230 | + .map_or_else(|| PathBuf::from("."), PathBuf::from); |
| 231 | + base.join("dpub").join("transcripts") |
| 232 | + } else { |
| 233 | + let home = std::env::var_os("HOME") |
| 234 | + .map_or_else(|| PathBuf::from("."), PathBuf::from); |
| 235 | + home.join(".cache").join("dpub").join("transcripts") |
| 236 | + } |
| 237 | +} |
| 238 | + |
| 239 | +#[cfg(test)] |
| 240 | +mod tests { |
| 241 | + use super::*; |
| 242 | + use dpub_whisper::Word; |
| 243 | + |
| 244 | + fn sample_segments() -> Vec<Segment> { |
| 245 | + vec![Segment { |
| 246 | + start_seconds: 0.0, |
| 247 | + end_seconds: 1.5, |
| 248 | + text: "Hello world.".into(), |
| 249 | + words: vec![ |
| 250 | + Word { |
| 251 | + start_seconds: 0.0, |
| 252 | + end_seconds: 0.5, |
| 253 | + text: "Hello".into(), |
| 254 | + }, |
| 255 | + Word { |
| 256 | + start_seconds: 0.5, |
| 257 | + end_seconds: 1.5, |
| 258 | + text: "world.".into(), |
| 259 | + }, |
| 260 | + ], |
| 261 | + }] |
| 262 | + } |
| 263 | + |
| 264 | + #[test] |
| 265 | + fn round_trip_envelope() { |
| 266 | + let dir = tempfile::tempdir().unwrap(); |
| 267 | + let path = dir.path().join("entry.json"); |
| 268 | + let env = Envelope { |
| 269 | + schema_version: SCHEMA_VERSION, |
| 270 | + audio_sha256: "aaaa".into(), |
| 271 | + model_sha256: "bbbb".into(), |
| 272 | + language: "nl".into(), |
| 273 | + dpub_whisper_version: "0.6.0".into(), |
| 274 | + segments: sample_segments(), |
| 275 | + }; |
| 276 | + write_cached(&path, &env).unwrap(); |
| 277 | + let got = read_cached(&path).expect("hit"); |
| 278 | + assert_eq!(got, env.segments); |
| 279 | + } |
| 280 | + |
| 281 | + #[test] |
| 282 | + fn missing_file_is_silent_miss() { |
| 283 | + let dir = tempfile::tempdir().unwrap(); |
| 284 | + let path = dir.path().join("nope.json"); |
| 285 | + assert!(read_cached(&path).is_none()); |
| 286 | + } |
| 287 | + |
| 288 | + #[test] |
| 289 | + fn corrupt_file_is_warning_miss() { |
| 290 | + let dir = tempfile::tempdir().unwrap(); |
| 291 | + let path = dir.path().join("bad.json"); |
| 292 | + fs::write(&path, b"not json").unwrap(); |
| 293 | + assert!(read_cached(&path).is_none()); |
| 294 | + } |
| 295 | + |
| 296 | + #[test] |
| 297 | + fn schema_mismatch_treated_as_miss() { |
| 298 | + let dir = tempfile::tempdir().unwrap(); |
| 299 | + let path = dir.path().join("v0.json"); |
| 300 | + let json = serde_json::json!({ |
| 301 | + "schema_version": SCHEMA_VERSION + 99, |
| 302 | + "audio_sha256": "a", |
| 303 | + "model_sha256": "b", |
| 304 | + "language": "nl", |
| 305 | + "dpub_whisper_version": "0.6.0", |
| 306 | + "segments": [], |
| 307 | + }); |
| 308 | + fs::write(&path, serde_json::to_vec(&json).unwrap()).unwrap(); |
| 309 | + assert!(read_cached(&path).is_none()); |
| 310 | + } |
| 311 | + |
| 312 | + #[test] |
| 313 | + fn hash_file_is_deterministic() { |
| 314 | + let dir = tempfile::tempdir().unwrap(); |
| 315 | + let p = dir.path().join("a.bin"); |
| 316 | + fs::write(&p, b"hello world").unwrap(); |
| 317 | + assert_eq!(hash_file(&p).unwrap(), hash_file(&p).unwrap()); |
| 318 | + } |
| 319 | + |
| 320 | + #[test] |
| 321 | + fn hash_file_distinguishes_inputs() { |
| 322 | + let dir = tempfile::tempdir().unwrap(); |
| 323 | + let a = dir.path().join("a.bin"); |
| 324 | + let b = dir.path().join("b.bin"); |
| 325 | + fs::write(&a, b"hello").unwrap(); |
| 326 | + fs::write(&b, b"world").unwrap(); |
| 327 | + assert_ne!(hash_file(&a).unwrap(), hash_file(&b).unwrap()); |
| 328 | + } |
| 329 | + |
| 330 | + #[test] |
| 331 | + fn combined_key_changes_when_any_input_changes() { |
| 332 | + let base = combined_key("aaaa", "bbbb", "nl"); |
| 333 | + assert_ne!(base, combined_key("zzzz", "bbbb", "nl")); |
| 334 | + assert_ne!(base, combined_key("aaaa", "zzzz", "nl")); |
| 335 | + assert_ne!(base, combined_key("aaaa", "bbbb", "en")); |
| 336 | + } |
| 337 | + |
| 338 | + #[test] |
| 339 | + fn cache_dir_ends_in_transcripts() { |
| 340 | + let dir = transcripts_cache_dir(); |
| 341 | + assert_eq!(dir.file_name().unwrap(), "transcripts"); |
| 342 | + let parent_name = dir.parent().unwrap().file_name().unwrap(); |
| 343 | + assert_eq!(parent_name, "dpub"); |
| 344 | + } |
| 345 | +} |
0 commit comments