From d2369c3dea94c09f45cb109696f8606bdd5cad3c Mon Sep 17 00:00:00 2001 From: laurent Date: Mon, 24 Jul 2023 15:13:47 +0100 Subject: [PATCH] Polish the llama2 wasm ui. --- candle-wasm-examples/llama2-c/Cargo.toml | 1 + candle-wasm-examples/llama2-c/src/app.rs | 19 +++++++++++++++++-- candle-wasm-examples/llama2-c/src/bin/app.rs | 1 + .../llama2-c/src/bin/worker.rs | 1 + candle-wasm-examples/llama2-c/src/model.rs | 2 +- candle-wasm-examples/llama2-c/src/worker.rs | 18 +++++++++++++----- 6 files changed, 34 insertions(+), 8 deletions(-) diff --git a/candle-wasm-examples/llama2-c/Cargo.toml b/candle-wasm-examples/llama2-c/Cargo.toml index 22d9cfe81c..6aae0e59b3 100644 --- a/candle-wasm-examples/llama2-c/Cargo.toml +++ b/candle-wasm-examples/llama2-c/Cargo.toml @@ -24,6 +24,7 @@ serde = { workspace = true } serde_json = { workspace = true } # Wasm specific crates. +console_error_panic_hook = "0.1.7" getrandom = { version = "0.2", features = ["js"] } gloo = "0.8" js-sys = "0.3.64" diff --git a/candle-wasm-examples/llama2-c/src/app.rs b/candle-wasm-examples/llama2-c/src/app.rs index 460ac053af..eab0aa6ef8 100644 --- a/candle-wasm-examples/llama2-c/src/app.rs +++ b/candle-wasm-examples/llama2-c/src/app.rs @@ -1,5 +1,6 @@ use crate::console_log; use crate::worker::{ModelData, Worker, WorkerInput, WorkerOutput}; +use std::str::FromStr; use wasm_bindgen::prelude::*; use wasm_bindgen_futures::JsFuture; use yew::{html, Component, Context, Html}; @@ -42,6 +43,7 @@ pub struct CurrentDecode { pub struct App { status: String, + temperature: std::rc::Rc>, generated: String, current_decode: Option, worker: Box>, @@ -73,6 +75,7 @@ impl Component for App { let worker = Worker::bridge(std::rc::Rc::new(cb)); Self { status, + temperature: std::rc::Rc::new(std::cell::RefCell::new(0.)), generated: String::new(), current_decode: None, worker, @@ -109,7 +112,10 @@ impl Component for App { self.current_decode = Some(CurrentDecode { start_time }); self.status = "generating...".to_string(); self.generated.clear(); - ctx.link().send_message(Msg::WorkerInMsg(WorkerInput::Run)) + let temp = *self.temperature.borrow(); + console_log!("temp: {}", temp); + ctx.link() + .send_message(Msg::WorkerInMsg(WorkerInput::Run(temp))) } true } @@ -151,8 +157,16 @@ impl Component for App { } fn view(&self, ctx: &Context) -> Html { + use yew::TargetCast; + let temperature = self.temperature.clone(); + let oninput = move |e: yew::InputEvent| { + let input: web_sys::HtmlInputElement = e.target_unchecked_into(); + if let Ok(temp) = f64::from_str(&input.value()) { + *temperature.borrow_mut() = temp + } + }; html! { -
+

{"Running "} {"llama2.c"} {" in the browser using rust/wasm with "} @@ -161,6 +175,7 @@ impl Component for App {

{"Once the weights have loaded, click on the run button to start generating content."}

+ {"temperature: "}

diff --git a/candle-wasm-examples/llama2-c/src/bin/app.rs b/candle-wasm-examples/llama2-c/src/bin/app.rs index 3428f6fff3..717eeafc6d 100644 --- a/candle-wasm-examples/llama2-c/src/bin/app.rs +++ b/candle-wasm-examples/llama2-c/src/bin/app.rs @@ -1,4 +1,5 @@ fn main() { wasm_logger::init(wasm_logger::Config::new(log::Level::Trace)); + console_error_panic_hook::set_once(); yew::Renderer::::new().render(); } diff --git a/candle-wasm-examples/llama2-c/src/bin/worker.rs b/candle-wasm-examples/llama2-c/src/bin/worker.rs index d8ca2172ee..accb51b735 100644 --- a/candle-wasm-examples/llama2-c/src/bin/worker.rs +++ b/candle-wasm-examples/llama2-c/src/bin/worker.rs @@ -1,4 +1,5 @@ use yew_agent::PublicWorker; fn main() { + console_error_panic_hook::set_once(); candle_wasm_example_llama2::Worker::register(); } diff --git a/candle-wasm-examples/llama2-c/src/model.rs b/candle-wasm-examples/llama2-c/src/model.rs index 13f939db10..8b0b3c3ea3 100644 --- a/candle-wasm-examples/llama2-c/src/model.rs +++ b/candle-wasm-examples/llama2-c/src/model.rs @@ -20,7 +20,7 @@ pub struct Cache { masks: Arc>>, pub use_kv_cache: bool, #[allow(clippy::type_complexity)] - kvs: Arc>>>, + pub kvs: Arc>>>, cos: Tensor, sin: Tensor, device: Device, diff --git a/candle-wasm-examples/llama2-c/src/worker.rs b/candle-wasm-examples/llama2-c/src/worker.rs index 9b0351d669..d64da8c6b0 100644 --- a/candle-wasm-examples/llama2-c/src/worker.rs +++ b/candle-wasm-examples/llama2-c/src/worker.rs @@ -107,9 +107,11 @@ impl LogitsProcessor { } impl Model { - fn run(&self, link: &WorkerLink, id: HandlerId) -> Result<()> { + fn run(&self, link: &WorkerLink, id: HandlerId, temp: f64) -> Result<()> { let dev = Device::Cpu; - let mut logits_processor = LogitsProcessor::new(299792458, None); + let temp = if temp <= 0. { None } else { Some(temp) }; + console_log!("{temp:?}"); + let mut logits_processor = LogitsProcessor::new(299792458, temp); let mut index_pos = 0; let mut tokens = vec![1u32]; @@ -299,7 +301,7 @@ pub struct Worker { #[derive(Serialize, Deserialize)] pub enum WorkerInput { ModelData(ModelData), - Run, + Run(f64), } #[derive(Serialize, Deserialize)] @@ -332,10 +334,16 @@ impl yew_agent::Worker for Worker { } Err(err) => Err(format!("model creation error {err:?}")), }, - WorkerInput::Run => match &self.model { + WorkerInput::Run(temp) => match &mut self.model { None => Err("model has not been set yet".to_string()), Some(model) => { - let result = model.run(&self.link, id).map_err(|e| e.to_string()); + { + let mut cache = model.cache.kvs.lock().unwrap(); + for elem in cache.iter_mut() { + *elem = None + } + } + let result = model.run(&self.link, id, temp).map_err(|e| e.to_string()); Ok(WorkerOutput::GenerationDone(result)) } },