whisper_cli 0.1.5

A command line interface for whisper-rs
Documentation
use crate::utils::download_file;
use dirs::cache_dir;
use std::{fmt::Display, fs, path::PathBuf};

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, clap::ValueEnum)]
pub enum Size {
    #[clap(name = "tiny.en")]
    TinyEnglish,
    #[clap(name = "tiny")]
    Tiny,
    #[clap(name = "base.en")]
    BaseEnglish,
    #[clap(name = "base")]
    Base,
    #[clap(name = "small.en")]
    SmallEnglish,
    #[clap(name = "small")]
    Small,
    #[clap(name = "medium.en")]
    MediumEnglish,
    #[clap(name = "medium")]
    Medium,
    #[clap(name = "large")]
    Large,
    #[clap(name = "large-v1")]
    LargeV1,
}

impl Size {
    pub fn get_path(self) -> PathBuf {
        let mut path = cache_dir().expect("Could not find cache directory");
        path.push("whisper");
        path.push("models");
        path.push(format!("{self}.bin"));

        path
    }

    pub const fn is_english_only(self) -> bool {
        matches!(
            self,
            Self::TinyEnglish | Self::BaseEnglish | Self::SmallEnglish | Self::MediumEnglish
        )
    }
}

impl Display for Size {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let key = match self {
            Self::TinyEnglish => "tiny.en",
            Self::Tiny => "tiny",
            Self::BaseEnglish => "base.en",
            Self::Base => "base",
            Self::SmallEnglish => "small.en",
            Self::Small => "small",
            Self::MediumEnglish => "medium.en",
            Self::Medium => "medium",
            Self::Large => "large",
            Self::LargeV1 => "large-v1",
        };

        write!(f, "{key}")
    }
}

pub struct Model {
    size: Size,
}

impl Model {
    pub const fn new(size: Size) -> Self {
        Self { size }
    }

    pub fn get_path(&self) -> PathBuf {
        self.size.get_path()
    }

    pub async fn download(&self) {
        let path = self.get_path();
        if path.exists() {
            return;
        }

        let cache_dir = path.parent().expect("Failed to get cache dir");
        if !cache_dir.exists() {
            fs::create_dir_all(cache_dir).expect("Failed to create cache dir.");
        }

        download_file(
            &format!(
                "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-{}.bin",
                self.size
            ),
            path.to_str().unwrap(),
        )
        .await;
    }
}