ferrum_cli/commands/
transcribe.rs

1//! Transcribe command - Speech-to-text using Whisper models
2
3use crate::config::CliConfig;
4use candle_core::{DType, Device as CandleDevice};
5use clap::Args;
6use colored::Colorize;
7use ferrum_models::source::{ModelFormat, ResolvedModelSource};
8use ferrum_models::{ConfigManager, HfDownloader, WhisperModelExecutor};
9use ferrum_types::Result;
10use std::path::PathBuf;
11
12/// Transcribe audio files using Whisper models
13#[derive(Args, Debug)]
14pub struct TranscribeCommand {
15    /// Whisper model name (e.g., whisper-tiny, openai/whisper-base)
16    #[arg(required = true)]
17    pub model: String,
18
19    /// Audio file path (WAV format)
20    #[arg(required = true)]
21    pub audio: String,
22
23    /// Language hint (e.g., en, zh, ja)
24    #[arg(short, long)]
25    pub language: Option<String>,
26
27    /// Backend: auto, cpu, metal (default: auto)
28    #[arg(short, long, default_value = "auto")]
29    pub backend: String,
30}
31
32pub async fn execute(cmd: TranscribeCommand, config: CliConfig) -> Result<()> {
33    let model_id = resolve_whisper_alias(&cmd.model);
34    let cache_dir = get_hf_cache_dir(&config);
35
36    eprintln!("{} {}", "Model:".dimmed(), model_id.cyan());
37
38    // Find or download model
39    let source = match find_cached_model(&cache_dir, &model_id) {
40        Some(source) => source,
41        None => {
42            eprintln!(
43                "{} Model '{}' not found locally, downloading...",
44                "📥".cyan(),
45                model_id
46            );
47            let token = std::env::var("HF_TOKEN")
48                .or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
49                .ok();
50            let downloader = HfDownloader::new(cache_dir, token)?;
51            let snapshot_path = downloader.download(&model_id, None).await?;
52            let format = detect_format(&snapshot_path);
53            if format == ModelFormat::Unknown {
54                return Err(ferrum_types::FerrumError::model(
55                    "Downloaded model has unknown format",
56                ));
57            }
58            ResolvedModelSource {
59                original: model_id.clone(),
60                local_path: snapshot_path,
61                format,
62                from_cache: false,
63            }
64        }
65    };
66
67    // Verify architecture
68    let mut config_manager = ConfigManager::new();
69    let model_def = config_manager.load_from_path(&source.local_path).await?;
70    if model_def.architecture != ferrum_models::Architecture::Whisper {
71        return Err(ferrum_types::FerrumError::model(format!(
72            "'{}' is not a Whisper model (detected: {:?})",
73            model_id, model_def.architecture
74        )));
75    }
76
77    let candle_device = select_candle_device(&cmd.backend)?;
78    eprintln!("{} {:?}", "Device:".dimmed(), &candle_device);
79    eprintln!("{}", "Loading Whisper model...".dimmed());
80    let executor = WhisperModelExecutor::from_path(
81        &source.local_path.to_string_lossy(),
82        candle_device,
83        DType::F32,
84    )?;
85    eprintln!("{}", "Model loaded.".green());
86
87    // Transcribe
88    eprintln!("{} {}", "Audio:".dimmed(), cmd.audio.cyan());
89    let start = std::time::Instant::now();
90    let text = executor.transcribe_file(&cmd.audio, cmd.language.as_deref())?;
91    let elapsed = start.elapsed();
92
93    println!("{}", text);
94    eprintln!("\n{} {:.2}s", "Time:".dimmed(), elapsed.as_secs_f64());
95
96    Ok(())
97}
98
99fn resolve_whisper_alias(name: &str) -> String {
100    match name.to_lowercase().as_str() {
101        "whisper-tiny" | "whisper:tiny" => "openai/whisper-tiny".to_string(),
102        "whisper-base" | "whisper:base" => "openai/whisper-base".to_string(),
103        "whisper-small" | "whisper:small" => "openai/whisper-small".to_string(),
104        "whisper-medium" | "whisper:medium" => "openai/whisper-medium".to_string(),
105        "whisper-large-v3" | "whisper:large-v3" => "openai/whisper-large-v3".to_string(),
106        "whisper-turbo" | "whisper:turbo" | "whisper-large-v3-turbo" => {
107            "openai/whisper-large-v3-turbo".to_string()
108        }
109        _ => name.to_string(),
110    }
111}
112
113fn get_hf_cache_dir(config: &CliConfig) -> PathBuf {
114    if let Ok(hf_home) = std::env::var("HF_HOME") {
115        return PathBuf::from(hf_home);
116    }
117    let configured = shellexpand::tilde(&config.models.download.hf_cache_dir).to_string();
118    PathBuf::from(configured)
119}
120
121fn find_cached_model(cache_dir: &PathBuf, model_id: &str) -> Option<ResolvedModelSource> {
122    let hub_dir = cache_dir.join("hub");
123    let model_dir_name = format!("models--{}", model_id.replace("/", "--"));
124    let model_dir = hub_dir.join(&model_dir_name);
125
126    if model_dir.exists() {
127        let snapshots_dir = model_dir.join("snapshots");
128        if snapshots_dir.exists() {
129            // Try refs/main first
130            let ref_main = model_dir.join("refs").join("main");
131            if let Ok(rev) = std::fs::read_to_string(&ref_main) {
132                let rev = rev.trim();
133                if !rev.is_empty() {
134                    let snapshot = snapshots_dir.join(rev);
135                    if snapshot.exists() {
136                        let format = detect_format(&snapshot);
137                        if format != ModelFormat::Unknown {
138                            return Some(ResolvedModelSource {
139                                original: model_id.to_string(),
140                                local_path: snapshot,
141                                format,
142                                from_cache: true,
143                            });
144                        }
145                    }
146                }
147            }
148            // Fallback: first snapshot
149            if let Ok(entries) = std::fs::read_dir(&snapshots_dir) {
150                for entry in entries.flatten() {
151                    let path = entry.path();
152                    if path.is_dir() && path.join("config.json").exists() {
153                        let format = detect_format(&path);
154                        if format != ModelFormat::Unknown {
155                            return Some(ResolvedModelSource {
156                                original: model_id.to_string(),
157                                local_path: path,
158                                format,
159                                from_cache: true,
160                            });
161                        }
162                    }
163                }
164            }
165        }
166    }
167
168    None
169}
170
171pub fn select_candle_device(backend: &str) -> ferrum_types::Result<CandleDevice> {
172    match backend.to_lowercase().as_str() {
173        "cpu" => Ok(CandleDevice::Cpu),
174        "metal" => {
175            #[cfg(all(target_os = "macos", feature = "metal"))]
176            {
177                return CandleDevice::new_metal(0)
178                    .map_err(|error| ferrum_types::FerrumError::device(error.to_string()));
179            }
180            #[allow(unreachable_code)]
181            {
182                Err(ferrum_types::FerrumError::unsupported(
183                    "Metal transcription requires a Metal-enabled Ferrum build",
184                ))
185            }
186        }
187        "cuda" => {
188            #[cfg(feature = "candle-cuda-compat")]
189            {
190                return CandleDevice::new_cuda(0)
191                    .map_err(|error| ferrum_types::FerrumError::device(error.to_string()));
192            }
193            #[allow(unreachable_code)]
194            {
195                Err(ferrum_types::FerrumError::unsupported(
196                    "CUDA transcription requires the candle-cuda-compat feature",
197                ))
198            }
199        }
200        "auto" | _ => {
201            #[cfg(feature = "candle-cuda-compat")]
202            {
203                if let Ok(d) = CandleDevice::new_cuda(0) {
204                    return Ok(d);
205                }
206            }
207            #[cfg(all(target_os = "macos", feature = "metal"))]
208            {
209                if let Ok(device) = CandleDevice::new_metal(0) {
210                    return Ok(device);
211                }
212            }
213            #[allow(unreachable_code)]
214            Ok(CandleDevice::Cpu)
215        }
216    }
217}
218
219fn detect_format(path: &PathBuf) -> ModelFormat {
220    if path.join("model.safetensors").exists() {
221        ModelFormat::SafeTensors
222    } else if std::fs::read_dir(path)
223        .map(|d| {
224            d.filter_map(|e| e.ok())
225                .any(|e| e.path().extension().is_some_and(|ext| ext == "safetensors"))
226        })
227        .unwrap_or(false)
228    {
229        ModelFormat::SafeTensors
230    } else if path.join("pytorch_model.bin").exists() {
231        ModelFormat::PyTorchBin
232    } else {
233        ModelFormat::Unknown
234    }
235}
ferrum_cli/commands/transcribe.rs

ferrum_cli/commands/
transcribe.rs
