添加 All-MiniLM-L6-v2-Embedding 模型
This commit is contained in:
@@ -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 |
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 模型
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
+42
@@ -428,6 +428,11 @@ fn resolve_load_spec_for_run(args: &RunArgs) -> anyhow::Result<LoadSpec> {
|
||||
|
||||
fn run_target_model_with_spec(args: &RunArgs, spec: &LoadSpec) -> anyhow::Result<bool> {
|
||||
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([
|
||||
|
||||
@@ -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<Self> {
|
||||
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::<SentenceBertConfig>(&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<ModuleEntry> = 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,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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<DType>,
|
||||
) -> Result<Self> {
|
||||
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<DType>) -> Result<Self> {
|
||||
let backend = AllMiniLML6V2Backend::load(path, device, dtype)?;
|
||||
Ok(Self { backend })
|
||||
}
|
||||
|
||||
pub fn init_onnx(onnx_path: &str, tokenizer_dir: Option<&str>) -> Result<Self> {
|
||||
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<DType>,
|
||||
) -> Result<Self> {
|
||||
let backend = AllMiniLML6V2Backend::load_gguf(gguf_path, tokenizer_dir, device, dtype)?;
|
||||
Ok(Self { backend })
|
||||
}
|
||||
|
||||
pub fn embed(&mut self, input: &[String]) -> Result<Vec<Vec<f32>>> {
|
||||
self.backend.embed_texts(input)
|
||||
}
|
||||
}
|
||||
|
||||
impl TextEmbeddingBackend for AllMiniLML6V2Model {
|
||||
fn embed_texts(&mut self, input: &[String]) -> Result<Vec<Vec<f32>>> {
|
||||
self.backend.embed_texts(input)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod config;
|
||||
pub mod generate;
|
||||
pub mod model;
|
||||
@@ -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<DType>) -> Result<Self> {
|
||||
Ok(Self::Safetensors(AllMiniLML6V2SafetensorsBackend::load(
|
||||
path, device, dtype,
|
||||
)?))
|
||||
}
|
||||
|
||||
pub fn load_onnx(onnx_path: &str, tokenizer_dir: Option<&str>) -> Result<Self> {
|
||||
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<DType>,
|
||||
) -> Result<Self> {
|
||||
Ok(Self::Gguf(AllMiniLML6V2GgufBackend::load(
|
||||
gguf_path,
|
||||
tokenizer_dir,
|
||||
device,
|
||||
dtype,
|
||||
)?))
|
||||
}
|
||||
|
||||
pub fn embed_texts(&mut self, input: &[String]) -> Result<Vec<Vec<f32>>> {
|
||||
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<DType>) -> Result<Self> {
|
||||
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<Vec<Vec<f32>>> {
|
||||
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<Vec<f32>> {
|
||||
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::<f32>()?;
|
||||
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<DType>,
|
||||
) -> Result<Self> {
|
||||
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<Vec<Vec<f32>>> {
|
||||
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<Vec<f32>> {
|
||||
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::<f32>()?;
|
||||
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<String>,
|
||||
input_descriptors: Vec<OnnxInputDescriptor>,
|
||||
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<i64>,
|
||||
kind: Option<OnnxTensorKind>,
|
||||
}
|
||||
|
||||
#[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<OnnxTensorKind> {
|
||||
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<Self> {
|
||||
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::<Vec<_>>(),
|
||||
map_tensor_kind(*ty),
|
||||
),
|
||||
_ => (Vec::new(), None),
|
||||
};
|
||||
OnnxInputDescriptor {
|
||||
name: input.name().to_string(),
|
||||
shape,
|
||||
kind,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
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<Vec<Vec<f32>>> {
|
||||
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<Vec<f32>> {
|
||||
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::<Vec<_>>();
|
||||
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::<Vec<_>>()
|
||||
.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<ort::value::DynValue> {
|
||||
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<i64> {
|
||||
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<i64>,
|
||||
) -> Result<ort::value::DynValue> {
|
||||
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::<Vec<_>>();
|
||||
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<Vec<f32>> {
|
||||
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<Vec<u32>> {
|
||||
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<String> {
|
||||
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::<Vec<_>>();
|
||||
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<PathBuf> {
|
||||
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<PathBuf> {
|
||||
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<R: std::io::Read + std::io::Seek>(
|
||||
gguf: &mut crate::models::common::gguf::Gguf<R>,
|
||||
cfg: &AllMiniLML6V2Config,
|
||||
device: &Device,
|
||||
dtype: DType,
|
||||
) -> Result<HashMap<String, Tensor>> {
|
||||
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<R: std::io::Read + std::io::Seek>(
|
||||
gguf: &mut crate::models::common::gguf::Gguf<R>,
|
||||
tensors: &mut HashMap<String, Tensor>,
|
||||
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<f32>],
|
||||
pooling: AllMiniLML6V2PoolingStrategy,
|
||||
) -> Result<Vec<f32>> {
|
||||
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<i64>,
|
||||
data: &[i64],
|
||||
) -> Result<ort::value::DynValue> {
|
||||
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::<Vec<_>>(),
|
||||
))?
|
||||
.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::<Vec<_>>(),
|
||||
))?
|
||||
.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<Vec<f32>> {
|
||||
if let Ok((shape, values)) = value.try_extract_tensor::<f32>() {
|
||||
let shape_vec = shape.iter().copied().collect::<Vec<_>>();
|
||||
return extract_embedding_from_shape(&shape_vec, values, seq_len, pooling);
|
||||
}
|
||||
if let Ok((shape, values)) = value.try_extract_tensor::<f16>() {
|
||||
let shape_vec = shape.iter().copied().collect::<Vec<_>>();
|
||||
let values = values
|
||||
.iter()
|
||||
.map(|value| value.to_f32())
|
||||
.collect::<Vec<_>>();
|
||||
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<Vec<f32>> {
|
||||
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
|
||||
)),
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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<ModelInstance<'a>> {
|
||||
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
|
||||
|
||||
+27
-5
@@ -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<Vec<Vec<f32>>> {
|
||||
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 => {
|
||||
|
||||
@@ -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<PathBuf> {
|
||||
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<PathBuf> {
|
||||
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::<Vec<_>>();
|
||||
|
||||
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(())
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user