Update web-rwkv to v0.3.9.

Ai00-X · Nov 2, 2023 · aab9e06 · aab9e06
1 parent 6c8871d
commit aab9e06
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 15 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ai00_server"
-version = "0.2.2"
+version = "0.2.3"
 edition = "2021"
 authors = ["Gu ZhenNiu <[email protected]>", "Zhang Zhenyuan <[email protected]>"]
 license = "MIT OR Apache-2.0"
@@ -19,8 +19,8 @@ tower = { version = "0.4", features = ["util"] }
 tower-http = { version = "0.4", features = ["full"] }
 tokio = { version = "1", features = ["full"] }
 pollster = "0.3"
-web-rwkv = "0.3.7"
-# web-rwkv = { git = "https://github.com/cryscan/web-rwkv", branch = "main" }
+# web-rwkv = "0.3.7"
+web-rwkv = { git = "https://github.com/cryscan/web-rwkv", branch = "main" }
 memmap = "0.7"
 bytemuck = "1"
 regex = "1.8"

diff --git a/src/main.rs b/src/main.rs
@@ -26,8 +26,8 @@ use tower_http::{cors::CorsLayer, services::ServeDir};
 use web_rwkv::{
     context::{Context, ContextBuilder, Instance},
     model::{
-        loader::Loader, FromBuilder, LayerFlags, Lora, LoraBlend, Model, ModelBuilder, ModelInfo,
-        ModelState, ModelVersion, Quantization, StateBuilder,
+        loader::Loader, FromBuilder, Lora, LoraBlend, Model, ModelBuilder, ModelInfo, ModelState,
+        ModelVersion, Quant, StateBuilder,
     },
     tokenizer::Tokenizer,
     wgpu::{Backends, PowerPreference},
@@ -235,10 +235,7 @@ where
         head_chunk_size,
         ..
     } = request;
-    let quant = match quant {
-        0 => Quantization::None,
-        x => Quantization::Int8(LayerFlags::from_bits_retain((1 << x) - 1)),
-    };
+    let quant = (0..quant).map(|layer| (layer, Quant::Int8)).collect();
 
     let lora: Vec<Lora> = lora
         .into_iter()