style: 简化 test-utils

YdrMaster · YdrMaster · commit d25e5d10c04a · 2025-02-11T10:46:35.000+08:00
Signed-off-by: YdrMaster &lt;ydrml@hotmail.com&gt;
diff --git a/Cargo.toml b/Cargo.toml
@@ -20,12 +20,15 @@ members = [
 ]
 resolver = "2"
 
+[workspace.package]
+edition = "2021"
+
 [workspace.dependencies]
 common.path = "common"
 gguf.path = "gguf"
 tensor.path = "tensor"
 causal-lm.path = "causal-lm"
-test-utils = { path = "test-utils", default-features = false }
+test-utils.path = "test-utils"
 
 ggus = "0.4"
 ndarray-layout = "0.0"
diff --git a/common/Cargo.toml b/common/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "common"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [dependencies]
diff --git a/gguf/Cargo.toml b/gguf/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "gguf"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [dependencies]
diff --git a/models/clip/common-cpu/Cargo.toml b/models/clip/common-cpu/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "clip-cpu"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/models/clip/common/Cargo.toml b/models/clip/common/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "clip"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/models/gpt2/common-cpu/Cargo.toml b/models/gpt2/common-cpu/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "gpt2-cpu"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["onenewcode <lovestudy@qq.com>", "YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -12,6 +12,6 @@ common.workspace = true
 operators = { workspace = true, features = ["common-cpu"] }
 
 [dev-dependencies]
-test-utils = { workspace = true, features = ["gpt2"] }
+test-utils.workspace = true
 gguf.workspace = true
 regex.workspace = true
diff --git a/models/gpt2/common-cpu/src/infer.rs b/models/gpt2/common-cpu/src/infer.rs
@@ -138,7 +138,17 @@ fn test_infer() {
 
         let senders = senders.into_boxed_slice();
         barrier.wait();
-        test_infer_paralle(&model, senders, eos, tokenizer, &prompt, max_steps)
+        test_infer_paralle(
+            senders,
+            test_utils::AboutToken {
+                tokenizer,
+                token_embd: model.token_embd,
+                nvoc: model.meta.nvoc,
+                eos,
+            },
+            &prompt,
+            max_steps,
+        )
     })
 }
 
diff --git a/models/gpt2/common/Cargo.toml b/models/gpt2/common/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "gpt2"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["onenewcode <lovestudy@qq.com>", "YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/models/gpt2/cuda/Cargo.toml b/models/gpt2/cuda/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "gpt2-cuda"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["onenewcode <lovestudy@qq.com>", "YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -18,6 +18,6 @@ search-cuda-tools.workspace = true
 search-corex-tools.workspace = true
 
 [dev-dependencies]
-test-utils = { workspace = true, features = ["gpt2"] }
+test-utils.workspace = true
 gguf.workspace = true
 regex.workspace = true
diff --git a/models/llama/common-cpu/Cargo.toml b/models/llama/common-cpu/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "llama-cpu"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -12,6 +12,6 @@ common.workspace = true
 operators = { workspace = true, features = ["common-cpu"] }
 
 [dev-dependencies]
-test-utils = { workspace = true, features = ["llama"] }
+test-utils.workspace = true
 gguf.workspace = true
 regex.workspace = true
diff --git a/models/llama/common-cpu/src/infer.rs b/models/llama/common-cpu/src/infer.rs
@@ -146,6 +146,16 @@ fn test_infer() {
 
         let senders = senders.into_boxed_slice();
         barrier.wait();
-        test_infer_paralle(&model, senders, eos, tokenizer, &prompt, max_steps)
+        test_infer_paralle(
+            senders,
+            test_utils::AboutToken {
+                tokenizer,
+                token_embd: model.token_embd,
+                nvoc: model.meta.nvoc,
+                eos,
+            },
+            &prompt,
+            max_steps,
+        )
     })
 }
diff --git a/models/llama/common/Cargo.toml b/models/llama/common/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "llama"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/models/llama/cuda/Cargo.toml b/models/llama/cuda/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "llama-cuda"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -18,6 +18,6 @@ search-cuda-tools.workspace = true
 search-corex-tools.workspace = true
 
 [dev-dependencies]
-test-utils = { workspace = true, features = ["llama"] }
+test-utils.workspace = true
 gguf.workspace = true
 regex.workspace = true
diff --git a/models/llama/cuda/src/nccl_parallel.rs b/models/llama/cuda/src/nccl_parallel.rs
@@ -181,6 +181,16 @@ fn test_infer() {
 
         let senders = senders.into_boxed_slice();
         barrier.wait();
-        test_infer_paralle(&model, senders, eos, tokenizer, &prompt, max_steps)
+        test_infer_paralle(
+            senders,
+            test_utils::AboutToken {
+                tokenizer,
+                token_embd: model.token_embd,
+                nvoc: model.meta.nvoc,
+                eos,
+            },
+            &prompt,
+            max_steps,
+        )
     })
 }
diff --git a/models/llama/infini/Cargo.toml b/models/llama/infini/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "llama-infini"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -17,6 +17,6 @@ build-script-cfg.workspace = true
 search-infini-tools.workspace = true
 
 [dev-dependencies]
-test-utils = { workspace = true, features = ["llama"] }
+test-utils.workspace = true
 gguf.workspace = true
 regex.workspace = true
diff --git a/models/llama/infini/src/infer.rs b/models/llama/infini/src/infer.rs
@@ -174,6 +174,16 @@ fn test_infer() {
 
         let senders = senders.into_boxed_slice();
         barrier.wait();
-        test_infer_paralle(&model, senders, eos, tokenizer, &prompt, max_steps)
+        test_infer_paralle(
+            senders,
+            test_utils::AboutToken {
+                tokenizer,
+                token_embd: model.token_embd,
+                nvoc: model.meta.nvoc,
+                eos,
+            },
+            &prompt,
+            max_steps,
+        )
     })
 }
diff --git a/models/llama/opencl/Cargo.toml b/models/llama/opencl/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "llama-cl"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/tensor/Cargo.toml b/tensor/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "tensor"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [dependencies]
diff --git a/test-utils/Cargo.toml b/test-utils/Cargo.toml
@@ -1,17 +1,11 @@
 [package]
 name = "test-utils"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
-[features]
-default = ["llama","gpt2"]
-
 [dependencies]
-llama = { path = "../models/llama/common", optional = true }
-gpt2 = { path = "../models/gpt2/common", optional = true }
 common.workspace = true
 gguf.workspace = true
-tensor.workspace = true
 env_logger.workspace = true
 cli-table = "0.4.9"
diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs
@@ -11,41 +11,7 @@ use std::{
     sync::{mpsc, Once},
     time::{Duration, Instant},
 };
-#[cfg(feature = "llama")]
-mod llama {
-    use crate::InferStorage;
-    use llama::LlamaStorage;
-    use tensor::Tensor;
 
-    impl InferStorage for &LlamaStorage<&[u8]> {
-        fn embd(&self, nt: usize) -> Tensor<usize> {
-            self.meta.embd(nt)
-        }
-        fn token_embd(&self) -> &[u8] {
-            self.token_embd
-        }
-    }
-}
-#[cfg(feature = "gpt2")]
-mod gpt2 {
-    use crate::InferStorage;
-    use gpt2::GPT2Storage;
-    use tensor::Tensor;
-
-    impl InferStorage for &GPT2Storage<&[u8]> {
-        fn embd(&self, nt: usize) -> Tensor<usize> {
-            self.meta.embd(nt)
-        }
-        fn token_embd(&self) -> &[u8] {
-            self.token_embd
-        }
-    }
-}
-
-pub trait InferStorage {
-    fn embd(&self, nt: usize) -> Tensor<usize>;
-    fn token_embd(&self) -> &[u8];
-}
 pub struct Inference {
     pub model: Box<[Mmap]>,
     pub devices: Option<String>,
@@ -70,7 +36,6 @@ mod env {
     pub const ROLL_CACHE_SIZE: &str = "ROLL_CACHE_SIZE";
 }
 use env::*;
-use tensor::Tensor;
 
 impl Inference {
     pub fn load() -> Option<Self> {
@@ -232,23 +197,32 @@ pub fn test_infer(
     }
 }
 
+pub struct AboutToken<'a> {
+    pub tokenizer: Tokenizer,
+    pub token_embd: &'a [u8],
+    pub nvoc: usize,
+    pub eos: utok,
+}
+
 pub fn test_infer_paralle(
-    model: impl InferStorage,
     senders: Box<[mpsc::Sender<Task>]>,
-    eos: utok,
-    tokenizer: Tokenizer,
+    about_token: AboutToken,
     prompt: &str,
     max_steps: usize,
 ) {
-    use tensor::Blob;
-
     let (next, next_recv) = mpsc::channel();
+    let AboutToken {
+        tokenizer,
+        token_embd,
+        nvoc,
+        eos,
+    } = about_token;
+    let d = token_embd.len() / nvoc;
+    let mut embd = Vec::new();
     test_infer(eos, tokenizer, prompt, max_steps, |input, pos| {
-        let mut embd = model.embd(input.len()).map(Blob::new).take();
-
-        let d = embd.len() / input.len();
+        embd.resize(d * input.len(), 0);
         for (i, &tok) in input.iter().enumerate() {
-            embd[i * d..][..d].copy_from_slice(&model.token_embd()[tok as usize * d..][..d]);
+            embd[i * d..][..d].copy_from_slice(&token_embd[tok as usize * d..][..d]);
         }
 
         for sender in &senders {
