update am

Author: yujonglee

apps/web/content/docs/developers/16.local-models.mdx

@@ -13,8 +13,19 @@ description: "Learn about local models in Hyprnote"
 
 - All STT models are stored in `~/Library/Application Support/hyprnote/models/stt/`
 
+### Manual Download
+
 | URL | Folder |
 | --- | --- |
 | https://huggingface.co/argmaxinc/whisperkit-pro/tree/main/openai_whisper-large-v3-v20240930_626MB | `openai_whisper-large-v3-v20240930_626MB` |
 | https://huggingface.co/argmaxinc/parakeetkit-pro/tree/main/nvidia_parakeet-v2_476MB | `nvidia_parakeet-v2_476MB` |
 | https://huggingface.co/argmaxinc/parakeetkit-pro/tree/main/nvidia_parakeet-v3_494MB | `nvidia_parakeet-v3_494MB` |
+| https://huggingface.co/argmaxinc/ctckit-pro/tree/main/parakeet-tdt_ctc-110m | `parakeet-tdt_ctc-110m` |
+
+If you have [Huggingface CLI](https://huggingface.co/docs/huggingface_hub/guides/cli) installed: 
+
+```bash
+hf download argmaxinc/ctckit-pro \
+  --include "parakeet-tdt_ctc-110m/*" \
+  --local-dir "$HOME/Library/Application Support/hyprnote/models/stt"
+```

crates/am/src/client.rs

@@ -33,13 +33,40 @@ impl Client {
         }
     }
 
+    pub async fn wait_for_ready(
+        &self,
+        max_wait_time: Option<u32>,
+        poll_interval: Option<u32>,
+    ) -> Result<ServerStatus, Error> {
+        let url = format!("{}/waitForReady", self.base_url);
+        let mut request = self.client.get(&url);
+
+        if let Some(max_wait) = max_wait_time {
+            request = request.query(&[("maxWaitTime", max_wait)]);
+        }
+        if let Some(poll) = poll_interval {
+            request = request.query(&[("pollInterval", poll)]);
+        }
+
+        let response = request.send().await?;
+
+        match response.status() {
+            StatusCode::OK => Ok(response.json().await?),
+            StatusCode::BAD_REQUEST | StatusCode::REQUEST_TIMEOUT => {
+                Err(self.handle_error_response(response).await)
+            }
+            _ => Err(Error::UnexpectedResponse),
+        }
+    }
+
     pub async fn init(&self, request: InitRequest) -> Result<InitResponse, Error> {
         if !request.api_key.starts_with("ax_") {
             return Err(Error::InvalidApiKey);
         }
 
         let url = format!("{}/init", self.base_url);
         let response = self.client.post(&url).json(&request).send().await?;
+        println!("{:?}", request);
 
         match response.status() {
             StatusCode::OK => Ok(response.json().await?),
@@ -100,9 +127,18 @@ impl InitRequest {
         Self {
             api_key: api_key.into(),
             model: None,
+            model_token: None,
+            download_base: None,
             model_repo: None,
             model_folder: None,
+            tokenizer_folder: None,
+            fast_load: None,
+            fast_load_encoder_compute_units: None,
+            fast_load_decoder_compute_units: None,
+            model_vad: None,
+            verbose: None,
             custom_vocabulary: None,
+            custom_vocabulary_model_folder: None,
         }
     }
 
@@ -124,6 +160,13 @@ impl InitRequest {
         match model {
             crate::AmModel::ParakeetV2 => {
                 self.custom_vocabulary = Some(vec![]);
+                self.custom_vocabulary_model_folder = Some(
+                    base_dir
+                        .as_ref()
+                        .join("parakeet-tdt_ctc-110m")
+                        .to_string_lossy()
+                        .to_string(),
+                );
             }
             _ => {
                 self.custom_vocabulary = None;

crates/am/src/lib.rs

@@ -18,13 +18,7 @@ mod tests {
         let status = client.status().await;
         println!("{:?}", status);
         client
-            .init(InitRequest {
-                api_key: "".to_string(),
-                model: Some("nvidia_parakeet-v2_476MB".to_string()),
-                model_repo: Some("argmaxinc/parakeetkit-pro".to_string()),
-                model_folder: None,
-                custom_vocabulary: Some(vec![]),
-            })
+            .init(InitRequest::new("").with_model(AmModel::ParakeetV2, "/tmp"))
             .await
             .unwrap();
         assert!(true);

crates/am/src/types.rs

@@ -14,10 +14,20 @@ common_derives! {
         pub model_state: ModelState,
         pub verbose: bool,
         #[serde(skip_serializing_if = "Option::is_none")]
+        pub download_progress: Option<DownloadProgress>,
+        #[serde(skip_serializing_if = "Option::is_none")]
         pub message: Option<String>,
     }
 }
 
+common_derives! {
+    #[serde(rename_all = "camelCase")]
+    pub struct DownloadProgress {
+        pub progress_percentage: f64,
+        pub is_downloading: bool,
+    }
+}
+
 common_derives! {
     #[derive(Eq, PartialEq)]
     #[serde(rename_all = "lowercase")]
@@ -51,11 +61,41 @@ common_derives! {
         #[serde(skip_serializing_if = "Option::is_none")]
         pub model: Option<String>,
         #[serde(skip_serializing_if = "Option::is_none")]
+        pub model_token: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub download_base: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
         pub model_repo: Option<String>,
         #[serde(skip_serializing_if = "Option::is_none")]
         pub model_folder: Option<String>,
         #[serde(skip_serializing_if = "Option::is_none")]
+        pub tokenizer_folder: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub fast_load: Option<bool>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub fast_load_encoder_compute_units: Option<ComputeUnits>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub fast_load_decoder_compute_units: Option<ComputeUnits>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub model_vad: Option<bool>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub verbose: Option<bool>,
+        #[serde(skip_serializing_if = "Option::is_none")]
         pub custom_vocabulary: Option<Vec<String>>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub custom_vocabulary_model_folder: Option<String>,
+    }
+}
+
+common_derives! {
+    #[serde(rename_all = "lowercase")]
+    pub enum ComputeUnits {
+        Cpu,
+        #[serde(rename = "cpuandgpu")]
+        CpuAndGpu,
+        #[serde(rename = "cpuandneuralengine")]
+        CpuAndNeuralEngine,
+        All,
     }
 }