test: 测试现在从环境变量获取模型路径

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Apr 12, 2024 · f89a3fc · f89a3fc
1 parent 75eca4f
commit f89a3fc
Show file tree

Hide file tree

Showing 11 changed files with 92 additions and 17 deletions.
diff --git a/common/src/lib.rs b/common/src/lib.rs
@@ -11,5 +11,6 @@ pub type utok = u32;
 pub type upos = u32;
 
 mod blob;
+pub mod test_model;
 
 pub use blob::Blob;
diff --git a/common/src/test_model.rs b/common/src/test_model.rs
@@ -0,0 +1,38 @@
+//!
+
+use std::{
+    env::var_os,
+    fs::canonicalize,
+    path::{Path, PathBuf},
+    process::Command,
+    str::from_utf8,
+};
+
+///
+pub fn find() -> Option<PathBuf> {
+    let model = var_os("TEST_MODEL")?;
+    if let Ok(path) = canonicalize(&model) {
+        return Some(path.into());
+    }
+
+    let output = Command::new(env!("CARGO"))
+        .arg("locate-project")
+        .arg("--workspace")
+        .arg("--message-format=plain")
+        .output()
+        .unwrap()
+        .stdout;
+    let workspace = Path::new(from_utf8(&output).unwrap().trim())
+        .parent()
+        .unwrap();
+    if let Ok(path) = canonicalize(workspace.join(model)) {
+        return Some(path.into());
+    }
+
+    None
+}
+
+#[test]
+fn test_find() {
+    println!("{:?}", find());
+}
diff --git a/nvidia/common/src/lib.rs b/nvidia/common/src/lib.rs
@@ -12,7 +12,7 @@ mod rms_norm;
 mod rotary_embedding;
 mod swiglu;
 
-pub use common::utok;
+pub use common::{test_model, utok};
 pub use tensor::{slice, udim, DataType, LocalSplitable, Tensor};
 
 use cublas::{Cublas, CublasSpore};

diff --git a/nvidia/distributed/src/lib.rs b/nvidia/distributed/src/lib.rs
@@ -249,6 +249,11 @@ fn test() {
     use simple_logger::SimpleLogger;
     use transformer::Transformer as _;
 
+    let Some(model_dir) = common_nv::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
     const N: usize = 1;
 
     cuda::init();
@@ -259,10 +264,7 @@ fn test() {
     SimpleLogger::new().with_level(Trace).init().unwrap();
 
     let time = Instant::now();
-    let transformer = Transformer::new(
-        "../../../TinyLlama-1.1B-Chat-v1.0_F16",
-        &[Device::fetch().unwrap()],
-    );
+    let transformer = Transformer::new(model_dir, &[Device::fetch().unwrap()]);
     info!("load {:?}", time.elapsed());
 
     let time = Instant::now();

diff --git a/nvidia/distributed/src/parameters.rs b/nvidia/distributed/src/parameters.rs
@@ -147,6 +147,11 @@ fn test_load() {
     use std::io::ErrorKind::NotFound;
     use transformer::{Memory, SafeTensorError};
 
+    let Some(model_dir) = common_nv::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
     const N: usize = 1;
 
     cuda::init();
@@ -157,7 +162,7 @@ fn test_load() {
     SimpleLogger::new().with_level(Trace).init().unwrap();
 
     let time = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir("../../../TinyLlama-1.1B-Chat-v1.0_F16");
+    let safetensors = Memory::load_safetensors_from_dir(model_dir);
     info!("mmap {:?}", time.elapsed());
 
     let model = match safetensors {

diff --git a/service/src/lib.rs b/service/src/lib.rs
@@ -118,14 +118,13 @@ fn tokenizer(model_dir: impl AsRef<Path>) -> Box<dyn Tokenizer + Send + Sync> {
 #[test]
 fn test() {
     use colored::{Color, Colorize};
-    use std::{io::Write, iter::zip, path::Path, time::Duration};
+    use std::{io::Write, iter::zip, time::Duration};
     use tokio::{runtime::Builder, task::JoinSet, time::sleep};
 
-    let model_dir = "../../TinyLlama-1.1B-Chat-v1.0_F16/";
-    if !Path::new(model_dir).exists() {
-        println!("model not exist");
+    let Some(model_dir) = common::test_model::find() else {
         return;
-    }
+    };
+    println!("model_dir: {}", model_dir.display());
 
     let runtime = Builder::new_current_thread().enable_time().build().unwrap();
     let _rt = runtime.enter();

diff --git a/tokenizer/src/bpe.rs b/tokenizer/src/bpe.rs
@@ -142,7 +142,12 @@ impl Tokenizer for BPE {
 
 #[test]
 fn read_tokenizer() {
-    if let Ok(bpe) = BPE::from_model_file("../../TinyLlama-1.1B-Chat-v1.0/tokenizer.model") {
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
+    if let Ok(bpe) = BPE::from_model_file(model_dir.join("tokenizer.model")) {
         for i in 0..bpe.offsets.len() {
             println!("{}: {}", bpe.get_piece(i as utok), bpe.get_score(i as utok));
         }
@@ -151,8 +156,13 @@ fn read_tokenizer() {
 
 #[test]
 fn once_upon_a_time() {
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
     use std::time::Instant;
-    if let Ok(bpe) = BPE::from_model_file("../../TinyLlama-1.1B-Chat-v1.0/tokenizer.model") {
+    if let Ok(bpe) = BPE::from_model_file(model_dir.join("tokenizer.model")) {
         const PROMPT: &str = "▁Once▁upon▁a▁time,";
         let tokens = bpe.encode(PROMPT);
         let t0 = Instant::now();

diff --git a/transformer-cpu/src/lib.rs b/transformer-cpu/src/lib.rs
@@ -325,8 +325,13 @@ fn test_build() {
     use std::{io::ErrorKind::NotFound, time::Instant};
     use transformer::{Memory, SafeTensorError};
 
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
     let t0 = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir("../../TinyLlama-1.1B-Chat-v1.0_F16");
+    let safetensors = Memory::load_safetensors_from_dir(model_dir);
     let t1 = Instant::now();
     println!("mmap {:?}", t1 - t0);
 

diff --git a/transformer/src/parameters/distribute.rs b/transformer/src/parameters/distribute.rs
@@ -286,8 +286,13 @@ fn test() {
     use super::{Memory, SafeTensorError};
     use std::{io::ErrorKind::NotFound, time::Instant};
 
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
     let time = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir("../../TinyLlama-1.1B-Chat-v1.0");
+    let safetensors = Memory::load_safetensors_from_dir(model_dir);
     println!("mmap {:?}", time.elapsed());
 
     let model = match safetensors {

diff --git a/transformer/src/parameters/memory.rs b/transformer/src/parameters/memory.rs
@@ -183,8 +183,13 @@ fn test_load() {
     use super::SafeTensorError;
     use std::{io::ErrorKind::NotFound, time::Instant};
 
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
     let t0 = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir("../../TinyLlama-1.1B-Chat-v1.0");
+    let safetensors = Memory::load_safetensors_from_dir(model_dir);
     let t1 = Instant::now();
     println!("mmap {:?}", t1 - t0);
 

diff --git a/transformer/src/parameters/safe_tensors.rs b/transformer/src/parameters/safe_tensors.rs
@@ -155,7 +155,12 @@ fn concat0(tensors: &[&Tensor<Storage>]) -> Tensor<Storage> {
 
 #[test]
 fn test_load() {
-    let file = match std::fs::File::open("../../TinyLlama-1.1B-Chat-v1.0/model.safetensors") {
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
+    let file = match std::fs::File::open(model_dir.join("model.safetensors")) {
         Ok(f) => f,
         Err(e) if e.kind() == std::io::ErrorKind::NotFound => return,
         Err(e) => panic!("{e:?}"),