Load model in test

src/lib.rs

@@ -1,4 +1,4 @@
-mod quantized_mistral;
+pub mod quantized_mistral;
 mod utils;
 use wasm_bindgen::prelude::*;
 

tests/web.rs

@@ -7,6 +7,7 @@ use std::fmt::format;
 use std::println;
 
 use js_sys::Uint8Array;
+use transformers_wasm::quantized_mistral::Model;
 use wasm_bindgen::{prelude::*, JsValue};
 use wasm_bindgen_futures::JsFuture;
 use wasm_bindgen_test::*;
@@ -53,14 +54,24 @@ async fn pass() -> Result<(), JsValue> {
     let tokenizer_url = "http://localhost:45678/tokenizer.json";
     let model_url = "http://localhost:45678/tinymistral-248m.q4_k_m.gguf";
 
-    let tokenizer: Vec<u8> = load_binary(&tokenizer_url).await?;
-    let tokenizer_len = format!("{}", &tokenizer.len());
-    console::log_2(&"tokenizer size".into(), &tokenizer_len.into());
+    let tokenizer_blob: Vec<u8> = load_binary(&tokenizer_url).await?;
+    let tokenizer_blob_len = format!("{}", &tokenizer_blob.len());
+    console::log_2(&"tokenizer blob size".into(), &tokenizer_blob_len.into());
 
-    let model: Vec<u8> = load_binary(&model_url).await?;
-    let model_len = format!("{}", &model.len());
-    console::log_2(&"model size".into(), &model_len.into());
+    let model_blob: Vec<u8> = load_binary(&model_url).await?;
+    let model_blob_len = format!("{}", &model_blob.len());
+    console::log_2(&"model blob size".into(), &model_blob_len.into());
 
+    let mut model = Model::new(tokenizer_blob, model_blob)?;
+
+    let prompt: String = String::from("What is a good recipe for onion soup");
+    let temp: f64 = 0.8;
+    let top_p: f64 = 1.;
+    let repeat_penalty: f32 = 1.1;
+    let seed: u64 = 203948203948;
+    let first_result: String = model.init_with_prompt(prompt, temp, top_p, repeat_penalty, seed)?;
+
+    console::log_2(&"first prompt result".into(), &first_result.into());
     assert_eq!(1 + 1, 2);
     Ok(())
 }