Merge pull request #1 from Gitzoz/downloader

Download whisper model from hugging face

.gitignore

@@ -15,6 +15,8 @@ Cargo.lock
 
 .vscode/*
 
-models/
+models/*
+
+whisper.**
 
 .DS_Store

Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "transkriboi"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+whisper-rs = "0.8.0"
+reqwest = { version = "0.11", features = ["json"] }
+tokio = { version = "1", features = ["full"] }

src/downloader/mod.rs

@@ -0,0 +1,2 @@
+pub mod models;
+pub mod web;

src/downloader/models.rs

@@ -0,0 +1,43 @@
+use core::fmt;
+
+pub enum ModelFormat {
+    GGML,
+    COREML,
+}
+
+impl fmt::Display for ModelFormat {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            ModelFormat::COREML => write!(f, "whisper.cpp-coreml"),
+            ModelFormat::GGML => write!(f, "whisper.cpp"),
+        }
+    }
+}
+
+pub enum ModelSize {
+    TinyEn,
+    Tiny,
+    BaseEn,
+    Base,
+    SmallEn,
+    Small,
+    Medium,
+    MediumEn,
+    Large,
+}
+
+impl fmt::Display for ModelSize {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            ModelSize::Base => write!(f, "base"),
+            ModelSize::BaseEn => write!(f, "base.en"),
+            ModelSize::Tiny => write!(f, "tiny"),
+            ModelSize::TinyEn => write!(f, "tiny.en"),
+            ModelSize::Small => write!(f, "small"),
+            ModelSize::SmallEn => write!(f, "small.en"),
+            ModelSize::Medium => write!(f, "medium"),
+            ModelSize::MediumEn => write!(f, "medium.en"),
+            ModelSize::Large => write!(f, "large"),
+        }
+    }
+}

src/downloader/web.rs

@@ -0,0 +1,23 @@
+use crate::downloader::models::{ModelFormat, ModelSize};
+use tokio::{fs::File, io::AsyncWriteExt};
+
+const MODEL_SRC_URL: &str = "https://huggingface.co/ggerganov/";
+const PFX: &str = "resolve/main/ggml";
+
+pub async fn download_model(
+    output_path: &str,
+    model_size: ModelSize,
+    model_format: ModelFormat,
+) -> Result<String, Box<dyn std::error::Error>> {
+    let url = format!(
+        "{}{}/{}-{}.bin",
+        MODEL_SRC_URL, model_format, PFX, model_size
+    );
+    let file_name = format!("{}{}_{}", output_path, model_format, model_size);
+    let mut response = reqwest::get(url).await?;
+    let mut file = File::create(file_name.clone()).await?;
+    while let Some(chunk) = response.chunk().await? {
+        file.write_all(&chunk).await?;
+    }
+    Ok(file_name)
+}

src/main.rs

@@ -0,0 +1,14 @@
+mod downloader;
+use crate::downloader::models::*;
+use crate::downloader::web::download_model;
+
+#[tokio::main]
+async fn main() {
+    println!("Hello, world!");
+    let op = download_model("./models/", ModelSize::Tiny, ModelFormat::GGML).await;
+    match op {
+        Ok(n) => println!("Ok {}", n),
+        Err(e) => println!("Error {}", e)
+    }
+
+}