use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use serde::Deserialize;
#[derive(Clone, Copy, Debug)]
pub enum ModelChoice {
Tiny,
Small,
Medium,
Large,
}
impl ModelChoice {
pub fn cli_name(self) -> &'static str {
match self {
Self::Tiny => "tiny",
Self::Small => "small",
Self::Medium => "medium",
Self::Large => "large",
}
}
pub fn runtime_name(self) -> &'static str {
match self {
Self::Tiny => "tiny",
Self::Small => "small",
Self::Medium => "medium",
Self::Large => "large-v3",
}
}
pub fn repo_id(self) -> &'static str {
match self {
Self::Tiny => "Systran/faster-whisper-tiny",
Self::Small => "Systran/faster-whisper-small",
Self::Medium => "Systran/faster-whisper-medium",
Self::Large => "Systran/faster-whisper-large-v3",
}
}
pub fn cache_dir_name(self) -> &'static str {
match self {
Self::Tiny => "tiny",
Self::Small => "small",
Self::Medium => "medium",
Self::Large => "large",
}
}
}
#[derive(Debug, Deserialize)]
pub(crate) struct HfModelResponse {
#[serde(default)]
pub(crate) siblings: Vec<HfSibling>,
}
#[derive(Debug, Deserialize)]
pub(crate) struct HfSibling {
pub(crate) rfilename: String,
pub(crate) size: Option<u64>,
pub(crate) lfs: Option<HfLfs>,
}
#[derive(Debug, Deserialize)]
pub(crate) struct HfLfs {
pub(crate) size: Option<u64>,
}
impl HfSibling {
pub(crate) fn expected_size(&self) -> Option<u64> {
self.size
.or_else(|| self.lfs.as_ref().and_then(|lfs| lfs.size))
}
}
#[derive(Debug, Deserialize)]
pub struct ModelConfig {
pub model_type: Option<String>,
pub d_model: Option<u32>,
pub num_mel_bins: Option<u32>,
pub vocab_size: Option<u32>,
pub encoder_layers: Option<u32>,
pub decoder_layers: Option<u32>,
pub encoder_attention_heads: Option<u32>,
pub decoder_attention_heads: Option<u32>,
pub max_source_positions: Option<u32>,
pub max_target_positions: Option<u32>,
}
pub fn model_directory(choice: ModelChoice, models_root: Option<&Path>) -> Result<PathBuf> {
let models_dir = model_root_directory(models_root)?;
Ok(models_dir.join(choice.cache_dir_name()))
}
pub fn binary_directory() -> Result<PathBuf> {
let executable = std::env::current_exe().context("failed to resolve executable path")?;
executable
.parent()
.map(Path::to_path_buf)
.context("failed to resolve executable directory")
}
fn model_root_directory(models_root: Option<&Path>) -> Result<PathBuf> {
if let Some(models_root) = models_root {
return Ok(models_root.to_path_buf());
}
Ok(binary_directory()?.join("models"))
}
pub fn read_model_config(model_dir: &Path) -> Result<ModelConfig> {
let config_path = model_dir.join("config.json");
let config_contents = std::fs::read_to_string(&config_path)
.with_context(|| format!("failed to read `{}`", config_path.display()))?;
serde_json::from_str(&config_contents)
.with_context(|| format!("failed to parse `{}`", config_path.display()))
}
pub fn remove_model_with_artifacts(choice: ModelChoice, models_root: &Path) -> Result<usize> {
if !models_root.exists() {
return Ok(0);
}
let model_name = choice.cache_dir_name();
let mut removed = 0;
for entry in std::fs::read_dir(models_root)
.with_context(|| format!("failed to read `{}`", models_root.display()))?
{
let entry =
entry.with_context(|| format!("failed to inspect `{}`", models_root.display()))?;
let path = entry.path();
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
if !matches_model_artifact(file_name, model_name) {
continue;
}
remove_path(&path)?;
removed += 1;
}
Ok(removed)
}
pub fn remove_all_models(models_root: &Path) -> Result<bool> {
if !models_root.exists() {
return Ok(false);
}
std::fs::remove_dir_all(models_root)
.with_context(|| format!("failed to remove `{}`", models_root.display()))?;
Ok(true)
}
fn matches_model_artifact(file_name: &str, model_name: &str) -> bool {
file_name == model_name
|| file_name.starts_with(&format!("{model_name}."))
|| file_name.starts_with(&format!("{model_name}-"))
}
fn remove_path(path: &Path) -> Result<()> {
let metadata = std::fs::symlink_metadata(path)
.with_context(|| format!("failed to inspect `{}`", path.display()))?;
if metadata.is_dir() {
std::fs::remove_dir_all(path)
.with_context(|| format!("failed to remove directory `{}`", path.display()))?;
} else {
std::fs::remove_file(path)
.with_context(|| format!("failed to remove file `{}`", path.display()))?;
}
Ok(())
}