chore: done

Author: FL33TW00D

.cargo/config.toml

@@ -2,3 +2,4 @@
 [build]
 rustflags = [ "--cfg=web_sys_unstable_apis" ]
 rustdocflags = [ "--cfg=web_sys_unstable_apis" ]
+#target = "wasm32-unknown-unknown"

crates/ratchet-models/src/whisper/mod.rs

@@ -15,6 +15,7 @@ pub mod tokenizer;
 pub mod transcribe;
 pub mod transcript;
 
+pub use config::Config;
 pub use decoder::WhisperDecoder;
 pub use encoder::WhisperEncoder;
 pub use model::Whisper;

crates/ratchet-models/src/whisper/model.rs

@@ -11,7 +11,7 @@ use ratchet::NDArrayExt;
 use hf_hub::api::sync::Api;
 
 #[cfg(target_arch = "wasm32")]
-use {ratchet_hub::ApiBuilder, ratchet_hub::RepoType, wasm_bindgen::JsError};
+use {js_sys::Uint8Array, ratchet_hub::ApiBuilder, ratchet_hub::RepoType, wasm_bindgen::JsError};
 
 use crate::registry::WhisperVariants;
 use crate::whisper::{options::Language, task::DecodingTask, tokenizer::WhisperTokenizer};
@@ -30,6 +30,7 @@ pub struct Whisper {
 }
 
 impl Whisper {
+    #[cfg(not(target_arch = "wasm32"))]
     pub fn load<R: BufRead + Seek>(
         header: Header,
         reader: &mut R,
@@ -57,14 +58,43 @@ impl Whisper {
         })
     }
 
+    #[cfg(target_arch = "wasm32")]
+    pub async fn load<R: BufRead + Seek>(
+        header: Header,
+        reader: &mut R,
+        device: Device,
+    ) -> Result<Self, JsError> {
+        let mel_bytes = Self::fetch_resource(WhisperVariants::Tiny, "melfilters.bytes").await?;
+        let mut mel_filters = vec![0f32; mel_bytes.len() / 4];
+        <byteorder::LittleEndian as byteorder::ByteOrder>::read_f32_into(
+            &mel_bytes,
+            &mut mel_filters,
+        );
+        let specgen = SpectrogramGenerator::new(mel_filters);
+
+        let config: Config = serde_json::from_slice(
+            &Self::fetch_resource(WhisperVariants::Tiny, "config.json").await?,
+        )?;
+        let encoder = WhisperEncoder::load(&header, &config, reader, &device).unwrap();
+        let decoder = WhisperDecoder::load(&header, &config, reader, &device).unwrap();
+
+        Ok(Self {
+            specgen,
+            encoder,
+            decoder,
+            config,
+            device,
+        })
+    }
+
     #[cfg(target_arch = "wasm32")]
     pub async fn fetch_resource(
         variant: WhisperVariants,
         resource: &str,
     ) -> Result<Vec<u8>, JsError> {
         let repo_id = variant.repo_id();
         let model_repo = ApiBuilder::from_hf(repo_id, RepoType::Model).build();
-        model_repo.get(resource).await?
+        Ok(model_repo.get(resource).await?.to_vec())
     }
 
     #[cfg(not(target_arch = "wasm32"))]

crates/ratchet-models/src/whisper/transcribe.rs

@@ -90,7 +90,7 @@ pub async fn transcribe(
     callback: Option<impl Fn(StreamedSegment)>,
 ) -> anyhow::Result<TranscriptionResult> {
     let runtime = Instant::now();
-    let n_mels = model.hparams.n_mels as usize;
+    let n_mels = model.config.n_mels as usize;
     let mel = model.specgen.generate(audio)?.to(&model.device).await?;
     let content_frames = mel.shape()[mel.rank() - 1] - N_FRAMES;
 

crates/ratchet-models/tests/whisper.rs

@@ -5,7 +5,8 @@ use ndarray::{s, Axis};
 use ndarray_stats::QuantileExt;
 use ratchet::{shape, Device, DeviceRequest, Tensor};
 use ratchet_hub::{ApiBuilder, RepoType};
-use ratchet_models::whisper::{Whisper, WhisperDecoder, WhisperEncoder};
+use ratchet_loader::gguf::gguf;
+use ratchet_models::whisper::{Config, Whisper, WhisperDecoder, WhisperEncoder};
 use ratchet_nn::Module;
 use std::path::PathBuf;
 use wasm_bindgen::prelude::*;
@@ -21,22 +22,24 @@ fn log_init() {
 async fn tiny_encoder() -> Result<(), JsValue> {
     log_init();
     let model_repo = ApiBuilder::from_hf("FL33TW00D-HF/whisper-tiny", RepoType::Model).build();
-    let model_data = model_repo.get("tiny_f32.bin").await?;
+    let model_data = model_repo.get("tiny_f32.gguf").await?;
+    let config_data = model_repo.get("config.json").await?;
 
     let ground_repo = ApiBuilder::from_hf("FL33TW00D-HF/ratchet-util", RepoType::Dataset).build();
     let input_npy = ground_repo.get("jfk_tiny_encoder_input.npy").await?;
     let ground_npy = ground_repo.get("jfk_tiny_encoder_hs.npy").await?;
 
     let mut reader = std::io::BufReader::new(std::io::Cursor::new(model_data.to_vec()));
-    let gg = Whisper::load_ggml(&mut reader).unwrap();
+    let header = gguf::Header::read(&mut reader).unwrap();
+    let config: Config = serde_json::from_slice(&config_data.to_vec()).unwrap();
 
     let device = Device::request_device(DeviceRequest::GPU).await.unwrap();
 
     let input_data = &input_npy.to_vec();
     let input = Tensor::from_npy_bytes::<f32>(input_data, &device).unwrap();
     let ground = Tensor::from_npy_bytes::<f32>(&ground_npy.to_vec(), &Device::CPU).unwrap();
 
-    let encoder = WhisperEncoder::load(&gg, &mut reader, &device).unwrap();
+    let encoder = WhisperEncoder::load(&header, &config, &mut reader, &device).unwrap();
     let result = encoder.schedule(input).unwrap().resolve().unwrap();
     let ours = result.to(&Device::CPU).await.unwrap();
     ground.all_close(&ours, 1e-3, 1e-3).unwrap();
@@ -46,18 +49,19 @@ async fn tiny_encoder() -> Result<(), JsValue> {
 #[wasm_bindgen_test]
 async fn tiny_decoder() -> Result<(), JsValue> {
     let model_repo = ApiBuilder::from_hf("FL33TW00D-HF/whisper-tiny", RepoType::Model).build();
-    let model_data = model_repo.get("tiny_f32.bin").await?;
+    let model_data = model_repo.get("tiny_f32.gguf").await?;
+    let config_data = model_repo.get("config.json").await?;
 
     let ground_repo = ApiBuilder::from_hf("FL33TW00D-HF/ratchet-util", RepoType::Dataset).build();
     let hs_data = ground_repo.get("jfk_tiny_encoder_hs.npy").await?;
 
     let mut reader = std::io::BufReader::new(std::io::Cursor::new(model_data.to_vec()));
-    let gg_disk = Whisper::load_ggml(&mut reader).unwrap();
-    assert_eq!(gg_disk.tensors.len(), 167);
+    let header = gguf::Header::read(&mut reader).unwrap();
+    let config: Config = serde_json::from_slice(&config_data.to_vec()).unwrap();
 
     let device = Device::request_device(DeviceRequest::GPU).await.unwrap();
     let audio_ctx = Tensor::from_npy_bytes::<f32>(&hs_data.to_vec(), &device).unwrap();
-    let mut decoder = WhisperDecoder::load(&gg_disk, &mut reader, &device).unwrap();
+    let mut decoder = WhisperDecoder::load(&header, &config, &mut reader, &device).unwrap();
 
     let mut tokens = vec![50258, 50259, 50359];
     let mut all_tokens = tokens.clone();