From b6b970fd258e2fbf0a3d6638922804b8525b6cbe Mon Sep 17 00:00:00 2001 From: "273265088@qq.com" <273265088@qq.com> Date: Thu, 26 Mar 2026 13:40:24 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=20All-MiniLM-L6-v2-Embedding?= =?UTF-8?q?=20=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 10 + README.zh-CN.md | 10 + docs/supported-models.md | 2 +- docs/supported-models.zh-CN.md | 2 +- src/exec/all_minilm_l6_v2.rs | 39 + src/exec/mod.rs | 1 + src/main.rs | 42 + src/models/all_minilm_l6_v2/config.rs | 105 +++ src/models/all_minilm_l6_v2/generate.rs | 72 ++ src/models/all_minilm_l6_v2/mod.rs | 3 + src/models/all_minilm_l6_v2/model.rs | 956 ++++++++++++++++++++ src/models/core/artifact.rs | 5 + src/models/core/registry.rs | 15 +- src/models/mod.rs | 32 +- tests/test_all_minilm_l6_v2_multi_format.rs | 353 ++++++++ tests/test_load_spec.rs | 33 + 16 files changed, 1670 insertions(+), 10 deletions(-) create mode 100644 src/exec/all_minilm_l6_v2.rs create mode 100644 src/models/all_minilm_l6_v2/config.rs create mode 100644 src/models/all_minilm_l6_v2/generate.rs create mode 100644 src/models/all_minilm_l6_v2/mod.rs create mode 100644 src/models/all_minilm_l6_v2/model.rs create mode 100644 tests/test_all_minilm_l6_v2_multi_format.rs diff --git a/README.md b/README.md index 5bf72ac..d83b85a 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,15 @@ aha -m qwen3asr-0.6b # Run inference directly (without starting service) aha run -m qwen3asr-0.6b -i "audio.wav" +# Run local all-MiniLM-L6-v2 embedding (native safetensors) +aha run -m all-minilm-l6-v2 -i "Rust embedding test" --weight-path D:\model_download\all-MiniLM-L6-v2 + +# Run local all-MiniLM-L6-v2 embedding (GGUF) +aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format gguf --gguf-path D:\model_download\All-MiniLM-L6-v2-Embedding-GGUF --tokenizer-dir D:\model_download\all-MiniLM-L6-v2 + +# Run local all-MiniLM-L6-v2 embedding (ONNX) +aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format onnx --onnx-path D:\model_download\all-MiniLM-L6-v2\onnx --tokenizer-dir D:\model_download\all-MiniLM-L6-v2 + # Start service only (model already downloaded) aha serv -m qwen3asr-0.6b -p 10100 @@ -126,6 +135,7 @@ curl http://localhost:10100/chat/completions \ | Category | Models | |----------|--------| | **Text** | Qwen3, MiniCPM4 | +| **Embedding** | Qwen3-Embedding, all-MiniLM-L6-v2 | | **Vision** | Qwen2.5-VL, Qwen3-VL | | **OCR** | DeepSeek-OCR, Hunyuan-OCR, PaddleOCR-VL | | **ASR** | GLM-ASR-Nano, Fun-ASR-Nano, Qwen3-ASR | diff --git a/README.zh-CN.md b/README.zh-CN.md index b7b0f8d..cc31391 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -96,6 +96,15 @@ aha -m qwen3asr-0.6b # 直接运行推理（无需启动服务） aha run -m qwen3asr-0.6b -i "audio.wav" +# 本地运行 all-MiniLM-L6-v2 向量模型（原生 safetensors） +aha run -m all-minilm-l6-v2 -i "Rust embedding test" --weight-path D:\model_download\all-MiniLM-L6-v2 + +# 本地运行 all-MiniLM-L6-v2 向量模型（GGUF） +aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format gguf --gguf-path D:\model_download\All-MiniLM-L6-v2-Embedding-GGUF --tokenizer-dir D:\model_download\all-MiniLM-L6-v2 + +# 本地运行 all-MiniLM-L6-v2 向量模型（ONNX） +aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format onnx --onnx-path D:\model_download\all-MiniLM-L6-v2\onnx --tokenizer-dir D:\model_download\all-MiniLM-L6-v2 + # 仅启动服务（模型已下载） aha serv -m qwen3asr-0.6b -p 10100 @@ -125,6 +134,7 @@ curl http://localhost:10100/chat/completions \ | 类别 | 模型 | |------|------| | **文本** | Qwen3, MiniCPM4 | +| **向量** | Qwen3-Embedding, all-MiniLM-L6-v2 | | **视觉** | Qwen2.5-VL, Qwen3-VL | | **OCR** | DeepSeek-OCR, Hunyuan-OCR, PaddleOCR-VL | | **ASR** | GLM-ASR-Nano, Fun-ASR-Nano,Qwen3-ASR | diff --git a/docs/supported-models.md b/docs/supported-models.md index e57355e..96478c9 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -16,6 +16,7 @@ aha supports a growing collection of state-of-the-art AI models across multiple | **Qwen3-Embedding-0.6B** | 0.6B | Text embedding (safetensors / gguf / onnx) | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | | **Qwen3-Embedding-4B** | 4B | Text embedding (safetensors / gguf / onnx) | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | | **Qwen3-Embedding-8B** | 8B | Text embedding (safetensors / gguf / onnx) | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | +| **all-MiniLM-L6-v2** | 22M | Sentence-transformers embedding (safetensors / gguf / onnx) | [Apache 2.0](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/blob/main/LICENSE) | ## Reranker @@ -123,7 +124,6 @@ The following repositories are now cataloged for future integration, but are **n - jinaai/jina-embeddings-v5-text-nano - jinaai/jina-embeddings-v5-text-small-text-matching - jinaai/jina-embeddings-v5-text-small-text-matching-GGUF -- sentence-transformers/all-MiniLM-L6-v2 - sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 ### Reranker Models diff --git a/docs/supported-models.zh-CN.md b/docs/supported-models.zh-CN.md index 0223124..d35308f 100644 --- a/docs/supported-models.zh-CN.md +++ b/docs/supported-models.zh-CN.md @@ -16,6 +16,7 @@ aha 支持多个领域的最先进 AI 模型集合。 | **Qwen3-Embedding-0.6B** | 0.6B | 文本向量（safetensors / gguf / onnx） | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | | **Qwen3-Embedding-4B** | 4B | 文本向量（safetensors / gguf / onnx） | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | | **Qwen3-Embedding-8B** | 8B | 文本向量（safetensors / gguf / onnx） | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) | +| **all-MiniLM-L6-v2** | 22M | sentence-transformers 文本向量（safetensors / gguf / onnx） | [Apache 2.0](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/blob/main/LICENSE) | ## Reranker @@ -123,7 +124,6 @@ aha 支持多个领域的最先进 AI 模型集合。 - jinaai/jina-embeddings-v5-text-nano - jinaai/jina-embeddings-v5-text-small-text-matching - jinaai/jina-embeddings-v5-text-small-text-matching-GGUF -- sentence-transformers/all-MiniLM-L6-v2 - sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 ### Reranker 模型 diff --git a/src/exec/all_minilm_l6_v2.rs b/src/exec/all_minilm_l6_v2.rs new file mode 100644 index 0000000..07d7717 --- /dev/null +++ b/src/exec/all_minilm_l6_v2.rs @@ -0,0 +1,39 @@ +use anyhow::Result; + +use crate::exec::ExecModel; +use crate::models::{LoadSpec, all_minilm_l6_v2::generate::AllMiniLML6V2Model}; +use crate::utils::get_file_path; + +pub struct AllMiniLML6V2Exec; + +impl AllMiniLML6V2Exec { + pub fn run_with_spec(input: &[String], output: Option<&str>, spec: &LoadSpec) -> Result<()> { + let input_text = input + .first() + .ok_or_else(|| anyhow::anyhow!("embedding run requires one text input"))?; + let text = if input_text.starts_with("file://") { + let path = get_file_path(input_text)?; + std::fs::read_to_string(path)? + } else { + input_text.clone() + }; + + let mut model = AllMiniLML6V2Model::init_from_spec(spec, None, None)?; + let embedding = model.embed(&[text])?; + + let output_json = serde_json::to_string_pretty(&embedding)?; + println!("{}", output_json); + if let Some(out) = output { + std::fs::write(out, output_json)?; + println!("Output saved to: {}", out); + } + Ok(()) + } +} + +impl ExecModel for AllMiniLML6V2Exec { + fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { + let spec = LoadSpec::for_safetensors(crate::models::WhichModel::AllMiniLML6V2, weight_path); + Self::run_with_spec(input, output, &spec) + } +} diff --git a/src/exec/mod.rs b/src/exec/mod.rs index fa28273..e41ef3c 100644 --- a/src/exec/mod.rs +++ b/src/exec/mod.rs @@ -3,6 +3,7 @@ //! This module provides model-specific exec implementations for the `run` subcommand. //! Each model has its own exec module that handles input/output parsing and model invocation. +pub mod all_minilm_l6_v2; pub mod deepseek_ocr; pub mod fun_asr_nano; pub mod glm_asr_nano; diff --git a/src/main.rs b/src/main.rs index 5313d9a..e76357c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -428,6 +428,11 @@ fn resolve_load_spec_for_run(args: &RunArgs) -> anyhow::Result { fn run_target_model_with_spec(args: &RunArgs, spec: &LoadSpec) -> anyhow::Result { match args.model { + WhichModel::AllMiniLML6V2 => { + use aha::exec::all_minilm_l6_v2::AllMiniLML6V2Exec; + AllMiniLML6V2Exec::run_with_spec(&args.input, args.output.as_deref(), spec)?; + Ok(true) + } WhichModel::Qwen3_0_6B => { use aha::exec::qwen3::Qwen3Exec; Qwen3Exec::run_with_spec(&args.input, args.output.as_deref(), spec)?; @@ -620,6 +625,10 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> { None => get_default_weight_path(model), }; match model { + WhichModel::AllMiniLML6V2 => { + use aha::exec::all_minilm_l6_v2::AllMiniLML6V2Exec; + AllMiniLML6V2Exec::run(&input, output.as_deref(), &weight_path)?; + } WhichModel::MiniCPM4_0_5B => { use aha::exec::minicpm4::MiniCPM4Exec; MiniCPM4Exec::run(&input, output.as_deref(), &weight_path)?; @@ -975,6 +984,39 @@ mod tests { ); } + #[test] + fn parse_all_minilm_run_onnx_flags() { + let cli = Cli::try_parse_from([ + "aha", + "run", + "--model", + "all-minilm-l6-v2", + "--input", + "hello", + "--artifact-format", + "onnx", + "--onnx-path", + "D:\\model_download\\all-MiniLM-L6-v2\\onnx", + "--tokenizer-dir", + "D:\\model_download\\all-MiniLM-L6-v2", + ]) + .expect("run args should parse"); + + let Some(Commands::Run(args)) = cli.command else { + panic!("expected run subcommand"); + }; + assert!(matches!(args.artifact_format, Some(ArtifactArg::Onnx))); + assert_eq!(args.model, WhichModel::AllMiniLML6V2); + assert_eq!( + args.onnx_path.as_deref(), + Some("D:\\model_download\\all-MiniLM-L6-v2\\onnx") + ); + assert_eq!( + args.tokenizer_dir.as_deref(), + Some("D:\\model_download\\all-MiniLM-L6-v2") + ); + } + #[test] fn parse_serv_onnx_flags() { let cli = Cli::try_parse_from([ diff --git a/src/models/all_minilm_l6_v2/config.rs b/src/models/all_minilm_l6_v2/config.rs new file mode 100644 index 0000000..389ae27 --- /dev/null +++ b/src/models/all_minilm_l6_v2/config.rs @@ -0,0 +1,105 @@ +use std::path::Path; + +use anyhow::Result; +use candle_transformers::models::bert::Config as BertConfig; +use serde::Deserialize; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AllMiniLML6V2PoolingStrategy { + Cls, + Mean, + Max, + MeanSqrtLen, +} + +#[derive(Debug, Clone)] +pub struct AllMiniLML6V2Config { + pub base: BertConfig, + pub pooling: AllMiniLML6V2PoolingStrategy, + pub normalize: bool, + pub max_seq_length: usize, + pub do_lower_case: bool, +} + +#[derive(Debug, Deserialize)] +struct SentenceBertConfig { + #[serde(default = "default_max_seq_length")] + max_seq_length: usize, + #[serde(default)] + do_lower_case: bool, +} + +#[derive(Debug, Deserialize)] +struct PoolingConfig { + #[serde(default)] + pooling_mode_cls_token: bool, + #[serde(default)] + pooling_mode_mean_tokens: bool, + #[serde(default)] + pooling_mode_max_tokens: bool, + #[serde(default)] + pooling_mode_mean_sqrt_len_tokens: bool, +} + +#[derive(Debug, Deserialize)] +struct ModuleEntry { + #[serde(default)] + r#type: String, +} + +fn default_max_seq_length() -> usize { + 256 +} + +impl AllMiniLML6V2Config { + pub fn load(path: &str) -> Result { + let config_path = Path::new(path).join("config.json"); + let base: BertConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + + let sentence_bert_path = Path::new(path).join("sentence_bert_config.json"); + let sentence_bert = if sentence_bert_path.exists() { + serde_json::from_slice::(&std::fs::read(sentence_bert_path)?)? + } else { + SentenceBertConfig { + max_seq_length: default_max_seq_length(), + do_lower_case: false, + } + }; + + let pooling_path = Path::new(path).join("1_Pooling").join("config.json"); + let pooling = if pooling_path.exists() { + let cfg: PoolingConfig = serde_json::from_slice(&std::fs::read(pooling_path)?)?; + if cfg.pooling_mode_mean_tokens { + AllMiniLML6V2PoolingStrategy::Mean + } else if cfg.pooling_mode_cls_token { + AllMiniLML6V2PoolingStrategy::Cls + } else if cfg.pooling_mode_max_tokens { + AllMiniLML6V2PoolingStrategy::Max + } else if cfg.pooling_mode_mean_sqrt_len_tokens { + AllMiniLML6V2PoolingStrategy::MeanSqrtLen + } else { + AllMiniLML6V2PoolingStrategy::Mean + } + } else { + AllMiniLML6V2PoolingStrategy::Mean + }; + + let modules_path = Path::new(path).join("modules.json"); + let normalize = if modules_path.exists() { + let modules: Vec = serde_json::from_slice(&std::fs::read(modules_path)?)?; + modules + .iter() + .any(|module| module.r#type.ends_with(".Normalize")) + } else { + false + }; + + Ok(Self { + base, + pooling, + normalize, + max_seq_length: sentence_bert.max_seq_length, + do_lower_case: sentence_bert.do_lower_case, + }) + } +} diff --git a/src/models/all_minilm_l6_v2/generate.rs b/src/models/all_minilm_l6_v2/generate.rs new file mode 100644 index 0000000..e867f23 --- /dev/null +++ b/src/models/all_minilm_l6_v2/generate.rs @@ -0,0 +1,72 @@ +use anyhow::Result; +use candle_core::{DType, Device}; + +use crate::models::{ + all_minilm_l6_v2::model::AllMiniLML6V2Backend, + artifact::{ArtifactKind, LoadSpec}, + common::retrieval::TextEmbeddingBackend, +}; + +pub struct AllMiniLML6V2Model { + backend: AllMiniLML6V2Backend, +} + +impl AllMiniLML6V2Model { + pub fn init_from_spec( + spec: &LoadSpec, + device: Option<&Device>, + dtype: Option, + ) -> Result { + match spec.resolved_artifact() { + ArtifactKind::Safetensors => { + let path = spec.paths.weight_dir.as_deref().ok_or_else(|| { + anyhow::anyhow!("weight_path is required for all-minilm-l6-v2 safetensors") + })?; + Self::init(path, device, dtype) + } + ArtifactKind::Gguf => { + let path = spec.paths.gguf_path.as_deref().ok_or_else(|| { + anyhow::anyhow!("gguf_path is required for all-minilm-l6-v2 gguf") + })?; + Self::init_gguf(path, spec.paths.tokenizer_dir.as_deref(), device, dtype) + } + ArtifactKind::Onnx => { + let path = spec.paths.onnx_path.as_deref().ok_or_else(|| { + anyhow::anyhow!("onnx_path is required for all-minilm-l6-v2 onnx") + })?; + Self::init_onnx(path, spec.paths.tokenizer_dir.as_deref()) + } + ArtifactKind::Auto => unreachable!("artifact kind should be resolved before init"), + } + } + + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let backend = AllMiniLML6V2Backend::load(path, device, dtype)?; + Ok(Self { backend }) + } + + pub fn init_onnx(onnx_path: &str, tokenizer_dir: Option<&str>) -> Result { + let backend = AllMiniLML6V2Backend::load_onnx(onnx_path, tokenizer_dir)?; + Ok(Self { backend }) + } + + pub fn init_gguf( + gguf_path: &str, + tokenizer_dir: Option<&str>, + device: Option<&Device>, + dtype: Option, + ) -> Result { + let backend = AllMiniLML6V2Backend::load_gguf(gguf_path, tokenizer_dir, device, dtype)?; + Ok(Self { backend }) + } + + pub fn embed(&mut self, input: &[String]) -> Result>> { + self.backend.embed_texts(input) + } +} + +impl TextEmbeddingBackend for AllMiniLML6V2Model { + fn embed_texts(&mut self, input: &[String]) -> Result>> { + self.backend.embed_texts(input) + } +} diff --git a/src/models/all_minilm_l6_v2/mod.rs b/src/models/all_minilm_l6_v2/mod.rs new file mode 100644 index 0000000..7fca417 --- /dev/null +++ b/src/models/all_minilm_l6_v2/mod.rs @@ -0,0 +1,3 @@ +pub mod config; +pub mod generate; +pub mod model; diff --git a/src/models/all_minilm_l6_v2/model.rs b/src/models/all_minilm_l6_v2/model.rs new file mode 100644 index 0000000..05aa0b9 --- /dev/null +++ b/src/models/all_minilm_l6_v2/model.rs @@ -0,0 +1,956 @@ +use anyhow::{Result, anyhow}; +use candle_core::{DType, Device, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::models::bert::BertModel; +#[cfg(feature = "onnx-runtime")] +use half::f16; +#[cfg(feature = "onnx-runtime")] +use ndarray::{Array, IxDyn}; +use std::{ + collections::HashMap, + path::{Path, PathBuf}, +}; + +use crate::{ + models::{ + all_minilm_l6_v2::config::{AllMiniLML6V2Config, AllMiniLML6V2PoolingStrategy}, + common::{ + gguf::load_gguf_file, + onnx::{create_session, resolve_tokenizer_dir}, + retrieval::l2_normalize, + }, + }, + tokenizer::TokenizerModel, + utils::{find_type_files, get_device, get_dtype}, +}; + +pub enum AllMiniLML6V2Backend { + Safetensors(AllMiniLML6V2SafetensorsBackend), + Gguf(AllMiniLML6V2GgufBackend), + Onnx(AllMiniLML6V2OnnxBackend), +} + +impl AllMiniLML6V2Backend { + pub fn load(path: &str, device: Option<&Device>, dtype: Option) -> Result { + Ok(Self::Safetensors(AllMiniLML6V2SafetensorsBackend::load( + path, device, dtype, + )?)) + } + + pub fn load_onnx(onnx_path: &str, tokenizer_dir: Option<&str>) -> Result { + Ok(Self::Onnx(AllMiniLML6V2OnnxBackend::load( + onnx_path, + tokenizer_dir, + )?)) + } + + pub fn load_gguf( + gguf_path: &str, + tokenizer_dir: Option<&str>, + device: Option<&Device>, + dtype: Option, + ) -> Result { + Ok(Self::Gguf(AllMiniLML6V2GgufBackend::load( + gguf_path, + tokenizer_dir, + device, + dtype, + )?)) + } + + pub fn embed_texts(&mut self, input: &[String]) -> Result>> { + match self { + Self::Safetensors(backend) => backend.embed_texts(input), + Self::Gguf(backend) => backend.embed_texts(input), + Self::Onnx(backend) => backend.embed_texts(input), + } + } +} + +pub struct AllMiniLML6V2SafetensorsBackend { + tokenizer: TokenizerModel, + model: BertModel, + device: Device, + pooling: AllMiniLML6V2PoolingStrategy, + normalize: bool, + max_seq_length: usize, + do_lower_case: bool, +} + +impl AllMiniLML6V2SafetensorsBackend { + pub fn load(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let tokenizer = TokenizerModel::init(path)?; + let cfg = AllMiniLML6V2Config::load(path)?; + let device = get_device(device); + let dtype = get_dtype(dtype, "float32"); + let model_list = find_type_files(path, "safetensors")?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? }; + let model = BertModel::load(vb, &cfg.base)?; + Ok(Self { + tokenizer, + model, + device, + pooling: cfg.pooling, + normalize: cfg.normalize, + max_seq_length: cfg.max_seq_length, + do_lower_case: cfg.do_lower_case, + }) + } + + pub fn embed_texts(&mut self, input: &[String]) -> Result>> { + if input.is_empty() { + return Err(anyhow!("embedding input cannot be empty")); + } + let mut out = Vec::with_capacity(input.len()); + for text in input { + out.push(self.embed_one(text)?); + } + Ok(out) + } + + fn embed_one(&mut self, text: &str) -> Result> { + let token_ids = prepare_token_ids( + &self.tokenizer, + text, + self.max_seq_length, + self.do_lower_case, + )?; + let seq_len = token_ids.len(); + let input_ids = Tensor::from_slice(&token_ids, (1, seq_len), &self.device)?; + let token_type_ids = Tensor::zeros((1, seq_len), DType::U32, &self.device)?; + let attention_mask = Tensor::ones((1, seq_len), DType::U32, &self.device)?; + let hidden = self + .model + .forward(&input_ids, &token_type_ids, Some(&attention_mask))? + .squeeze(0)? + .to_dtype(DType::F32)?; + let hidden_vec = hidden.to_vec2::()?; + let mut pooled = pool_hidden_state(&hidden_vec, self.pooling)?; + if self.normalize { + l2_normalize(&mut pooled); + } + Ok(pooled) + } +} + +pub struct AllMiniLML6V2GgufBackend { + tokenizer: TokenizerModel, + model: BertModel, + device: Device, + pooling: AllMiniLML6V2PoolingStrategy, + normalize: bool, + max_seq_length: usize, + do_lower_case: bool, +} + +impl AllMiniLML6V2GgufBackend { + pub fn load( + gguf_path: &str, + tokenizer_dir: Option<&str>, + device: Option<&Device>, + dtype: Option, + ) -> Result { + let tokenizer_dir = resolve_all_minilm_tokenizer_dir(gguf_path, tokenizer_dir)?; + let tokenizer = TokenizerModel::init(&tokenizer_dir.to_string_lossy())?; + let cfg = AllMiniLML6V2Config::load(&tokenizer_dir.to_string_lossy())?; + let device = get_device(device); + let dtype = dtype.unwrap_or(DType::F32); + let mut gguf = load_gguf_file(resolve_all_minilm_gguf_file(gguf_path)?.as_ref(), &device)?; + let tensors = load_all_minilm_gguf_tensors(&mut gguf, &cfg, &device, dtype)?; + let vb = VarBuilder::from_tensors(tensors, dtype, &device); + let model = BertModel::load(vb, &cfg.base)?; + Ok(Self { + tokenizer, + model, + device, + pooling: cfg.pooling, + normalize: cfg.normalize, + max_seq_length: cfg.max_seq_length, + do_lower_case: cfg.do_lower_case, + }) + } + + pub fn embed_texts(&mut self, input: &[String]) -> Result>> { + if input.is_empty() { + return Err(anyhow!("embedding input cannot be empty")); + } + let mut out = Vec::with_capacity(input.len()); + for text in input { + out.push(self.embed_one(text)?); + } + Ok(out) + } + + fn embed_one(&mut self, text: &str) -> Result> { + let token_ids = prepare_token_ids( + &self.tokenizer, + text, + self.max_seq_length, + self.do_lower_case, + )?; + let seq_len = token_ids.len(); + let input_ids = Tensor::from_slice(&token_ids, (1, seq_len), &self.device)?; + let token_type_ids = Tensor::zeros((1, seq_len), DType::U32, &self.device)?; + let attention_mask = Tensor::ones((1, seq_len), DType::U32, &self.device)?; + let hidden = self + .model + .forward(&input_ids, &token_type_ids, Some(&attention_mask))? + .squeeze(0)? + .to_dtype(DType::F32)?; + let hidden_vec = hidden.to_vec2::()?; + let mut pooled = pool_hidden_state(&hidden_vec, self.pooling)?; + if self.normalize { + l2_normalize(&mut pooled); + } + Ok(pooled) + } +} + +#[cfg_attr(not(feature = "onnx-runtime"), allow(dead_code))] +pub struct AllMiniLML6V2OnnxBackend { + tokenizer: TokenizerModel, + #[cfg(feature = "onnx-runtime")] + session: ort::session::Session, + #[cfg(not(feature = "onnx-runtime"))] + _session: (), + output_names: Vec, + input_descriptors: Vec, + pooling: AllMiniLML6V2PoolingStrategy, + normalize: bool, + max_seq_length: usize, + do_lower_case: bool, +} + +#[cfg_attr(not(feature = "onnx-runtime"), allow(dead_code))] +#[derive(Clone)] +struct OnnxInputDescriptor { + name: String, + shape: Vec, + kind: Option, +} + +#[cfg_attr(not(feature = "onnx-runtime"), allow(dead_code))] +#[derive(Clone, Copy)] +enum OnnxTensorKind { + Bool, + I32, + I64, + F16, + F32, +} + +#[cfg(feature = "onnx-runtime")] +fn map_tensor_kind(ty: ort::value::TensorElementType) -> Option { + match ty { + ort::value::TensorElementType::Bool => Some(OnnxTensorKind::Bool), + ort::value::TensorElementType::Int32 => Some(OnnxTensorKind::I32), + ort::value::TensorElementType::Int64 => Some(OnnxTensorKind::I64), + ort::value::TensorElementType::Float16 => Some(OnnxTensorKind::F16), + ort::value::TensorElementType::Float32 => Some(OnnxTensorKind::F32), + _ => None, + } +} + +impl AllMiniLML6V2OnnxBackend { + pub fn load(onnx_path: &str, tokenizer_dir: Option<&str>) -> Result { + let tokenizer_dir = + resolve_tokenizer_dir(onnx_path, tokenizer_dir, &["tokenizer.json", "config.json"])?; + let tokenizer = TokenizerModel::init(&tokenizer_dir.to_string_lossy())?; + let cfg = AllMiniLML6V2Config::load(&tokenizer_dir.to_string_lossy())?; + let bundle = create_session(onnx_path, None)?; + #[cfg(feature = "onnx-runtime")] + { + let input_descriptors = bundle + .session + .inputs() + .iter() + .map(|input| { + let (shape, kind) = match input.dtype() { + ort::value::ValueType::Tensor { ty, shape, .. } => ( + shape.iter().copied().collect::>(), + map_tensor_kind(*ty), + ), + _ => (Vec::new(), None), + }; + OnnxInputDescriptor { + name: input.name().to_string(), + shape, + kind, + } + }) + .collect::>(); + Ok(Self { + tokenizer, + session: bundle.session, + output_names: bundle.output_names, + input_descriptors, + pooling: cfg.pooling, + normalize: cfg.normalize, + max_seq_length: cfg.max_seq_length, + do_lower_case: cfg.do_lower_case, + }) + } + #[cfg(not(feature = "onnx-runtime"))] + { + let _ = bundle; + Ok(Self { + tokenizer, + _session: (), + output_names: Vec::new(), + input_descriptors: Vec::new(), + pooling: cfg.pooling, + normalize: cfg.normalize, + max_seq_length: cfg.max_seq_length, + do_lower_case: cfg.do_lower_case, + }) + } + } + + pub fn embed_texts(&mut self, input: &[String]) -> Result>> { + if input.is_empty() { + return Err(anyhow!("embedding input cannot be empty")); + } + let mut out = Vec::with_capacity(input.len()); + for text in input { + out.push(self.embed_one(text)?); + } + Ok(out) + } + + #[cfg(feature = "onnx-runtime")] + fn embed_one(&mut self, text: &str) -> Result> { + let token_ids = prepare_token_ids( + &self.tokenizer, + text, + self.max_seq_length, + self.do_lower_case, + )?; + let seq_len = token_ids.len(); + let input_ids = token_ids.iter().map(|id| *id as i64).collect::>(); + let attention_mask = vec![1_i64; seq_len]; + let token_type_ids = vec![0_i64; seq_len]; + + let mut inputs = Vec::with_capacity(self.input_descriptors.len()); + for desc in &self.input_descriptors { + let value = + self.build_input_value(desc, &input_ids, &attention_mask, &token_type_ids)?; + inputs.push((desc.name.clone(), value)); + } + + let outputs = self.session.run(inputs).map_err(|err| { + let names = self + .input_descriptors + .iter() + .map(|desc| desc.name.clone()) + .collect::>() + .join(", "); + anyhow!("failed to run all-minilm-l6-v2 onnx session: {err}; inputs={names}") + })?; + + let output_value = outputs + .get("sentence_embedding") + .or_else(|| outputs.get("last_hidden_state")) + .or_else(|| outputs.get("token_embeddings")) + .or_else(|| self.output_names.first().and_then(|name| outputs.get(name))) + .ok_or_else(|| anyhow!("onnx output for all-minilm-l6-v2 not found"))?; + + let mut embedding = extract_embedding_output(output_value, seq_len, self.pooling)?; + if self.normalize { + l2_normalize(&mut embedding); + } + Ok(embedding) + } + + #[cfg(feature = "onnx-runtime")] + fn build_input_value( + &self, + desc: &OnnxInputDescriptor, + input_ids: &[i64], + attention_mask: &[i64], + token_type_ids: &[i64], + ) -> Result { + let shape = self.resolve_shape(desc, input_ids.len()); + match desc.name.as_str() { + "input_ids" => return build_i64_like_input(desc, shape, input_ids), + "attention_mask" => return build_i64_like_input(desc, shape, attention_mask), + "token_type_ids" => return build_i64_like_input(desc, shape, token_type_ids), + _ => {} + } + self.build_zero_tensor(desc, shape) + } + + #[cfg(feature = "onnx-runtime")] + fn resolve_shape(&self, desc: &OnnxInputDescriptor, seq_len: usize) -> Vec { + if desc.shape.is_empty() { + return vec![1_i64]; + } + desc.shape + .iter() + .enumerate() + .map(|(idx, dim)| { + if *dim >= 0 { + *dim + } else if idx == 0 { + 1 + } else { + seq_len as i64 + } + }) + .collect() + } + + #[cfg(feature = "onnx-runtime")] + fn build_zero_tensor( + &self, + desc: &OnnxInputDescriptor, + shape: Vec, + ) -> Result { + let kind = desc.kind.ok_or_else(|| { + anyhow!( + "unsupported all-minilm-l6-v2 onnx input dtype for {}", + desc.name + ) + })?; + let elem_count = shape.iter().try_fold(1_usize, |acc, dim| { + if *dim < 0 { + Err(anyhow!( + "cannot resolve dynamic onnx shape for input {}: {:?}", + desc.name, + shape + )) + } else { + Ok(acc.saturating_mul(*dim as usize)) + } + })?; + let has_zero_dim = shape.contains(&0); + let make_ndarray = || { + let dims = shape.iter().map(|dim| *dim as usize).collect::>(); + IxDyn(&dims) + }; + let value = match kind { + OnnxTensorKind::Bool => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![false; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![false; elem_count]))?.into_dyn() + } + } + OnnxTensorKind::I32 => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![0_i32; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![0_i32; elem_count]))?.into_dyn() + } + } + OnnxTensorKind::I64 => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![0_i64; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![0_i64; elem_count]))?.into_dyn() + } + } + OnnxTensorKind::F16 => { + if has_zero_dim { + let arr = Array::from_shape_vec( + make_ndarray(), + vec![f16::from_f32(0.0); elem_count], + )?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![f16::from_f32(0.0); elem_count]))? + .into_dyn() + } + } + OnnxTensorKind::F32 => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![0_f32; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![0_f32; elem_count]))?.into_dyn() + } + } + }; + Ok(value) + } + + #[cfg(not(feature = "onnx-runtime"))] + fn embed_one(&mut self, _text: &str) -> Result> { + Err(anyhow!( + "onnx runtime support is not enabled; rebuild with --features onnx-runtime" + )) + } +} + +fn prepare_token_ids( + tokenizer: &TokenizerModel, + text: &str, + max_seq_length: usize, + do_lower_case: bool, +) -> Result> { + let text = if do_lower_case { + text.to_lowercase() + } else { + text.to_string() + }; + let mut token_ids = tokenizer.text_encode_vec(text, true)?; + if token_ids.len() > max_seq_length { + token_ids.truncate(max_seq_length); + if let Some(sep_id) = tokenizer.tokenizer.token_to_id("[SEP]") + && let Some(last) = token_ids.last_mut() + { + *last = sep_id; + } + } + if token_ids.is_empty() { + return Err(anyhow!("embedding tokenized input cannot be empty")); + } + Ok(token_ids) +} + +fn resolve_all_minilm_gguf_file(path: &str) -> Result { + let model_path = Path::new(path); + if !model_path.exists() { + return Err(anyhow!("gguf model path not found: {}", path)); + } + if model_path.is_file() { + if model_path + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case("gguf")) + { + return Ok(model_path.to_string_lossy().to_string()); + } + return Err(anyhow!( + "gguf model path does not point to a .gguf file: {}", + path + )); + } + + let mut matches = std::fs::read_dir(model_path)? + .flatten() + .map(|entry| entry.path()) + .filter(|candidate| { + candidate.is_file() + && candidate + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case("gguf")) + }) + .collect::>(); + matches.sort(); + matches + .into_iter() + .next() + .map(|path| path.to_string_lossy().to_string()) + .ok_or_else(|| anyhow!("no .gguf file found in {}", model_path.display())) +} + +fn resolve_all_minilm_tokenizer_dir( + gguf_path: &str, + tokenizer_dir: Option<&str>, +) -> Result { + fn has_required_files(path: &Path) -> bool { + path.join("tokenizer.json").exists() && path.join("config.json").exists() + } + + fn without_suffix(path: &Path, suffix: &str) -> Option { + let name = path.file_name()?.to_string_lossy(); + let base = name.strip_suffix(suffix)?; + if base.is_empty() { + None + } else { + Some(path.with_file_name(base)) + } + } + + let mut candidates = Vec::new(); + if let Some(dir) = tokenizer_dir { + candidates.push(PathBuf::from(dir)); + } + + let gguf_file = PathBuf::from(resolve_all_minilm_gguf_file(gguf_path)?); + if let Some(parent) = gguf_file.parent() { + candidates.push(parent.to_path_buf()); + if let Some(grand) = parent.parent() { + candidates.push(grand.to_path_buf()); + } + if let Some(sibling) = without_suffix(parent, "-Embedding-GGUF") { + candidates.push(sibling); + } + if let Some(sibling) = without_suffix(parent, "-GGUF") { + candidates.push(sibling); + } + if let Some(sibling) = without_suffix(parent, "-gguf") { + candidates.push(sibling); + } + } + + let mut unique = Vec::new(); + for candidate in candidates { + if !unique + .iter() + .any(|existing: &PathBuf| existing == &candidate) + { + unique.push(candidate); + } + } + + for candidate in unique { + if has_required_files(&candidate) { + return Ok(candidate); + } + } + + Err(anyhow!( + "unable to infer tokenizer directory for gguf artifact {}; provide --tokenizer-dir", + gguf_file.display() + )) +} + +fn load_all_minilm_gguf_tensors( + gguf: &mut crate::models::common::gguf::Gguf, + cfg: &AllMiniLML6V2Config, + device: &Device, + dtype: DType, +) -> Result> { + let mut tensors = HashMap::new(); + + insert_gguf_tensor( + gguf, + &mut tensors, + "token_embd.weight", + "embeddings.word_embeddings.weight", + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + "token_types.weight", + "embeddings.token_type_embeddings.weight", + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + "position_embd.weight", + "embeddings.position_embeddings.weight", + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + "token_embd_norm.weight", + "embeddings.LayerNorm.weight", + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + "token_embd_norm.bias", + "embeddings.LayerNorm.bias", + device, + dtype, + )?; + + for layer_idx in 0..cfg.base.num_hidden_layers { + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_q.weight"), + &format!("encoder.layer.{layer_idx}.attention.self.query.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_q.bias"), + &format!("encoder.layer.{layer_idx}.attention.self.query.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_k.weight"), + &format!("encoder.layer.{layer_idx}.attention.self.key.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_k.bias"), + &format!("encoder.layer.{layer_idx}.attention.self.key.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_v.weight"), + &format!("encoder.layer.{layer_idx}.attention.self.value.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_v.bias"), + &format!("encoder.layer.{layer_idx}.attention.self.value.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_output.weight"), + &format!("encoder.layer.{layer_idx}.attention.output.dense.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_output.bias"), + &format!("encoder.layer.{layer_idx}.attention.output.dense.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_output_norm.weight"), + &format!("encoder.layer.{layer_idx}.attention.output.LayerNorm.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.attn_output_norm.bias"), + &format!("encoder.layer.{layer_idx}.attention.output.LayerNorm.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.ffn_up.weight"), + &format!("encoder.layer.{layer_idx}.intermediate.dense.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.ffn_up.bias"), + &format!("encoder.layer.{layer_idx}.intermediate.dense.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.ffn_down.weight"), + &format!("encoder.layer.{layer_idx}.output.dense.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.ffn_down.bias"), + &format!("encoder.layer.{layer_idx}.output.dense.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.layer_output_norm.weight"), + &format!("encoder.layer.{layer_idx}.output.LayerNorm.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + gguf, + &mut tensors, + &format!("blk.{layer_idx}.layer_output_norm.bias"), + &format!("encoder.layer.{layer_idx}.output.LayerNorm.bias"), + device, + dtype, + )?; + } + + Ok(tensors) +} + +fn insert_gguf_tensor( + gguf: &mut crate::models::common::gguf::Gguf, + tensors: &mut HashMap, + gguf_name: &str, + bert_name: &str, + device: &Device, + dtype: DType, +) -> Result<()> { + let tensor = gguf + .get_dequantized(gguf_name) + .map_err(|err| anyhow!("failed to load gguf tensor {}: {}", gguf_name, err))? + .to_device(device) + .map_err(|err| anyhow!("failed to move gguf tensor {}: {}", gguf_name, err))?; + let tensor = tensor + .to_dtype(dtype) + .map_err(|err| anyhow!("failed to convert gguf tensor {}: {}", gguf_name, err))?; + tensors.insert(bert_name.to_string(), tensor); + Ok(()) +} + +fn pool_hidden_state( + hidden: &[Vec], + pooling: AllMiniLML6V2PoolingStrategy, +) -> Result> { + let first = hidden + .first() + .ok_or_else(|| anyhow!("embedding hidden state is empty"))?; + let width = first.len(); + for row in hidden { + if row.len() != width { + return Err(anyhow!("inconsistent embedding width in hidden state")); + } + } + match pooling { + AllMiniLML6V2PoolingStrategy::Cls => Ok(first.clone()), + AllMiniLML6V2PoolingStrategy::Mean => { + let mut pooled = vec![0.0f32; width]; + for row in hidden { + for (idx, value) in row.iter().enumerate() { + pooled[idx] += *value; + } + } + let inv = 1.0f32 / hidden.len() as f32; + for value in &mut pooled { + *value *= inv; + } + Ok(pooled) + } + AllMiniLML6V2PoolingStrategy::Max => { + let mut pooled = first.clone(); + for row in hidden.iter().skip(1) { + for (idx, value) in row.iter().enumerate() { + pooled[idx] = pooled[idx].max(*value); + } + } + Ok(pooled) + } + AllMiniLML6V2PoolingStrategy::MeanSqrtLen => { + let mut pooled = vec![0.0f32; width]; + for row in hidden { + for (idx, value) in row.iter().enumerate() { + pooled[idx] += *value; + } + } + let inv = 1.0f32 / (hidden.len() as f32).sqrt(); + for value in &mut pooled { + *value *= inv; + } + Ok(pooled) + } + } +} + +#[cfg(feature = "onnx-runtime")] +fn build_i64_like_input( + desc: &OnnxInputDescriptor, + shape: Vec, + data: &[i64], +) -> Result { + let kind = desc.kind.ok_or_else(|| { + anyhow!( + "unsupported all-minilm-l6-v2 onnx input dtype for {}", + desc.name + ) + })?; + match kind { + OnnxTensorKind::I32 => Ok(ort::value::Tensor::from_array(( + shape, + data.iter().map(|v| *v as i32).collect::>(), + ))? + .into_dyn()), + OnnxTensorKind::I64 => { + Ok(ort::value::Tensor::from_array((shape, data.to_vec()))?.into_dyn()) + } + OnnxTensorKind::Bool => Ok(ort::value::Tensor::from_array(( + shape, + data.iter().map(|v| *v != 0).collect::>(), + ))? + .into_dyn()), + _ => Err(anyhow!( + "unsupported integer-like onnx input dtype for {}", + desc.name + )), + } +} + +#[cfg(feature = "onnx-runtime")] +fn extract_embedding_output( + value: &ort::value::DynValue, + seq_len: usize, + pooling: AllMiniLML6V2PoolingStrategy, +) -> Result> { + if let Ok((shape, values)) = value.try_extract_tensor::() { + let shape_vec = shape.iter().copied().collect::>(); + return extract_embedding_from_shape(&shape_vec, values, seq_len, pooling); + } + if let Ok((shape, values)) = value.try_extract_tensor::() { + let shape_vec = shape.iter().copied().collect::>(); + let values = values + .iter() + .map(|value| value.to_f32()) + .collect::>(); + return extract_embedding_from_shape(&shape_vec, &values, seq_len, pooling); + } + Err(anyhow!( + "all-minilm-l6-v2 onnx output must be a f32/f16 tensor" + )) +} + +#[cfg(feature = "onnx-runtime")] +fn extract_embedding_from_shape( + shape: &[i64], + values: &[f32], + seq_len: usize, + pooling: AllMiniLML6V2PoolingStrategy, +) -> Result> { + match shape { + [1, token_count, hidden_size] => { + let token_count = *token_count as usize; + let hidden_size = *hidden_size as usize; + let effective_tokens = token_count.min(seq_len); + let mut hidden = Vec::with_capacity(effective_tokens); + for chunk in values.chunks(hidden_size).take(effective_tokens) { + hidden.push(chunk.to_vec()); + } + pool_hidden_state(&hidden, pooling) + } + [token_count, hidden_size] if *token_count as usize == seq_len => { + let hidden_size = *hidden_size as usize; + let mut hidden = Vec::with_capacity(seq_len); + for chunk in values.chunks(hidden_size).take(seq_len) { + hidden.push(chunk.to_vec()); + } + pool_hidden_state(&hidden, pooling) + } + [1, hidden_size] => Ok(values.iter().take(*hidden_size as usize).copied().collect()), + [hidden_size] => Ok(values.iter().take(*hidden_size as usize).copied().collect()), + _ => Err(anyhow!( + "unexpected all-minilm-l6-v2 onnx output shape: {:?}", + shape + )), + } +} diff --git a/src/models/core/artifact.rs b/src/models/core/artifact.rs index 0d27bad..dcc386a 100644 --- a/src/models/core/artifact.rs +++ b/src/models/core/artifact.rs @@ -28,6 +28,11 @@ pub struct LoadSpec { pub fn supported_artifacts(model: WhichModel) -> &'static [ArtifactKind] { match model { + WhichModel::AllMiniLML6V2 => &[ + ArtifactKind::Safetensors, + ArtifactKind::Gguf, + ArtifactKind::Onnx, + ], WhichModel::Qwen3_0_6B => &[ ArtifactKind::Safetensors, ArtifactKind::Gguf, diff --git a/src/models/core/registry.rs b/src/models/core/registry.rs index 084fd66..87584a6 100644 --- a/src/models/core/registry.rs +++ b/src/models/core/registry.rs @@ -1,15 +1,16 @@ use anyhow::Result; use crate::models::{ - ModelInstance, WhichModel, load_model_legacy, qwen3::generate::Qwen3GenerateModel, - qwen3_5::generate::Qwen3_5GenerateModel, qwen3_embedding::generate::Qwen3EmbeddingModel, - qwen3_reranker::generate::Qwen3RerankerModel, + ModelInstance, WhichModel, all_minilm_l6_v2::generate::AllMiniLML6V2Model, load_model_legacy, + qwen3::generate::Qwen3GenerateModel, qwen3_5::generate::Qwen3_5GenerateModel, + qwen3_embedding::generate::Qwen3EmbeddingModel, qwen3_reranker::generate::Qwen3RerankerModel, }; use super::artifact::LoadSpec; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ModelLoaderFamily { + AllMiniLML6V2, Qwen3, Qwen3Embedding, Qwen3Reranker, @@ -19,6 +20,7 @@ enum ModelLoaderFamily { fn resolve_model_loader_family(model: WhichModel) -> ModelLoaderFamily { match model { + WhichModel::AllMiniLML6V2 => ModelLoaderFamily::AllMiniLML6V2, WhichModel::Qwen3_0_6B => ModelLoaderFamily::Qwen3, WhichModel::Qwen3Embedding0_6B | WhichModel::Qwen3Embedding4B @@ -44,6 +46,9 @@ fn resolve_model_loader_family(model: WhichModel) -> ModelLoaderFamily { pub fn load_model_from_spec<'a>(spec: &LoadSpec) -> Result> { spec.validate()?; let model = match resolve_model_loader_family(spec.model) { + ModelLoaderFamily::AllMiniLML6V2 => { + ModelInstance::AllMiniLML6V2(AllMiniLML6V2Model::init_from_spec(spec, None, None)?) + } ModelLoaderFamily::Qwen3 => { ModelInstance::Qwen3(Qwen3GenerateModel::init_from_spec(spec, None, None)?) } @@ -73,6 +78,10 @@ mod tests { #[test] fn registry_routes_qwen3_family_models() { + assert_eq!( + resolve_model_loader_family(WhichModel::AllMiniLML6V2), + ModelLoaderFamily::AllMiniLML6V2 + ); assert_eq!( resolve_model_loader_family(WhichModel::Qwen3_0_6B), ModelLoaderFamily::Qwen3 diff --git a/src/models/mod.rs b/src/models/mod.rs index 1556e42..d61a26a 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,3 +1,4 @@ +pub mod all_minilm_l6_v2; pub mod bigvgan; pub mod campplus; pub mod common; @@ -30,6 +31,7 @@ use anyhow::{Result, anyhow}; use rocket::futures::Stream; use crate::models::{ + all_minilm_l6_v2::generate::AllMiniLML6V2Model, deepseek_ocr::generate::DeepseekOCRGenerateModel, fun_asr_nano::generate::FunAsrNanoGenerateModel, glm_asr_nano::generate::GlmAsrNanoGenerateModel, glm_ocr::generate::GlmOcrGenerateModel, @@ -63,6 +65,8 @@ pub enum WhichModel { Qwen2_5vl7B, #[value(name = "qwen3-0.6b", hide = true)] Qwen3_0_6B, + #[value(name = "all-minilm-l6-v2", hide = true)] + AllMiniLML6V2, #[value(name = "qwen3-embedding-0.6b", hide = true)] Qwen3Embedding0_6B, #[value(name = "qwen3-embedding-4b", hide = true)] @@ -138,6 +142,7 @@ pub const LISTED_MODELS: &[WhichModel] = &[ WhichModel::Qwen2_5vl3B, WhichModel::Qwen2_5vl7B, WhichModel::Qwen3_0_6B, + WhichModel::AllMiniLML6V2, WhichModel::Qwen3Embedding0_6B, WhichModel::Qwen3Embedding4B, WhichModel::Qwen3Embedding8B, @@ -188,10 +193,11 @@ impl WhichModel { } pub fn is_download_managed(self) -> bool { - !matches!( - self.artifact_format(), - ModelArtifactFormat::Gguf | ModelArtifactFormat::Onnx - ) + !matches!(self, WhichModel::AllMiniLML6V2) + && !matches!( + self.artifact_format(), + ModelArtifactFormat::Gguf | ModelArtifactFormat::Onnx + ) } pub fn openai_model_id(self) -> &'static str { @@ -200,6 +206,7 @@ impl WhichModel { WhichModel::Qwen2_5vl3B => "qwen2.5vl-3b", WhichModel::Qwen2_5vl7B => "qwen2.5vl-7b", WhichModel::Qwen3_0_6B => "qwen3-0.6b", + WhichModel::AllMiniLML6V2 => "all-minilm-l6-v2", WhichModel::Qwen3Embedding0_6B => "qwen3-embedding-0.6b", WhichModel::Qwen3Embedding4B => "qwen3-embedding-4b", WhichModel::Qwen3Embedding8B => "qwen3-embedding-8b", @@ -241,6 +248,7 @@ impl WhichModel { match self { WhichModel::MiniCPM4_0_5B => "OpenBMB", WhichModel::Qwen2_5vl3B | WhichModel::Qwen2_5vl7B => "Qwen", + WhichModel::AllMiniLML6V2 => "sentence-transformers", WhichModel::Qwen3_0_6B | WhichModel::Qwen3Embedding0_6B | WhichModel::Qwen3Embedding4B @@ -282,6 +290,7 @@ impl WhichModel { WhichModel::Qwen2_5vl3B => "Qwen/Qwen2.5-VL-3B-Instruct", WhichModel::Qwen2_5vl7B => "Qwen/Qwen2.5-VL-7B-Instruct", WhichModel::Qwen3_0_6B => "Qwen/Qwen3-0.6B", + WhichModel::AllMiniLML6V2 => "sentence-transformers/all-MiniLM-L6-v2", WhichModel::Qwen3Embedding0_6B => "Qwen/Qwen3-Embedding-0.6B", WhichModel::Qwen3Embedding4B => "Qwen/Qwen3-Embedding-4B", WhichModel::Qwen3Embedding8B => "Qwen/Qwen3-Embedding-8B", @@ -324,7 +333,8 @@ impl WhichModel { match self { // LLM models WhichModel::MiniCPM4_0_5B | WhichModel::Qwen3_0_6B => "llm", - WhichModel::Qwen3Embedding0_6B + WhichModel::AllMiniLML6V2 + | WhichModel::Qwen3Embedding0_6B | WhichModel::Qwen3Embedding4B | WhichModel::Qwen3Embedding8B => "embedding", WhichModel::Qwen3Reranker0_6B @@ -384,6 +394,7 @@ pub enum ModelInstance<'a> { MiniCPM4(MiniCPMGenerateModel<'a>), Qwen2_5VL(Qwen2_5VLGenerateModel<'a>), Qwen3(Qwen3GenerateModel<'a>), + AllMiniLML6V2(AllMiniLML6V2Model), Qwen3Embedding(Qwen3EmbeddingModel), Qwen3Reranker(Qwen3RerankerModel), Qwen3_5(Qwen3_5GenerateModel<'a>), @@ -405,6 +416,9 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::MiniCPM4(model) => model.generate(mes), ModelInstance::Qwen2_5VL(model) => model.generate(mes), ModelInstance::Qwen3(model) => model.generate(mes), + ModelInstance::AllMiniLML6V2(_) => { + Err(anyhow!("embedding model does not support chat completions")) + } ModelInstance::Qwen3Embedding(_) => { Err(anyhow!("embedding model does not support chat completions")) } @@ -440,6 +454,9 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::MiniCPM4(model) => model.generate_stream(mes), ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes), ModelInstance::Qwen3(model) => model.generate_stream(mes), + ModelInstance::AllMiniLML6V2(_) => Err(anyhow!( + "embedding model does not support streaming chat completions" + )), ModelInstance::Qwen3Embedding(_) => Err(anyhow!( "embedding model does not support streaming chat completions" )), @@ -464,6 +481,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { impl<'a> ModelInstance<'a> { pub fn embedding(&mut self, input: &[String]) -> Result>> { match self { + ModelInstance::AllMiniLML6V2(model) => model.embed(input), ModelInstance::Qwen3Embedding(model) => model.embed(input), _ => Err(anyhow!("current model does not support embeddings")), } @@ -504,6 +522,10 @@ pub fn load_model_legacy<'a>( let model = Qwen3GenerateModel::init(path, None, None)?; ModelInstance::Qwen3(model) } + WhichModel::AllMiniLML6V2 => { + let model = AllMiniLML6V2Model::init(path, None, None)?; + ModelInstance::AllMiniLML6V2(model) + } WhichModel::Qwen3Embedding0_6B | WhichModel::Qwen3Embedding4B | WhichModel::Qwen3Embedding8B => { diff --git a/tests/test_all_minilm_l6_v2_multi_format.rs b/tests/test_all_minilm_l6_v2_multi_format.rs new file mode 100644 index 0000000..3c6dab8 --- /dev/null +++ b/tests/test_all_minilm_l6_v2_multi_format.rs @@ -0,0 +1,353 @@ +use std::path::{Path, PathBuf}; + +use aha::models::{ + ArtifactKind, LoadSpec, ModelPaths, WhichModel, + all_minilm_l6_v2::generate::AllMiniLML6V2Model, + common::{onnx::ensure_ort_dylib_path, retrieval::cosine_similarity}, +}; +use anyhow::{Context, Result, anyhow}; +#[cfg(feature = "onnx-runtime")] +use ort::session::Session; + +const ALL_MINILM_L6_V2_DIR: &str = r"D:\model_download\all-MiniLM-L6-v2"; +const ALL_MINILM_L6_V2_GGUF_DIR: &str = r"D:\model_download\All-MiniLM-L6-v2-Embedding-GGUF"; +const ALL_MINILM_L6_V2_ONNX_DIR: &str = r"D:\model_download\all-MiniLM-L6-v2\onnx"; + +fn require_existing_dir(path: &str) -> Result<()> { + let dir = Path::new(path); + if !dir.exists() { + return Err(anyhow!("model dir not found: {}", path)); + } + if !dir.is_dir() { + return Err(anyhow!("path is not a directory: {}", path)); + } + Ok(()) +} + +fn first_file_with_extension_recursive(dir: &str, extension: &str) -> Result { + require_existing_dir(dir)?; + + let mut stack = vec![PathBuf::from(dir)]; + let mut matches = Vec::new(); + while let Some(current) = stack.pop() { + for entry in std::fs::read_dir(¤t)? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + stack.push(path); + continue; + } + if path + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case(extension)) + { + matches.push(path); + } + } + } + + matches.sort(); + matches + .into_iter() + .next() + .ok_or_else(|| anyhow!("no .{} file found (recursive) in {}", extension, dir)) +} + +fn first_file_with_extension(dir: &str, extension: &str) -> Result { + require_existing_dir(dir)?; + + let mut matches = std::fs::read_dir(dir)? + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + path.is_file() + && path + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case(extension)) + }) + .collect::>(); + + matches.sort(); + matches + .into_iter() + .next() + .ok_or_else(|| anyhow!("no .{} file found in {}", extension, dir)) +} + +#[test] +fn all_minilm_l6_v2_safetensors_can_load() -> Result<()> { + require_existing_dir(ALL_MINILM_L6_V2_DIR)?; + + let _model = AllMiniLML6V2Model::init(ALL_MINILM_L6_V2_DIR, None, None).with_context(|| { + format!( + "failed to init all-minilm-l6-v2 safetensors model from {}", + ALL_MINILM_L6_V2_DIR + ) + })?; + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_gguf_file_can_load() -> Result<()> { + let gguf_path = first_file_with_extension(ALL_MINILM_L6_V2_GGUF_DIR, "gguf")?; + let metadata = std::fs::metadata(&gguf_path) + .with_context(|| format!("failed to read gguf metadata: {}", gguf_path.display()))?; + if metadata.len() == 0 { + return Err(anyhow!("gguf file is empty: {}", gguf_path.display())); + } + let _bytes = std::fs::read(&gguf_path) + .with_context(|| format!("failed to read gguf file: {}", gguf_path.display()))?; + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_onnx_file_can_load() -> Result<()> { + let onnx_path = first_file_with_extension_recursive(ALL_MINILM_L6_V2_ONNX_DIR, "onnx")?; + let metadata = std::fs::metadata(&onnx_path) + .with_context(|| format!("failed to read onnx metadata: {}", onnx_path.display()))?; + if metadata.len() == 0 { + return Err(anyhow!("onnx file is empty: {}", onnx_path.display())); + } + let _bytes = std::fs::read(&onnx_path) + .with_context(|| format!("failed to read onnx file: {}", onnx_path.display()))?; + Ok(()) +} + +#[cfg(feature = "onnx-runtime")] +#[test] +fn all_minilm_l6_v2_onnxruntime_can_create_session() -> Result<()> { + let onnx_path = first_file_with_extension_recursive(ALL_MINILM_L6_V2_ONNX_DIR, "onnx")?; + + if let Err(err) = ensure_ort_dylib_path() { + println!("skip onnxruntime session test: {err}"); + return Ok(()); + } + + let session = Session::builder() + .context("failed to create onnxruntime session builder")? + .commit_from_file(&onnx_path) + .with_context(|| { + format!( + "failed to create onnxruntime session from {}", + onnx_path.display() + ) + })?; + + if session.inputs().is_empty() { + return Err(anyhow!("onnxruntime session has no inputs")); + } + if session.outputs().is_empty() { + return Err(anyhow!("onnxruntime session has no outputs")); + } + + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_safetensors_init_from_spec_can_embed() -> Result<()> { + require_existing_dir(ALL_MINILM_L6_V2_DIR)?; + + let spec = LoadSpec { + model: WhichModel::AllMiniLML6V2, + artifact: ArtifactKind::Safetensors, + paths: ModelPaths { + weight_dir: Some(ALL_MINILM_L6_V2_DIR.to_string()), + ..Default::default() + }, + }; + + let mut model = AllMiniLML6V2Model::init_from_spec(&spec, None, None)?; + let output = model.embed(&["test safetensors embedding".to_string()])?; + assert_eq!(output.len(), 1); + assert_eq!(output[0].len(), 384); + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_gguf_init_from_spec_can_embed() -> Result<()> { + let gguf_path = match first_file_with_extension(ALL_MINILM_L6_V2_GGUF_DIR, "gguf") { + Ok(path) => path, + Err(err) => { + println!("skip gguf init_from_spec test: {err}"); + return Ok(()); + } + }; + + let spec = LoadSpec { + model: WhichModel::AllMiniLML6V2, + artifact: ArtifactKind::Gguf, + paths: ModelPaths { + gguf_path: Some(gguf_path.to_string_lossy().to_string()), + tokenizer_dir: Some(ALL_MINILM_L6_V2_DIR.to_string()), + ..Default::default() + }, + }; + + let mut model = AllMiniLML6V2Model::init_from_spec(&spec, None, None)?; + let output = model.embed(&["test gguf embedding".to_string()])?; + assert_eq!(output.len(), 1); + assert_eq!(output[0].len(), 384); + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_onnx_init_from_spec_can_embed() -> Result<()> { + if let Err(err) = ensure_ort_dylib_path() { + println!("skip onnx init test: {err}"); + return Ok(()); + } + + let spec = LoadSpec { + model: WhichModel::AllMiniLML6V2, + artifact: ArtifactKind::Onnx, + paths: ModelPaths { + onnx_path: Some(ALL_MINILM_L6_V2_ONNX_DIR.to_string()), + tokenizer_dir: Some(ALL_MINILM_L6_V2_DIR.to_string()), + ..Default::default() + }, + }; + + let mut model = AllMiniLML6V2Model::init_from_spec(&spec, None, None)?; + let output = model.embed(&["test onnx embedding".to_string()])?; + assert_eq!(output.len(), 1); + assert_eq!(output[0].len(), 384); + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_load_spec_accepts_gguf() { + let spec = LoadSpec { + model: WhichModel::AllMiniLML6V2, + artifact: ArtifactKind::Gguf, + paths: ModelPaths { + gguf_path: Some(format!( + "{ALL_MINILM_L6_V2_GGUF_DIR}\\all-MiniLM-L6-v2-ggml-model-f16.gguf" + )), + tokenizer_dir: Some(ALL_MINILM_L6_V2_DIR.to_string()), + ..Default::default() + }, + }; + + spec.validate() + .expect("all-minilm-l6-v2 should accept gguf artifact"); +} + +#[test] +fn all_minilm_l6_v2_load_spec_accepts_onnx() { + let spec = LoadSpec { + model: WhichModel::AllMiniLML6V2, + artifact: ArtifactKind::Onnx, + paths: ModelPaths { + onnx_path: Some(ALL_MINILM_L6_V2_ONNX_DIR.to_string()), + tokenizer_dir: Some(ALL_MINILM_L6_V2_DIR.to_string()), + ..Default::default() + }, + }; + + spec.validate() + .expect("all-minilm-l6-v2 should accept onnx artifact"); +} + +#[test] +fn all_minilm_l6_v2_native_and_gguf_embeddings_are_close() -> Result<()> { + let gguf_path = match first_file_with_extension(ALL_MINILM_L6_V2_GGUF_DIR, "gguf") { + Ok(path) => path, + Err(err) => { + println!("skip native/gguf similarity test: {err}"); + return Ok(()); + } + }; + + let text = "Rust provides strong ownership guarantees for concurrent systems."; + let mut native_model = AllMiniLML6V2Model::init(ALL_MINILM_L6_V2_DIR, None, None)?; + let mut gguf_model = AllMiniLML6V2Model::init_gguf( + &gguf_path.to_string_lossy(), + Some(ALL_MINILM_L6_V2_DIR), + None, + None, + )?; + + let native_embedding = native_model.embed(&[text.to_string()])?; + let gguf_embedding = gguf_model.embed(&[text.to_string()])?; + + let similarity = cosine_similarity(&native_embedding[0], &gguf_embedding[0])?; + assert!( + similarity > 0.98, + "native/gguf embedding similarity too low: {similarity}" + ); + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_native_and_onnx_embeddings_are_close() -> Result<()> { + if let Err(err) = ensure_ort_dylib_path() { + println!("skip cross-backend similarity test: {err}"); + return Ok(()); + } + + let text = "Rust provides strong ownership guarantees for concurrent systems."; + let mut native_model = AllMiniLML6V2Model::init(ALL_MINILM_L6_V2_DIR, None, None)?; + let mut onnx_model = + AllMiniLML6V2Model::init_onnx(ALL_MINILM_L6_V2_ONNX_DIR, Some(ALL_MINILM_L6_V2_DIR))?; + + let native_embedding = native_model.embed(&[text.to_string()])?; + let onnx_embedding = onnx_model.embed(&[text.to_string()])?; + + let similarity = cosine_similarity(&native_embedding[0], &onnx_embedding[0])?; + assert!( + similarity > 0.98, + "native/onnx embedding similarity too low: {similarity}" + ); + Ok(()) +} + +#[test] +fn all_minilm_l6_v2_real_texts_similarity() -> Result<()> { + require_existing_dir(ALL_MINILM_L6_V2_DIR)?; + + let query = "How do I send asynchronous HTTP requests in Rust?"; + let documents = vec![ + "In Rust, reqwest with tokio is a common way to send async HTTP requests.".to_string(), + "The weather is sunny today, which is great for a walk in the park.".to_string(), + "PostgreSQL indexes can improve database query latency on large tables.".to_string(), + "A guitar usually has six strings and is used in many styles of music.".to_string(), + ]; + + let mut model = AllMiniLML6V2Model::init(ALL_MINILM_L6_V2_DIR, None, None)?; + + let mut inputs = vec![query.to_string()]; + inputs.extend(documents.clone()); + let embeddings = model.embed(&inputs)?; + + if embeddings.len() != inputs.len() { + return Err(anyhow!( + "embedding count mismatch: got {}, expect {}", + embeddings.len(), + inputs.len() + )); + } + + let query_embedding = &embeddings[0]; + let mut best_idx = 0usize; + let mut best_score = f32::NEG_INFINITY; + for (doc_idx, doc_emb) in embeddings.iter().enumerate().skip(1) { + let score = cosine_similarity(query_embedding, doc_emb)?; + println!( + "similarity(query, doc_{}) = {:.6}, doc = {}", + doc_idx - 1, + score, + documents[doc_idx - 1] + ); + if score > best_score { + best_score = score; + best_idx = doc_idx - 1; + } + } + + println!("best_match_doc_index={}", best_idx); + println!("best_match_doc={}", documents[best_idx]); + println!("best_match_score={:.6}", best_score); + + assert_eq!(best_idx, 0, "unexpected top match for query"); + Ok(()) +} diff --git a/tests/test_load_spec.rs b/tests/test_load_spec.rs index c9a48dd..2a659dd 100644 --- a/tests/test_load_spec.rs +++ b/tests/test_load_spec.rs @@ -15,6 +15,39 @@ fn load_spec_auto_resolves_to_safetensors_default() { assert!(spec.validate().is_ok()); } +#[test] +fn load_spec_all_minilm_accepts_onnx() { + let spec = LoadSpec { + model: WhichModel::AllMiniLML6V2, + artifact: ArtifactKind::Onnx, + paths: ModelPaths { + onnx_path: Some("D:/model_download/all-MiniLM-L6-v2/onnx".to_string()), + tokenizer_dir: Some("D:/model_download/all-MiniLM-L6-v2".to_string()), + ..Default::default() + }, + }; + + assert!(spec.validate().is_ok()); +} + +#[test] +fn load_spec_all_minilm_accepts_gguf() { + let spec = LoadSpec { + model: WhichModel::AllMiniLML6V2, + artifact: ArtifactKind::Gguf, + paths: ModelPaths { + gguf_path: Some( + "D:/model_download/All-MiniLM-L6-v2-Embedding-GGUF/all-MiniLM-L6-v2-ggml-model-f16.gguf" + .to_string(), + ), + tokenizer_dir: Some("D:/model_download/all-MiniLM-L6-v2".to_string()), + ..Default::default() + }, + }; + + assert!(spec.validate().is_ok()); +} + #[test] fn load_spec_gguf_requires_gguf_path() { let spec = LoadSpec {