feat(dictation): add LOCAL_WHISPER_LANGUAGE for multilingual local transcription (#10634)

Co-authored-by: Anshu Saurabh <677936+anshusaurav@users.noreply.github.com>
This commit is contained in:
Anshu Saurabh
2026-08-01 02:04:06 +05:30
committed by GitHub
parent eea560973a
commit 37d5ab5b2c
+63 -1
View File
@@ -9,6 +9,9 @@
use crate::config::paths::Paths;
pub const LOCAL_WHISPER_MODEL_CONFIG_KEY: &str = "LOCAL_WHISPER_MODEL";
pub const LOCAL_WHISPER_LANGUAGE_CONFIG_KEY: &str = "LOCAL_WHISPER_LANGUAGE";
const ENGLISH_LANGUAGE_TOKEN: u32 = 50259;
const LANGUAGE_TOKEN_COUNT: u32 = 99;
use anyhow::{Context, Result};
use candle_core::{Device, IndexOp, Tensor};
use candle_nn::ops::log_softmax;
@@ -255,6 +258,8 @@ impl WhisperTranscriber {
let tokenizer = Self::load_tokenizer(model_path_ref, Some(bundled_tokenizer))?;
tracing::debug!("tokenizer loaded successfully");
let language_token = Self::resolve_language_token(&tokenizer);
Ok(Self {
model,
config,
@@ -263,11 +268,20 @@ impl WhisperTranscriber {
tokenizer,
eot_token: 50257,
no_timestamps_token: 50363,
language_token: 50259,
language_token,
max_initial_timestamp_index: 50,
})
}
fn resolve_language_token(tokenizer: &Tokenizer) -> u32 {
let configured = crate::config::Config::global()
.get(LOCAL_WHISPER_LANGUAGE_CONFIG_KEY, false)
.ok()
.and_then(|v| v.as_str().map(|s| s.to_string()));
language_token(tokenizer, configured.as_deref())
}
fn load_tokenizer(model_dir: &Path, bundled_tokenizer: Option<&str>) -> Result<Tokenizer> {
let tokenizer_path = model_dir
.parent()
@@ -302,6 +316,8 @@ impl WhisperTranscriber {
return Ok(String::new());
}
self.language_token = Self::resolve_language_token(&self.tokenizer);
let (mel_tensor, actual_content_frames) = self.prepare_audio_input(audio_data)?;
let (_, _, padded_frames) = mel_tensor.dims3()?;
@@ -718,6 +734,35 @@ impl WhisperTranscriber {
}
}
/// Resolve an ISO 639-1 language code to its Whisper language token.
///
/// Whisper tokenizers expose one `<|xx|>` token per supported language, so the code is looked up
/// in the tokenizer rather than mapped through a table that would drift from the model.
/// Unset or unrecognized codes fall back to English.
///
/// The result is bounded to the language block because the tokenizer also contains task and
/// control tokens such as `<|translate|>`, which would otherwise resolve here and produce an
/// invalid decoder prompt.
fn language_token(tokenizer: &Tokenizer, code: Option<&str>) -> u32 {
let Some(code) = code else {
return ENGLISH_LANGUAGE_TOKEN;
};
let code = code.trim().to_lowercase();
let languages = ENGLISH_LANGUAGE_TOKEN..ENGLISH_LANGUAGE_TOKEN + LANGUAGE_TOKEN_COUNT;
match tokenizer.token_to_id(&format!("<|{}|>", code)) {
Some(token) if languages.contains(&token) => token,
_ => {
tracing::warn!(
language = %code,
"unsupported LOCAL_WHISPER_LANGUAGE, transcribing as English"
);
ENGLISH_LANGUAGE_TOKEN
}
}
}
/// Remove repeated phrases from transcribed text.
///
/// Whisper models (especially smaller/quantized ones) tend to loop, producing output like
@@ -1103,6 +1148,23 @@ mod tests {
const TS: u32 = 50364; // A timestamp token for tests
#[test_case(None, ENGLISH_LANGUAGE_TOKEN ; "unset falls back to english")]
#[test_case(Some("en"), ENGLISH_LANGUAGE_TOKEN ; "english")]
#[test_case(Some("de"), 50261 ; "german")]
#[test_case(Some("ru"), 50263 ; "russian")]
#[test_case(Some("DE"), 50261 ; "uppercase code is normalized")]
#[test_case(Some("klingon"), ENGLISH_LANGUAGE_TOKEN ; "unsupported falls back to english")]
#[test_case(Some("su"), 50357 ; "last language in the block")]
#[test_case(Some("translate"), ENGLISH_LANGUAGE_TOKEN ; "task token is not a language")]
#[test_case(Some("notimestamps"), ENGLISH_LANGUAGE_TOKEN ; "control token is not a language")]
#[test_case(Some("startoftranscript"), ENGLISH_LANGUAGE_TOKEN ; "sot token is not a language")]
fn test_language_token(code: Option<&str>, expected: u32) {
let tokenizer =
Tokenizer::from_bytes(include_bytes!("whisper_data/tokens.json")).expect("tokenizer");
assert_eq!(language_token(&tokenizer, code), expected);
}
// detect_repetition_impl tests
// sample_begin=3 means tokens[0..3] are SOT, language, transcribe
// timestamp_begin=50364 means tokens >= 50364 are timestamps