diff --git a/Cargo.lock b/Cargo.lock index 3c56c66..58f18ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -21,7 +21,7 @@ dependencies = [ [[package]] name = "aha" -version = "0.2.3" +version = "0.2.4" dependencies = [ "aha_openai_dive", "ahash", diff --git a/Cargo.toml b/Cargo.toml index f3b6db0..ab4596c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,10 +1,10 @@ [package] name = "aha" -version = "0.2.3" +version = "0.2.4" edition = "2024" repository = "https://github.com/jhqxxx/aha" license = "Apache-2.0" -description = "aha model inference library, now supports Qwen(2.5VL/3/3VL/3.5/ASR), MiniCPM4, VoxCPM/1.5, DeepSeek-OCR/2, Hunyuan-OCR, PaddleOCR-VL/1.5, RMBG2.0, GLM(ASR-Nano-2512/OCR), Fun-ASR-Nano-2512" +description = "aha model inference library, now supports Qwen(2.5VL/3/3VL/3.5/ASR), MiniCPM4, VoxCPM/1.5, DeepSeek-OCR/2, Hunyuan-OCR, PaddleOCR-VL/1.5, RMBG2.0, GLM(ASR-Nano-2512/OCR), Fun-ASR-Nano-2512, LFM(2/2.5)" [dependencies] candle-core = { version = "0.9.2" } diff --git a/README.md b/README.md index 5bf72ac..5ab8250 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,10 @@ aha is a high-performance, cross-platform AI inference engine built with Rust and the Candle framework. It brings state-of-the-art AI models to your local machine—no API keys, no cloud dependencies, just pure, fast AI running directly on your hardware. ## Changelog +### v0.2.4 (2026-03-23) +- add LFM2.5-1.2B-Instruct +- add LFM2-1.2B + ### v0.2.3 (2026-03-18) - add DeepSeek-OCR-2 @@ -49,9 +53,6 @@ aha is a high-performance, cross-platform AI inference engine built with Rust an ### v0.2.1 (2026-03-05) - Added Qwen3.5 model -### 2026-03-01 -- update interpolate.rs - **[View full changelog](docs/changelog.md)** → @@ -125,9 +126,9 @@ curl http://localhost:10100/chat/completions \ | Category | Models | |----------|--------| -| **Text** | Qwen3, MiniCPM4 | -| **Vision** | Qwen2.5-VL, Qwen3-VL | -| **OCR** | DeepSeek-OCR, Hunyuan-OCR, PaddleOCR-VL | +| **Text** | Qwen3, MiniCPM4,
LFM2-1.2B, LFM2.5-1.2B-Instruct | +| **Vision** | Qwen2.5-VL, Qwen3-VL, Qwen3.5 | +| **OCR** | DeepSeek-OCR, DeepSeek-OCR-2 ,
Hunyuan-OCR, PaddleOCR-VL, PaddleOCR-VL1.5 | | **ASR** | GLM-ASR-Nano, Fun-ASR-Nano, Qwen3-ASR | | **Audio** | VoxCPM, VoxCPM1.5 | | **Image** | RMBG-2.0 (background removal) | diff --git a/README.zh-CN.md b/README.zh-CN.md index b7b0f8d..00bd0fb 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -25,6 +25,10 @@ aha 是一款基于 Rust 和 Candle 框架构建的高性能跨平台 AI 推理引擎。将最先进的 AI 模型带到您的本地机器——无需 API 密钥，无需云依赖，纯粹、快速的 AI，直接在您的硬件上运行。 ## 更新日志 +### v0.2.4 (2026-03-23) +- 新增 LFM2.5-1.2B-Instruct +- 新增 LFM2-1.2B + ### v0.2.3 (2026-03-18) - 新增 DeepSeek-OCR-2 @@ -124,9 +128,9 @@ curl http://localhost:10100/chat/completions \ | 类别 | 模型 | |------|------| -| **文本** | Qwen3, MiniCPM4 | +| **文本** | Qwen3, MiniCPM4,
LFM2-1.2B, LFM2.5-1.2B-Instruct | | **视觉** | Qwen2.5-VL, Qwen3-VL | -| **OCR** | DeepSeek-OCR, Hunyuan-OCR, PaddleOCR-VL | +| **OCR** | DeepSeek-OCR, DeepSeek-OCR-2 ,
Hunyuan-OCR, PaddleOCR-VL, PaddleOCR-VL1.5 | | **ASR** | GLM-ASR-Nano, Fun-ASR-Nano,Qwen3-ASR | | **音频** | VoxCPM, VoxCPM1.5 | | **图像** | RMBG-2.0 (背景移除) | diff --git a/docs/changelog.md b/docs/changelog.md index 4c47692..f3ea4e3 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -5,6 +5,10 @@ All notable changes to aha will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +### v0.2.4 (2026-03-23) +- add LFM2.5-1.2B-Instruct +- add LFM2-1.2B + ### v0.2.3 (2026-03-18) - add DeepSeek-OCR-2 diff --git a/docs/changelog.zh-CN.md b/docs/changelog.zh-CN.md index 5f5a4e3..5355518 100644 --- a/docs/changelog.zh-CN.md +++ b/docs/changelog.zh-CN.md @@ -5,6 +5,10 @@ 格式基于 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.0.0/)， 本项目遵循 [语义化版本](https://semver.org/lang/zh-CN/spec/v2.0.0.html)。 +### v0.2.4 (2026-03-23) +- 新增 LFM2.5-1.2B-Instruct +- 新增 LFM2-1.2B + ### v0.2.3 (2026-03-18) - 新增 DeepSeek-OCR-2 diff --git a/docs/supported-models.md b/docs/supported-models.md index 436ba73..606b772 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -8,6 +8,8 @@ aha supports a growing collection of state-of-the-art AI models across multiple |-------|-----------|-------------|----------|---------| | **Qwen3-0.6B** | 0.6B | Latest generation | Advanced reasoning | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | | **MiniCPM4-0.5B** | 0.5B | Efficient lightweight | Edge deployment | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | +| **LFM2-1.2B** | 1.2B | Efficient lightweight | Edge deployment | [lfm1.0](https://huggingface.co/LiquidAI/LFM2-1.2B/blob/main/LICENSE) | +| **LFM2.5-1.2B-Instruct** | 1.2B | Efficient lightweight | Edge deployment | [lfm1.0](https://huggingface.co/LiquidAI/LFM2.5-1.2B-Instruct/blob/main/LICENSE) | ## Vision & Multimodal diff --git a/docs/supported-models.zh-CN.md b/docs/supported-models.zh-CN.md index f24305d..a1b26de 100644 --- a/docs/supported-models.zh-CN.md +++ b/docs/supported-models.zh-CN.md @@ -8,6 +8,8 @@ aha 支持多个领域的最先进 AI 模型集合。 |------|--------|------|----------|---------| | **Qwen3-0.6B** | 0.6B | 最新一代 | 高级推理 | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | | **MiniCPM4-0.5B** | 0.5B | 高效轻量级 | 边缘部署 | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | +| **LFM2-1.2B** | 1.2B | 高效轻量级 | 边缘部署 | [lfm1.0](https://huggingface.co/LiquidAI/LFM2-1.2B/blob/main/LICENSE) | +| **LFM2.5-1.2B-Instruct** | 1.2B | 高效轻量级 | 边缘部署 | [lfm1.0](https://huggingface.co/LiquidAI/LFM2.5-1.2B-Instruct/blob/main/LICENSE) | ## 视觉与多模态 diff --git a/src/api/mod.rs b/src/api/mod.rs index 73074b9..4ade969 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -247,6 +247,8 @@ struct ErrorResponse { fn which_model_to_id(which_model: WhichModel) -> &'static str { match which_model { WhichModel::MiniCPM4_0_5B => "minicpm4-0.5b", + WhichModel::LFM2_1_2B => "lfm2-1.2b", + WhichModel::LFM2_5_1_2BInstruct => "lfm2.5-1.2b-instruct", WhichModel::Qwen2_5vl3B => "qwen2.5vl-3b", WhichModel::Qwen2_5vl7B => "qwen2.5vl-7b", WhichModel::Qwen3_0_6B => "qwen3-0.6b", @@ -297,6 +299,7 @@ fn which_model_to_owner(which_model: WhichModel) -> &'static str { WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "OpenBMB", WhichModel::GlmASRNano2512 | WhichModel::GlmOCR => "ZhipuAI", WhichModel::FunASRNano2512 => "FunAudioLLM", + WhichModel::LFM2_1_2B | WhichModel::LFM2_5_1_2BInstruct => "LiquidAI", } } diff --git a/src/exec/lfm2.rs b/src/exec/lfm2.rs new file mode 100644 index 0000000..41986ff --- /dev/null +++ b/src/exec/lfm2.rs @@ -0,0 +1,61 @@ +//! MiniCPM4-0.5B exec implementation for CLI `run` subcommand + +use std::time::Instant; + +use anyhow::{Ok, Result}; + +use crate::exec::ExecModel; +use crate::models::GenerateModel; +use crate::models::lfm2::generate::Lfm2GenerateModel; +use crate::utils::get_file_path; + +pub struct Lfm2Exec; + +impl ExecModel for Lfm2Exec { + fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { + let input_text = &input[0]; + let target_text = if input_text.starts_with("file://") { + // let path = &input[7..]; + let path = get_file_path(input_text)?; + std::fs::read_to_string(path)? + } else { + input_text.to_string() + }; + + let i_start = Instant::now(); + let mut model = Lfm2GenerateModel::init(weight_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let message = format!( + r#"{{ + "temperature": 0.3, + "top_p": 0.8, + "model": "lfm2", + "messages": [ + {{ + "role": "user", + "content": "{}" + }} + ] + }}"#, + target_text.replace('"', "\\\"") + ); + let mes = serde_json::from_str(&message)?; + + let i_start = Instant::now(); + let result = model.generate(mes)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + + // Print result + println!("Result: {:?}", result); + + if let Some(out) = output { + std::fs::write(out, format!("{:?}", result))?; + println!("Output saved to: {}", out); + } + + Ok(()) + } +} diff --git a/src/exec/mod.rs b/src/exec/mod.rs index 1e2c08e..191024a 100644 --- a/src/exec/mod.rs +++ b/src/exec/mod.rs @@ -18,6 +18,7 @@ pub mod qwen3vl; pub mod rmbg2_0; pub mod voxcpm; pub mod voxcpm1_5; +pub mod lfm2; use anyhow::Result; diff --git a/src/main.rs b/src/main.rs index 99ae64e..299d77a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -244,6 +244,8 @@ struct ModelInfo { fn run_list(args: ListArgs) -> anyhow::Result<()> { let models = [ WhichModel::MiniCPM4_0_5B, + WhichModel::LFM2_1_2B, + WhichModel::LFM2_5_1_2BInstruct, WhichModel::Qwen2_5vl3B, WhichModel::Qwen2_5vl7B, WhichModel::Qwen3_0_6B, @@ -461,6 +463,14 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> { use aha::exec::minicpm4::MiniCPM4Exec; MiniCPM4Exec::run(&input, output.as_deref(), &weight_path)?; } + WhichModel::LFM2_1_2B => { + use aha::exec::lfm2::Lfm2Exec; + Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::LFM2_5_1_2BInstruct => { + use aha::exec::lfm2::Lfm2Exec; + Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; + } WhichModel::Qwen2_5vl3B => { use aha::exec::qwen2_5vl::Qwen2_5vlExec; Qwen2_5vlExec::run(&input, output.as_deref(), &weight_path)?; diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 8d871c2..d859dc5 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -411,6 +411,145 @@ impl QKVCatAttention { } } +pub struct QKNormAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_attention_heads: usize, + num_key_value_heads: usize, + num_kv_groups: usize, + head_dim: usize, + scaling: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl QKNormAttention { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + head_dim: Option, + num_key_value_heads: Option, + attention_bias: bool, + rms_norm_eps: f64, + q_proj_pp_name: Option<&str>, + k_proj_pp_name: Option<&str>, + v_proj_pp_name: Option<&str>, + o_proj_pp_name: Option<&str>, + q_norm_pp_name: Option<&str>, + k_norm_pp_name: Option<&str>, + ) -> Result { + let head_dim = head_dim.unwrap_or(hidden_size / num_attention_heads); + let num_key_value_heads = num_key_value_heads.unwrap_or(num_attention_heads); + let num_kv_groups = num_attention_heads / num_key_value_heads; + let scaling = 1f64 / f64::sqrt(head_dim as f64); + let q_proj_pp_name = q_proj_pp_name.unwrap_or("q_proj"); + let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj"); + let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj"); + let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); + let q_norm_pp_name = q_norm_pp_name.unwrap_or("q_norm"); + let k_norm_pp_name = k_norm_pp_name.unwrap_or("k_norm"); + let q_proj = linear_b( + hidden_size, + num_attention_heads * head_dim, + attention_bias, + vb.pp(q_proj_pp_name), + )?; + let k_proj = linear_b( + hidden_size, + num_key_value_heads * head_dim, + attention_bias, + vb.pp(k_proj_pp_name), + )?; + let v_proj = linear_b( + hidden_size, + num_key_value_heads * head_dim, + attention_bias, + vb.pp(v_proj_pp_name), + )?; + let o_proj = linear_b( + num_attention_heads * head_dim, + hidden_size, + attention_bias, + vb.pp(o_proj_pp_name), + )?; + let q_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(q_norm_pp_name))?; + let k_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(k_norm_pp_name))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + num_attention_heads, + num_key_value_heads, + num_kv_groups, + head_dim, + scaling, + kv_cache: None, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + let query_states = self.q_proj.forward(xs)?.reshape(( + b_sz, + q_len, + self.num_attention_heads, + self.head_dim, + ))?; + let query_states = self.q_norm.forward(&query_states)?.transpose(1, 2)?; + let key_states = self.k_proj.forward(xs)?.reshape(( + b_sz, + q_len, + self.num_key_value_heads, + self.head_dim, + ))?; + let key_states = self.k_norm.forward(&key_states)?.transpose(1, 2)?; + let value_states = self.v_proj.forward(xs)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + let (query_states, key_states) = + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?; + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + Some(self.num_kv_groups), + attention_mask, + self.scaling, + )?; + let attn_output = + attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + pub struct NaiveAttnTwoLinearMLPBlock { self_attn: NaiveAttention, mlp: TwoLinearMLP, diff --git a/src/models/lfm2/config.rs b/src/models/lfm2/config.rs new file mode 100644 index 0000000..8ec24c1 --- /dev/null +++ b/src/models/lfm2/config.rs @@ -0,0 +1,88 @@ +use serde::{Deserialize, Serialize}; +use anyhow::{anyhow, Result}; +#[derive(Debug, PartialEq, Deserialize, Serialize)] +pub struct Lfm2Config { + pub architectures: Vec, + pub block_auto_adjust_ff_dim: bool, + pub block_dim: usize, + pub block_ff_dim: usize, + pub block_ffn_dim_multiplier: f64, + pub block_mlp_init_scale: f64, + pub block_multiple_of: usize, + pub block_norm_eps: f64, + pub block_out_init_scale: f64, + pub block_use_swiglu: bool, + pub block_use_xavier_init: bool, + pub bos_token_id: u32, + #[serde[rename="conv_L_cache"]] + pub conv_l_cache: usize, + pub conv_bias: bool, + pub conv_dim: usize, + pub conv_dim_out: Option, + pub conv_use_xavier_init: bool, + pub eos_token_id: u32, + pub full_attn_idxs: Option>, + pub layer_types: Option>, + pub hidden_size: usize, + pub initializer_range: f64, + pub intermediate_size: Option, + pub max_position_embeddings: usize, + pub model_type: String, + pub norm_eps: f64, + pub num_attention_heads: usize, + pub num_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub pad_token_id: u32, + pub rope_theta: f32, + pub torch_dtype: Option, + pub dtype: Option, + pub use_cache: bool, + pub use_pos_enc: bool, + pub vocab_size: usize, + pub tie_embedding: Option, +} + + +impl Lfm2Config { + pub fn full_attn_idx2layer_type(&mut self) { + if self.layer_types.is_none() + && let Some(idx) = &self.full_attn_idxs + { + let mut layer_types = vec![]; + for i in 0..self.num_hidden_layers { + if idx.contains(&i) { + layer_types.push("full_attention".to_string()); + } else { + layer_types.push("conv".to_string()); + } + } + self.layer_types = Some(layer_types); + } + } + + pub fn get_layer_types(&self) -> Result> { + if let Some(types) = &self.layer_types { + Ok(types.clone()) + } else if let Some(idx) = &self.full_attn_idxs { + let mut layer_types = vec![]; + for i in 0..self.num_hidden_layers { + if idx.contains(&i) { + layer_types.push("full_attention".to_string()); + } else { + layer_types.push("conv".to_string()); + } + } + Ok(layer_types) + } else { + Err(anyhow!("layer_types full_attn_idxs cannot be none at the same time")) + } + } +} + +#[derive(Debug, PartialEq, Deserialize, Serialize)] +pub struct Lfm2GenerateConfig { + pub bos_token_id: u32, + pub eos_token_id: u32, + pub pad_token_id: u32 +} \ No newline at end of file diff --git a/src/models/lfm2/generate.rs b/src/models/lfm2/generate.rs new file mode 100644 index 0000000..1c58aff --- /dev/null +++ b/src/models/lfm2/generate.rs @@ -0,0 +1,163 @@ +use crate::{ + chat_template::ChatTemplate, + models::{GenerateModel, lfm2::{ + config::{Lfm2Config, Lfm2GenerateConfig}, + model::Lfm2Model, + }}, + tokenizer::TokenizerModel, + utils::{ + build_completion_response, find_type_files, get_device, get_dtype, get_logit_processor, + }, +}; +use crate::utils::build_completion_chunk_response; +use aha_openai_dive::v1::resources::chat::{ChatCompletionParameters, ChatCompletionResponse}; +use anyhow::Result; +use candle_core::{DType, Device, Tensor}; +use candle_nn::VarBuilder; +use rocket::async_stream::stream; + +pub struct Lfm2GenerateModel<'a> { + chat_template: ChatTemplate<'a>, + tokenizer: TokenizerModel, + device: Device, + model: Lfm2Model, + eos_token_id: u32, + model_name: String, +} +impl<'a> Lfm2GenerateModel<'a> { + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let chat_template = ChatTemplate::init(path)?; + let tokenizer = TokenizerModel::init(path)?; + let device = get_device(device); + let gen_cfg_path = path.to_string() + "/generation_config.json"; + let gen_cfg: Lfm2GenerateConfig = serde_json::from_slice(&std::fs::read(gen_cfg_path)?)?; + let cfg_path = path.to_string() + "/config.json"; + let cfg: Lfm2Config = serde_json::from_slice(&std::fs::read(cfg_path)?)?; + let model_path = find_type_files(path, "safetensors")?; + let cfg_dtype = if let Some(dtype) = &cfg.dtype { + dtype.clone() + } else if let Some(dtype) = &cfg.torch_dtype { + dtype.clone() + } else { + "bfloat16".to_string() + }; + let dtype = get_dtype(dtype, &cfg_dtype); + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_path, dtype, &device)? }; + let model = Lfm2Model::new(vb, &cfg)?; + let eos_token_id = gen_cfg.eos_token_id; + let model_name = std::path::Path::new(path) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or("lfm2") + .to_string(); + Ok(Self { + chat_template, + tokenizer, + device, + model, + eos_token_id, + model_name, + }) + } + + +} + +impl<'a> GenerateModel for Lfm2GenerateModel<'a> { + fn generate(&mut self, mes: ChatCompletionParameters) -> Result { + let mes_render = self.chat_template.apply_chat_template(&mes)?; + let mut logits = get_logit_processor( + mes.temperature, + mes.top_p, + None, + mes.seed.unwrap_or(34562) as u64, + ); + let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?; + let mut seq_len = input_ids.dim(1)?; + let prompt_tokens = seq_len as u32; + let mut seqlen_offset = 0; + let mut generate = vec![]; + let sample_len = mes.max_tokens.unwrap_or(1024); + for _ in 0..sample_len { + let logit = self.model.forward(&input_ids, seqlen_offset)?; + let logit = logit.squeeze(0)?.squeeze(0)?; + let next_token = logits.sample(&logit)?; + generate.push(next_token); + if next_token == self.eos_token_id { + break; + } + input_ids = Tensor::new(vec![next_token], &self.device)?.unsqueeze(0)?; + seqlen_offset += seq_len; + seq_len = 1; + } + self.model.clear_cache(); + let completion_tokens = generate.len() as u32; + let decode = self.tokenizer.token_decode(generate)?; + let mes = build_completion_response( + decode, + &self.model_name, + Some(completion_tokens), + Some(prompt_tokens), + ); + Ok(mes) + } + + fn generate_stream( + &mut self, + mes: ChatCompletionParameters, + ) -> Result< + Box< + dyn rocket::futures::Stream> + + Send + + Unpin + + '_, + >, + > { + let mes_render = self.chat_template.apply_chat_template(&mes)?; + let mut logits = get_logit_processor( + mes.temperature, + mes.top_p, + None, + mes.seed.unwrap_or(34562) as u64, + ); + let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?; + let mut seq_len = input_ids.dim(1)?; + let mut seqlen_offset = 0; + let sample_len = mes.max_tokens.unwrap_or(1024); + let stream = stream! { + let mut err_tokens = vec![]; + for _ in 0..sample_len { + let logit = self.model.forward(&input_ids, seqlen_offset)?; + let logit = logit.squeeze(0)?.squeeze(0)?; + let next_token = logits.sample(&logit)?; + let mut decode_ids = vec![]; + if !err_tokens.is_empty() { + decode_ids.extend_from_slice(&err_tokens); + } + decode_ids.push(next_token); + let decode = self.tokenizer.token_decode(decode_ids)?; + if decode.contains("�") { + err_tokens.push(next_token); + if err_tokens.len() > 3 { + err_tokens.clear(); + } + input_ids = Tensor::new(vec![next_token], &self.device)?.unsqueeze(0)?; + seqlen_offset += seq_len; + seq_len = 1; + continue; + } + err_tokens.clear(); + let chunk = build_completion_chunk_response(decode, &self.model_name, None, None); + yield Ok(chunk); + if next_token == self.eos_token_id { + break; + } + input_ids = Tensor::new(vec![next_token], &self.device)?.unsqueeze(0)?; + seqlen_offset += seq_len; + seq_len = 1; + } + self.model.clear_cache(); + }; + Ok(Box::new(Box::pin(stream))) + } +} \ No newline at end of file diff --git a/src/models/lfm2/mod.rs b/src/models/lfm2/mod.rs new file mode 100644 index 0000000..de0a4df --- /dev/null +++ b/src/models/lfm2/mod.rs @@ -0,0 +1,3 @@ +pub mod config; +pub mod model; +pub mod generate; \ No newline at end of file diff --git a/src/models/lfm2/model.rs b/src/models/lfm2/model.rs new file mode 100644 index 0000000..db7d4d1 --- /dev/null +++ b/src/models/lfm2/model.rs @@ -0,0 +1,313 @@ +use crate::{ + models::{ + common::{GateUpDownMLP, QKNormAttention, conv1d_depthwise, get_conv1d}, + lfm2::config::Lfm2Config, + }, + position_embed::rope::RoPE, + utils::tensor_utils::prepare_causal_attention_mask, +}; +use anyhow::Result; +use candle_core::{D, Tensor}; +use candle_nn::{ + Conv1d, Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, linear_b, rms_norm, +}; + +pub struct Lfm2ShortConv { + l_cache: usize, + conv: Conv1d, + in_proj: Linear, + out_proj: Linear, + cache: Option, +} + +impl Lfm2ShortConv { + pub fn new(vb: VarBuilder, config: &Lfm2Config) -> Result { + let l_cache = config.conv_l_cache; + let bias = config.conv_bias; + let conv = get_conv1d( + vb.pp("conv"), + config.hidden_size, + config.hidden_size, + l_cache, + l_cache - 1, + 1, + 1, + config.hidden_size, + bias, + )?; + let in_proj = linear_b( + config.hidden_size, + config.hidden_size * 3, + bias, + vb.pp("in_proj"), + )?; + let out_proj = linear_b( + config.hidden_size, + config.hidden_size, + bias, + vb.pp("out_proj"), + )?; + Ok(Self { + l_cache, + conv, + in_proj, + out_proj, + cache: None, + }) + } + + pub fn forward(&mut self, xs: &Tensor) -> Result { + let seq_len = xs.dim(1)?; + let bc_x = self.in_proj.forward(xs)?.transpose(D::Minus1, D::Minus2)?; + let chunk = bc_x.chunk(3, D::Minus2)?; + let bx = chunk[0].mul(&chunk[2])?; + let c: &Tensor = &chunk[1]; + let conv_out = if self.cache.is_none() && seq_len > 1 { + let pad_num = self.l_cache as isize - seq_len as isize; + let conv_state = if pad_num > 0 { + bx.pad_with_zeros(D::Minus1, pad_num as usize, 0)? + } else { + bx.narrow(D::Minus1, pad_num.unsigned_abs(), self.l_cache)? + }; + self.cache = Some(conv_state); + let bx = bx.pad_with_zeros(D::Minus1, self.l_cache-1, self.l_cache-1)?; + let bx = conv1d_depthwise(&bx, self.conv.weight(), self.conv.bias())?; + bx.narrow(D::Minus1, 0, seq_len)? + } else { + let conv_state = self.cache.as_ref().unwrap(); + let conv_state = Tensor::cat(&[conv_state, &bx], D::Minus1)?; + let conv_state = conv_state.narrow(D::Minus1, 1, self.l_cache)?; + let conv_out = conv1d_depthwise(&conv_state, self.conv.weight(), self.conv.bias())?; + self.cache = Some(conv_state); + conv_out + }; + let y = c.mul(&conv_out)?; + let y = y.transpose(D::Minus1, D::Minus2)?.contiguous()?; + let y = self.out_proj.forward(&y)?; + Ok(y) + } + + pub fn clear_cache(&mut self) { + self.cache = None; + } +} + +enum LayerKind { + SelfAttn(QKNormAttention), + Conv(Lfm2ShortConv), +} + +impl LayerKind { + pub fn forward( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + match self { + LayerKind::SelfAttn(attn) => attn.forward(xs, cos, sin, attention_mask), + LayerKind::Conv(conv) => conv.forward(xs), + } + } +} + +pub struct Lfm2DecoderLayer { + layer: LayerKind, + feed_forward: GateUpDownMLP, + operator_norm: RmsNorm, + ffn_norm: RmsNorm, +} + +impl Lfm2DecoderLayer { + pub fn new(vb: VarBuilder, config: &Lfm2Config, layer_type: &str) -> Result { + let layer = if layer_type.eq("full_attention") { + let attn = QKNormAttention::new( + vb.pp("self_attn"), + config.hidden_size, + config.num_attention_heads, + None, + Some(config.num_key_value_heads), + false, + config.block_norm_eps, + Some("q_proj"), + Some("k_proj"), + Some("v_proj"), + Some("out_proj"), + Some("q_layernorm"), + Some("k_layernorm"), + )?; + LayerKind::SelfAttn(attn) + } else { + let conv = Lfm2ShortConv::new(vb.pp("conv"), config)?; + LayerKind::Conv(conv) + }; + let intermediate_size = if config.block_auto_adjust_ff_dim { + let inter_size = 2 * config.block_ff_dim / 3; + let inter_size = (config.block_ffn_dim_multiplier * inter_size as f64) as usize; + let inter_size = config.block_multiple_of + * ((inter_size + config.block_multiple_of - 1) / config.block_multiple_of); + inter_size + } else { + config.block_ff_dim + }; + let feed_forward = GateUpDownMLP::new( + vb.pp("feed_forward"), + config.hidden_size, + intermediate_size, + candle_nn::Activation::Silu, + false, + Some("w1"), + Some("w3"), + Some("w2"), + )?; + + let operator_norm = rms_norm(config.hidden_size, config.norm_eps, vb.pp("operator_norm"))?; + let ffn_norm = rms_norm(config.hidden_size, config.norm_eps, vb.pp("ffn_norm"))?; + Ok(Self { + layer, + feed_forward, + operator_norm, + ffn_norm, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let res = xs.clone(); + let xs = self.operator_norm.forward(xs)?; + let xs = self.layer.forward(&xs, cos, sin, attention_mask)?; + let res = xs.add(&res)?; + let xs = self.ffn_norm.forward(&res)?; + let xs = self.feed_forward.forward(&xs)?; + let xs = xs.add(&res)?; + Ok(xs) + } + + pub fn clear_cache(&mut self) { + match &mut self.layer { + LayerKind::SelfAttn(attn) => attn.clear_kv_cache(), + LayerKind::Conv(conv) => conv.clear_cache(), + } + } +} + +pub struct Lfm2Decoder { + embed_tokens: Embedding, + layers: Vec, + // rotary_emb: RoPE, + pos_emb: RoPE, + embedding_norm: RmsNorm, +} + +impl Lfm2Decoder { + pub fn new(vb: VarBuilder, config: &Lfm2Config) -> Result { + let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?; + let mut layers = vec![]; + let vb_layers = vb.pp("layers"); + // let layer_types = config.layer_types.as_ref().unwrap(); + let layer_types = config.get_layer_types()?; + for i in 0..config.num_hidden_layers { + let layer_type = layer_types.get(i).unwrap(); + let layer = Lfm2DecoderLayer::new(vb_layers.pp(i), config, layer_type)?; + layers.push(layer); + } + let dim = config.hidden_size / config.num_attention_heads; + let pos_emb = RoPE::new(dim, config.rope_theta, vb.device())?; + let embedding_norm = + rms_norm(config.hidden_size, config.norm_eps, vb.pp("embedding_norm"))?; + Ok(Self { + embed_tokens, + layers, + pos_emb, + embedding_norm, + }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + inputs_embeds: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let inputs_embeds = if let Some(embed) = inputs_embeds { + embed.clone() + } else { + self.embed_tokens.forward(input_ids)? + }; + + let (bs, seq_len, _) = inputs_embeds.dims3()?; + let attention_mask = if seq_len > 1 { + Some(prepare_causal_attention_mask( + bs, + seq_len, + seqlen_offset, + inputs_embeds.device(), + )?) + } else { + None + }; + let (cos, sin) = self + .pos_emb + .forward(seqlen_offset, seq_len, inputs_embeds.device())?; + let mut xs = inputs_embeds; + for layer in &mut self.layers { + xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?; + } + let xs = self.embedding_norm.forward(&xs)?; + Ok(xs) + } + + pub fn clear_cache(&mut self) { + for layer in &mut self.layers { + layer.clear_cache() + } + } +} + +pub struct Lfm2Model { + model: Lfm2Decoder, + lm_head: Linear, +} + +impl Lfm2Model { + pub fn new(vb: VarBuilder, config: &Lfm2Config) -> Result { + let model = Lfm2Decoder::new(vb.pp("model"), config)?; + let lm_head = if let Some(flag) = config.tie_embedding + && flag + { + Linear::new(model.embed_tokens.embeddings().clone(), None) + } else { + let linear = linear_b( + config.hidden_size, + config.vocab_size, + false, + vb.pp("lm_head"), + ); + match linear { + Ok(linear) => linear, + Err(_) => { + Linear::new(model.embed_tokens.embeddings().clone(), None) + } + } + }; + Ok(Self { model, lm_head }) + } + + pub fn forward(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result { + let xs = self.model.forward(input_ids, None, seqlen_offset)?; + let seq_len = xs.dim(1)?; + let last_xs = xs.narrow(1, seq_len - 1, 1)?; + let xs = self.lm_head.forward(&last_xs)?; + Ok(xs) + } + + pub fn clear_cache(&mut self) { + self.model.clear_cache(); + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index dc3fe87..c9ce259 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -19,6 +19,7 @@ pub mod qwen3vl; pub mod rmbg2_0; pub mod voxcpm; pub mod w2v_bert_2_0; +pub mod lfm2; use aha_openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, @@ -27,18 +28,15 @@ use anyhow::{Result, anyhow}; use rocket::futures::Stream; use crate::models::{ - deepseek_ocr::generate::DeepseekOCRGenerateModel, - fun_asr_nano::generate::FunAsrNanoGenerateModel, - glm_asr_nano::generate::GlmAsrNanoGenerateModel, glm_ocr::generate::GlmOcrGenerateModel, - hunyuan_ocr::generate::HunyuanOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel, - paddleocr_vl::generate::PaddleOCRVLGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, - qwen3::generate::Qwen3GenerateModel, qwen3_5::generate::Qwen3_5GenerateModel, - qwen3_asr::generate::Qwen3AsrGenerateModel, qwen3vl::generate::Qwen3VLGenerateModel, - rmbg2_0::generate::RMBG2_0Model, voxcpm::generate::VoxCPMGenerate, + deepseek_ocr::generate::DeepseekOCRGenerateModel, fun_asr_nano::generate::FunAsrNanoGenerateModel, glm_asr_nano::generate::GlmAsrNanoGenerateModel, glm_ocr::generate::GlmOcrGenerateModel, hunyuan_ocr::generate::HunyuanOCRGenerateModel, lfm2::generate::Lfm2GenerateModel, minicpm4::generate::MiniCPMGenerateModel, paddleocr_vl::generate::PaddleOCRVLGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, qwen3::generate::Qwen3GenerateModel, qwen3_5::generate::Qwen3_5GenerateModel, qwen3_asr::generate::Qwen3AsrGenerateModel, qwen3vl::generate::Qwen3VLGenerateModel, rmbg2_0::generate::RMBG2_0Model, voxcpm::generate::VoxCPMGenerate }; #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] pub enum WhichModel { + #[value(name = "lfm2-1.2b", hide = true)] + LFM2_1_2B, + #[value(name = "lfm2.5-1.2b-instruct", hide = true)] + LFM2_5_1_2BInstruct, #[value(name = "minicpm4-0.5b", hide = true)] MiniCPM4_0_5B, #[value(name = "qwen2.5vl-3b", hide = true)] @@ -97,6 +95,8 @@ impl WhichModel { /// Get the ModelScope model ID for this model variant pub fn model_id(self) -> &'static str { match self { + WhichModel::LFM2_1_2B => "LiquidAI/LFM2-1.2B", + WhichModel::LFM2_5_1_2BInstruct => "LiquidAI/LFM2.5-1.2B-Instruct", WhichModel::MiniCPM4_0_5B => "OpenBMB/MiniCPM4-0.5B", WhichModel::Qwen2_5vl3B => "Qwen/Qwen2.5-VL-3B-Instruct", WhichModel::Qwen2_5vl7B => "Qwen/Qwen2.5-VL-7B-Instruct", @@ -130,7 +130,9 @@ impl WhichModel { pub fn model_type(self) -> &'static str { match self { // LLM models - WhichModel::MiniCPM4_0_5B | WhichModel::Qwen3_0_6B => "llm", + WhichModel::MiniCPM4_0_5B | WhichModel::Qwen3_0_6B + | WhichModel::LFM2_1_2B + | WhichModel::LFM2_5_1_2BInstruct => "llm", WhichModel::Qwen2_5vl3B | WhichModel::Qwen2_5vl7B | WhichModel::Qwen3vl2B @@ -177,6 +179,7 @@ pub trait GenerateModel { pub enum ModelInstance<'a> { MiniCPM4(MiniCPMGenerateModel<'a>), + Lfm2(Lfm2GenerateModel<'a>), Qwen2_5VL(Qwen2_5VLGenerateModel<'a>), Qwen3(Qwen3GenerateModel<'a>), Qwen3_5(Qwen3_5GenerateModel<'a>), @@ -196,6 +199,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { fn generate(&mut self, mes: ChatCompletionParameters) -> Result { match self { ModelInstance::MiniCPM4(model) => model.generate(mes), + ModelInstance::Lfm2(model) => model.generate(mes), ModelInstance::Qwen2_5VL(model) => model.generate(mes), ModelInstance::Qwen3(model) => model.generate(mes), ModelInstance::Qwen3_5(model) => model.generate(mes), @@ -225,6 +229,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { > { match self { ModelInstance::MiniCPM4(model) => model.generate_stream(mes), + ModelInstance::Lfm2(model) => model.generate_stream(mes), ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes), ModelInstance::Qwen3(model) => model.generate_stream(mes), ModelInstance::Qwen3_5(model) => model.generate_stream(mes), @@ -253,6 +258,14 @@ pub fn load_model<'a>( let model = MiniCPMGenerateModel::init(path, None, None)?; ModelInstance::MiniCPM4(model) } + WhichModel::LFM2_1_2B => { + let model = Lfm2GenerateModel::init(path, None, None)?; + ModelInstance::Lfm2(model) + } + WhichModel::LFM2_5_1_2BInstruct => { + let model = Lfm2GenerateModel::init(path, None, None)?; + ModelInstance::Lfm2(model) + } WhichModel::Qwen2_5vl3B => { let model = Qwen2_5VLGenerateModel::init(path, None, None)?; ModelInstance::Qwen2_5VL(model) diff --git a/src/models/qwen3/model.rs b/src/models/qwen3/model.rs index 9543ad0..bd67077 100644 --- a/src/models/qwen3/model.rs +++ b/src/models/qwen3/model.rs @@ -6,7 +6,7 @@ use candle_nn::{ use crate::{ models::{ - common::{GateUpDownMLP, eager_attention_forward}, + common::{GateUpDownMLP, QKNormAttention, eager_attention_forward}, qwen3::config::Qwen3Config, }, position_embed::rope::{RoPE, apply_rotary_pos_emb}, @@ -135,7 +135,8 @@ impl Qwen3Attention { } pub struct Qwen3DecoderLayer { - self_attn: Qwen3Attention, + // self_attn: Qwen3Attention, + self_attn: QKNormAttention, mlp: GateUpDownMLP, input_layernorm: RmsNorm, post_attention_layernorm: RmsNorm, @@ -143,7 +144,22 @@ pub struct Qwen3DecoderLayer { impl Qwen3DecoderLayer { pub fn new(config: &Qwen3Config, vb: VarBuilder) -> Result { - let self_attn = Qwen3Attention::new(config, vb.pp("self_attn"))?; + // let self_attn = Qwen3Attention::new(config, vb.pp("self_attn"))?; + let self_attn = QKNormAttention::new( + vb.pp("self_attn"), + config.hidden_size, + config.num_attention_heads, + Some(config.head_dim), + Some(config.num_key_value_heads), + config.attention_bias, + config.rms_norm_eps, + None, + None, + None, + None, + None, + None, + )?; let mlp = GateUpDownMLP::new( vb.pp("mlp"), config.hidden_size, diff --git a/src/models/qwen3_5/generate.rs b/src/models/qwen3_5/generate.rs index 992b450..1c4cfae 100644 --- a/src/models/qwen3_5/generate.rs +++ b/src/models/qwen3_5/generate.rs @@ -36,10 +36,9 @@ pub struct Qwen3_5GenerateModel<'a> { impl<'a> Qwen3_5GenerateModel<'a> { pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { - let model_name = path - .split("/") - .collect::>() - .pop() + let model_name = std::path::Path::new(path) + .file_name() + .and_then(|s| s.to_str()) .unwrap_or("qwen3.5"); let chat_template = ChatTemplate::init(path)?; let tokenizer = TokenizerModel::init(path)?; diff --git a/src/tokenizer/mod.rs b/src/tokenizer/mod.rs index 1e1d8e3..7d4c165 100644 --- a/src/tokenizer/mod.rs +++ b/src/tokenizer/mod.rs @@ -110,6 +110,14 @@ impl TokenizerModel { .map_err(|e| anyhow!(format!("tokenizer encode error{}", e)))?; Ok(decode) } + + pub fn token_decode_with_special(&self, tokens: Vec) -> Result { + let decode = self + .tokenizer + .decode(&tokens, false) + .map_err(|e| anyhow!(format!("tokenizer encode error{}", e)))?; + Ok(decode) + } } pub fn sentencepiece_encode( diff --git a/tests/config_tests.rs b/tests/config_tests.rs index 7b1dfcb..94409c8 100644 --- a/tests/config_tests.rs +++ b/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(()) +} diff --git a/tests/messy_test.rs b/tests/messy_test.rs index 43ab670..e80cccc 100644 --- a/tests/messy_test.rs +++ b/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)?; diff --git a/tests/test_lfm2.rs b/tests/test_lfm2.rs new file mode 100644 index 0000000..60e2c9d --- /dev/null +++ b/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(()) +} \ No newline at end of file diff --git a/tests/test_qwen3.rs b/tests/test_qwen3.rs index cc1b944..61bd73a 100644 --- a/tests/test_qwen3.rs +++ b/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 = diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 643b01d..cd949b8 100644 --- a/tests/weight_test.rs +++ b/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(()) +}