Add test using tiny.en model behind a feature flag

This commit is contained in:
Jonathan Soo 2023-05-04 19:18:51 -04:00
parent 908e96f9da
commit 44e34ba301
2 changed files with 24 additions and 0 deletions

View file

@ -23,6 +23,7 @@ hound = "3.5.0"
[features]
simd = []
coreml = ["whisper-rs-sys/coreml"]
test-with-tiny-model = []
[package.metadata.docs.rs]
features = ["simd"]

View file

@ -431,3 +431,26 @@ impl Drop for WhisperContext {
// see https://github.com/ggerganov/whisper.cpp/issues/32#issuecomment-1272790388
unsafe impl Send for WhisperContext {}
unsafe impl Sync for WhisperContext {}
#[cfg(test)]
#[cfg(feature = "test-with-tiny-model")]
mod test_with_tiny_model {
use super::*;
const MODEL_PATH: &str = "./sys/whisper.cpp/models/ggml-tiny.en.bin";
// These tests expect that the tiny.en model has been downloaded
// using the script `sys/whisper.cpp/models/download-ggml-model.sh tiny.en`
#[test]
fn test_tokenize_round_trip() {
let ctx = WhisperContext::new(MODEL_PATH).expect("Download the ggml-tiny.en model using 'sys/whisper.cpp/models/download-ggml-model.sh tiny.en'");
let text_in = " And so my fellow Americans, ask not what your country can do for you, ask what you can do for your country.";
let tokens = ctx.tokenize(text_in, 1024).unwrap();
let text_out = tokens
.into_iter()
.map(|t| ctx.token_to_str(t).unwrap())
.collect::<Vec<_>>()
.join("");
assert_eq!(text_in, text_out);
}
}