candle whisper basic version completed

Author: hubertshelley

examples/candle_whisper/Cargo.toml

@@ -26,10 +26,16 @@ serde = { version = "1.0", features = ["derive"] }
 silent = { path = "../../silent", features = ["full"] }
 symphonia = { version = "0.5.3", features = ["all"] }
 anyhow = "1.0.79"
+tokio = { version = "1.35.1", features = ["full"] }
+
+#candle-core = { version = "0.3.2" }
+#candle-nn = { version = "0.3.2" }
+#candle-transformers = { version = "0.3.2" }
+# version = "0.3.2" is not working for metal
+candle-core = { git = "https://github.com/huggingface/candle" }
+candle-nn = { git = "https://github.com/huggingface/candle" }
+candle-transformers = { git = "https://github.com/huggingface/candle" }
 
-candle-core = { version = "0.3.2" }
-candle-nn = { version = "0.3.2" }
-candle-transformers = { version = "0.3.2" }
 tokenizers = { version = "0.15.0", features = ["onig"] }
 rand = "0.8.5"
 serde_json = "1.0.109"

examples/candle_whisper/src/args.rs

@@ -26,39 +26,10 @@ pub(crate) struct Args {
     #[arg(long, default_value = "tiny.en")]
     pub(crate) model: WhichModel,
 
-    /// The input to be processed, in wav format, will default to `jfk.wav`. Alternatively
-    /// this can be set to sample:jfk, sample:gb1, ... to fetch a sample from the following
-    /// repo: https://huggingface.co/datasets/Narsil/candle_demo/
-    #[arg(long)]
-    pub(crate) input: String,
-
     /// The seed to use when generating random samples.
     #[arg(long, default_value_t = 299792458)]
     pub(crate) seed: u64,
 
-    /// Enable tracing (generates a trace-timestamp.json file).
-    #[arg(long)]
-    tracing: bool,
-
     #[arg(long)]
     quantized: bool,
-
-    /// Language.
-    #[arg(long)]
-    pub(crate) language: Option<String>,
-
-    /// Task, when no task is specified, the input tokens contain only the sot token which can
-    /// improve things when in no-timestamp mode.
-    #[arg(long)]
-    pub(crate) task: Option<Task>,
-
-    /// Timestamps mode, this is not fully implemented yet.
-    #[arg(long)]
-    pub(crate) timestamps: bool,
-
-    /// Print the full DecodingResult structure rather than just the text.
-    #[arg(long)]
-    pub(crate) verbose: bool,
-    #[arg(long, default_value_t = 0.0)]
-    pub(crate) temperature: f64,
 }

examples/candle_whisper/src/handlers.rs

@@ -20,7 +20,9 @@ use candle_nn::VarBuilder;
 use candle_transformers::models::whisper::{self as m, audio, Config};
 use silent::{Request, Response, Result as SilentResult, SilentError, StatusCode};
 use std::path::PathBuf;
+use std::sync::Arc;
 use tokenizers::Tokenizer;
+use tokio::sync::Mutex;
 
 use crate::pcm_decode::pcm_decode;
 use crate::types::{CreateTranscriptionRequest, CreateTranscriptionResponse};
@@ -40,6 +42,7 @@ pub(crate) struct WhisperModel {
     mel_filters: Vec<f32>,
     device: candle::Device,
 }
+
 pub(crate) fn init_model(args: Args) -> Result<WhisperModel> {
     let device = device(args.cpu)?;
     let model_id = args.model_id;
@@ -75,23 +78,11 @@ pub(crate) fn init_model(args: Args) -> Result<WhisperModel> {
     })
 }
 
-fn handle(req: CreateTranscriptionRequest) -> CreateTranscriptionResponse {
-    CreateTranscriptionResponse::new(vec![], req.response_format.clone())
-}
-
-async fn create_transcription(mut req: Request) -> SilentResult<Response> {
-    let req = req.form_data().await?.try_into().map_err(|e| {
-        SilentError::business_error(
-            StatusCode::BAD_REQUEST,
-            format!("failed to parse request: {}", e),
-        )
-    })?;
-    let res = handle(req);
-    Ok(res.into())
-}
-pub(crate) fn handle1(args: Args) -> Result<()> {
-    let mut whisper_model = init_model(args.clone())?;
-    let input = PathBuf::from(args.input);
+fn handle(
+    req: CreateTranscriptionRequest,
+    mut whisper_model: WhisperModel,
+) -> Result<CreateTranscriptionResponse> {
+    let input = req.file.path().clone();
 
     let pcm_data = pcm_decode(input);
     let config = whisper_model.config.clone();
@@ -111,7 +102,7 @@ pub(crate) fn handle1(args: Args) -> Result<()> {
     let mut model = whisper_model.model.clone();
     let tokenizer = whisper_model.tokenizer.clone();
     let device = whisper_model.device.clone();
-    let language_token = match (args.model.is_multilingual(), args.language) {
+    let language_token = match (req.model.is_multilingual(), req.language) {
         (true, None) => Some(multilingual::detect_language(&mut model, &tokenizer, &mel)?),
         (false, None) => None,
         (true, Some(language)) => match token_id(&tokenizer, &format!("<|{language}|>")) {
@@ -126,16 +117,40 @@ pub(crate) fn handle1(args: Args) -> Result<()> {
     let mut dc = Decoder::new(
         model,
         tokenizer,
-        args.seed,
+        299792458,
         &device,
         language_token,
-        args.task,
-        args.timestamps,
-        args.verbose,
-        args.temperature,
+        None,
+        req.response_format.has_timestamps(),
+        req.response_format.is_verbose(),
+        req.temperature,
     )?;
-    println!("starting decoding");
-    let s = dc.run(&mel)?;
-    println!("done: {:?}", s);
-    Ok(())
+    let segments = dc.run(&mel)?;
+    Ok(CreateTranscriptionResponse::new(
+        segments,
+        req.response_format.clone(),
+    ))
+}
+
+pub(crate) async fn create_transcription(mut req: Request) -> SilentResult<Response> {
+    let whisper_model = req
+        .configs()
+        .get::<Arc<Mutex<WhisperModel>>>()
+        .unwrap()
+        .lock()
+        .await
+        .clone();
+    let req = req.form_data().await?.try_into().map_err(|e| {
+        SilentError::business_error(
+            StatusCode::BAD_REQUEST,
+            format!("failed to parse request: {}", e),
+        )
+    })?;
+    let res = handle(req, whisper_model).map_err(|e| {
+        SilentError::business_error(
+            StatusCode::INTERNAL_SERVER_ERROR,
+            format!("failed to create transcription: {}", e),
+        )
+    })?;
+    Ok(res.into())
 }

examples/candle_whisper/src/main.rs

@@ -3,11 +3,11 @@ mod pcm_decode;
 mod types;
 
 use crate::args::Args;
-use crate::handlers::{handle1, init_model};
+use crate::handlers::{create_transcription, init_model};
 use clap::Parser;
 use silent::prelude::*;
-use std::path::PathBuf;
-use tokenizers::Tokenizer;
+use std::sync::Arc;
+use tokio::sync::Mutex;
 
 mod args;
 mod decoder;
@@ -19,7 +19,11 @@ fn main() {
     logger::fmt().with_max_level(Level::INFO).init();
     let args = Args::parse();
     let mut configs = Configs::default();
-    handle1(args).unwrap();
-    // let route = Route::new("").get(|_req| async { Ok("hello world") });
-    // Server::new().run(route);
+    let whisper_model = init_model(args.clone()).expect("failed to initialize model");
+    configs.insert(Arc::new(Mutex::new(whisper_model)));
+    let route = Route::new("/v1/audio/transcriptions").post(create_transcription);
+    Server::new()
+        .with_configs(configs)
+        .bind("0.0.0.0:8000".parse().unwrap())
+        .run(route);
 }

examples/candle_whisper/src/model.rs

@@ -1,3 +1,4 @@
+use anyhow::{bail, Error};
 use candle_core::{self as candle, Tensor};
 use candle_transformers::models::whisper::{self as m, Config};
 use clap::ValueEnum;
@@ -26,6 +27,29 @@ pub(crate) enum WhichModel {
     DistilLargeV2,
 }
 
+impl TryFrom<String> for WhichModel {
+    type Error = Error;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        Ok(match value {
+            s if s == "tiny" => Self::Tiny,
+            s if s == "tiny.en" => Self::TinyEn,
+            s if s == "base" => Self::Base,
+            s if s == "base.en" => Self::BaseEn,
+            s if s == "small" => Self::Small,
+            s if s == "small.en" => Self::SmallEn,
+            s if s == "medium" => Self::Medium,
+            s if s == "medium.en" => Self::MediumEn,
+            s if s == "large" => Self::Large,
+            s if s == "large-v2" => Self::LargeV2,
+            s if s == "large-v3" => Self::LargeV3,
+            s if s == "distil-medium.en" => Self::DistilMediumEn,
+            s if s == "distil-large-v2" => Self::DistilLargeV2,
+            _ => bail!("invalid model"),
+        })
+    }
+}
+
 impl WhichModel {
     pub(crate) fn is_multilingual(&self) -> bool {
         match self {
@@ -61,6 +85,7 @@ impl WhichModel {
         }
     }
 }
+
 #[derive(Clone, Debug)]
 pub enum Model {
     Normal(m::model::Whisper),

examples/candle_whisper/src/types.rs

@@ -1,5 +1,6 @@
 use crate::decoder::Segment;
 use crate::model::WhichModel;
+use anyhow::Error;
 use serde::Deserialize;
 use serde_json::json;
 use silent::prelude::{FilePart, FormData};
@@ -14,11 +15,39 @@ pub(crate) enum ResponseFormat {
     Vtt,
 }
 
+impl From<String> for ResponseFormat {
+    fn from(value: String) -> Self {
+        match value.as_str() {
+            "json" => Self::Json,
+            "text" => Self::Text,
+            "srt" => Self::Srt,
+            "verbose_json" => Self::VerboseJson,
+            "vtt" => Self::Vtt,
+            _ => Self::Json,
+        }
+    }
+}
+
+impl ResponseFormat {
+    pub(crate) fn is_verbose(&self) -> bool {
+        match self {
+            Self::Json | Self::Text => false,
+            Self::VerboseJson | Self::Srt | Self::Vtt => true,
+        }
+    }
+    pub(crate) fn has_timestamps(&self) -> bool {
+        match self {
+            Self::VerboseJson | Self::Json | Self::Text => false,
+            Self::Srt | Self::Vtt => true,
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct CreateTranscriptionRequest {
     // The audio file object (not file name) to transcribe, in one of these formats: wav.
     pub(crate) file: FilePart,
-    // ID of the model to use. Only whisper-large-v3 is currently available.
+    // ID of the model to use. Only large-v3 is currently available.
     pub(crate) model: WhichModel,
     // The language of the input audio. Supplying the input language in ISO-639-1 format will improve accuracy and latency.
     pub(crate) language: Option<String>,
@@ -42,16 +71,24 @@ impl TryFrom<&FormData> for CreateTranscriptionRequest {
                 StatusCode::BAD_REQUEST,
                 "file is required".to_string(),
             ))?;
-        let model = serde_json::from_str(&value.fields.get("model").cloned().ok_or(
-            SilentError::business_error(StatusCode::BAD_REQUEST, "model is required".to_string()),
-        )?)?;
-        let response_format = serde_json::from_str(
-            &value
-                .fields
-                .get("response_format")
-                .cloned()
-                .unwrap_or("json".to_string()),
-        )?;
+        let model = value
+            .fields
+            .get("model")
+            .cloned()
+            .ok_or(SilentError::business_error(
+                StatusCode::BAD_REQUEST,
+                "model is required".to_string(),
+            ))?
+            .try_into()
+            .map_err(|e: Error| {
+                SilentError::business_error(StatusCode::BAD_REQUEST, e.to_string())
+            })?;
+        let response_format = value
+            .fields
+            .get("response_format")
+            .cloned()
+            .unwrap_or("json".to_string())
+            .into();
         Ok(Self {
             file,
             model,
@@ -73,6 +110,7 @@ impl TryFrom<&FormData> for CreateTranscriptionRequest {
         })
     }
 }
+
 pub struct CreateTranscriptionResponse {
     segments: Vec<Segment>,
     format: ResponseFormat,
@@ -88,18 +126,19 @@ impl From<CreateTranscriptionResponse> for Response {
     fn from(value: CreateTranscriptionResponse) -> Self {
         match value.format {
             ResponseFormat::Json => json!({
-                "text": value.segments.iter().map(|s| s.text()).collect::<Vec<_>>(),
+                "text": value.segments.iter().map(|s| s.text()).collect::<Vec<_>>().join(""),
             })
             .into(),
             ResponseFormat::Text => value
                 .segments
                 .iter()
                 .map(|s| s.text())
                 .collect::<Vec<_>>()
+                .join("")
                 .into(),
             ResponseFormat::Srt => !unimplemented!("Srt"),
             ResponseFormat::VerboseJson => json!({
-                "text": value.segments.iter().map(|s| s.text()).collect::<Vec<_>>(),
+                "text": value.segments.iter().map(|s| s.text()).collect::<Vec<_>>().join(""),
             })
             .into(),
             ResponseFormat::Vtt => !unimplemented!("Vtt"),