add LFM2-1.2B, LFM2.5-1.2B-Instruct

tests/config_tests.rs

@@ -1,8 +1,5 @@
 use aha::models::{
-    deepseek_ocr::config::DeepseekOCRConfig, hunyuan_ocr::config::HunYuanVLConfig,
-    minicpm4::config::MiniCPM4Config, paddleocr_vl::config::PaddleOCRVLConfig,
-    qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig,
-    voxcpm::config::VoxCPMConfig,
+    deepseek_ocr::config::DeepseekOCRConfig, hunyuan_ocr::config::HunYuanVLConfig, lfm2::config::Lfm2Config, minicpm4::config::MiniCPM4Config, paddleocr_vl::config::PaddleOCRVLConfig, qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig
 };
 use anyhow::Result;
 
@@ -85,3 +82,16 @@ fn paddleocr_vl_config() -> Result<()> {
     println!("{:?}", config);
     Ok(())
 }
+
+#[test]
+fn lfm2_config() -> Result<()> {
+    // cargo test -F cuda --test config_tests lfm2_config -r -- --nocapture
+    let model_path = "/home/jhq/.aha/LiquidAI/LFM2-1.2B/";
+    // let model_path = "/home/jhq/.aha/LiquidAI/LFM2.5-1.2B-Instruct/";
+    let config_path = model_path.to_string() + "/config.json";
+    let mut config: Lfm2Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
+    println!("{:?}", config);
+    config.full_attn_idx2layer_type();
+    println!("{:?}", config);
+    Ok(())
+}

tests/messy_test.rs

@@ -35,14 +35,18 @@ async fn download_test() -> Result<()> {
 #[test]
 fn messy_test() -> Result<()> {
     // RUST_BACKTRACE=1 cargo test -F cuda --test messy_test messy_test -r -- --nocapture
-    let save_dir =
-        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
-    let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir);
-    let stem = std::path::Path::new(&model_path)
-        .file_stem() // 获取文件名主干（不含扩展名）
-        .and_then(|s| s.to_str())
-        .unwrap_or("qwen3.5");
-    println!("stem: {:?}", stem);
+    // let t1 = Tensor::randn(0.0, 1.0, (1, 2, 6), device)?;
+    // println!(" t1: {}", t1);
+    // let t2 = t1.pad_with_zeros(D::Minus1, -3, 0)?;
+    // println!(" t2: {}", t);
+    // let save_dir =
+    //     aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    // let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir);
+    // let stem = std::path::Path::new(&model_path)
+    //     .file_name() 
+    //     .and_then(|s| s.to_str())
+    //     .unwrap_or("qwen3.5");
+    // println!("stem: {:?}", stem);
     // let device = &candle_core::Device::Cpu;
     // let t1 = Tensor::randn(0.0, 1.0, (16, 9, 64, 128), device)?;
     // let t2 = Tensor::randn(0.0, 1.0, (16, 9, 128, 64), device)?;

tests/test_lfm2.rs

@@ -0,0 +1,91 @@
+use std::{pin::pin, time::Instant};
+use anyhow::Result;
+use aha::{chat::ChatCompletionParameters, models::{GenerateModel, lfm2::generate::Lfm2GenerateModel}};
+use rocket::futures::StreamExt;
+
+#[test]
+fn lfm2_generate() -> Result<()> {
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_lfm2 lfm2_generate -r -- --nocapture
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen3_0_6b_generate -r -- --nocapture
+
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    // let model_path = format!("{}/LiquidAI/LFM2-1.2B/", save_dir);
+    let model_path = format!("{}/LiquidAI/LFM2.5-1.2B-Instruct/", save_dir);
+    let message = r#"
+    {
+        "model": "lfm2",
+        "messages": [
+            {
+                "role": "user",
+                "content": "你如何看待AI"
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut model = Lfm2GenerateModel::init(&model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+
+    let i_start = Instant::now();
+    let result = model.generate(mes)?;
+    let i_duration = i_start.elapsed();
+    println!("generate: \n {:?}", result);
+    if let Some(usage) = &result.usage {
+        let num_token = usage.total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
+    println!("Time elapsed in generate is: {:?}", i_duration);
+
+    Ok(())
+}
+
+
+#[tokio::test]
+async fn lfm2_stream() -> Result<()> {
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_lfm2 lfm2_stream -r -- --nocapture
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen3_0_6b_generate -r -- --nocapture
+
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    // let model_path = format!("{}/LiquidAI/LFM2-1.2B/", save_dir);
+    let model_path = format!("{}/LiquidAI/LFM2.5-1.2B-Instruct/", save_dir);
+    let message = r#"
+    {
+        "model": "lfm2",
+        "messages": [
+            {
+                "role": "user",
+                "content": "你如何看待AI"
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut model = Lfm2GenerateModel::init(&model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+
+    let i_start = Instant::now();
+    // let result = model.generate(mes)?;
+    let mut stream = pin!(model.generate_stream(mes)?);
+    let i_duration = i_start.elapsed();
+    while let Some(token) = stream.next().await {
+        println!("generate: \n {:?}", token);
+    }
+    // println!("generate: \n {:?}", result);
+    // if let Some(usage) = &result.usage {
+    //     let num_token = usage.total_tokens;
+    //     let duration_secs = i_duration.as_secs_f64();
+    //     let tps = num_token as f64 / duration_secs;
+    //     println!("Tokens per second (TPS): {:.2}", tps);
+    // }
+    println!("Time elapsed in generate is: {:?}", i_duration);
+
+    Ok(())
+}

tests/test_qwen3.rs

@@ -7,7 +7,7 @@ use rocket::futures::StreamExt;
 
 #[test]
 fn qwen3_0_6b_generate() -> Result<()> {
-    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3_0_6b_generate -r -- --nocapture
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_qwen3 qwen3_0_6b_generate -r -- --nocapture
     // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen3_0_6b_generate -r -- --nocapture
 
     let save_dir =

tests/weight_test.rs

@@ -265,3 +265,25 @@ fn deepseekocrv2_weight() -> Result<()> {
     println!("model_list: {:?}", model_list);
     Ok(())
 }
+
+#[test]
+fn lfm2_weight() -> Result<()> {
+    // cargo test -F cuda --test weight_test lfm2_weight -r -- --nocapture
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let model_path = format!("{}/LiquidAI/LFM2-1.2B/", save_dir);
+    let model_list = find_type_files(&model_path, "safetensors")?;
+
+    let device = Device::Cpu;
+    for m in &model_list {
+        let weights = safetensors::load(m, &device)?;
+        for (key, tensor) in weights.iter() {
+            // if key.contains("lm_head") {
+            //     println!("=== {} === {:?}", key, tensor.shape());
+            // }
+            println!("=== {} === {:?}", key, tensor.shape());
+        }
+    }
+    println!("model_list: {:?}", model_list);
+    Ok(())
+}