273265088@qq.com
73 lines
2.5 KiB
Rust
73 lines
2.5 KiB
Rust
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)
|
|
}
|
|
}
|