97 lines
2.6 KiB
Rust
97 lines
2.6 KiB
Rust
// BGE-M3 编码管线 — 通过 HTTP 调用本地 ONNX BGE 服务 (localhost:8000)
|
|
// 使用 OpenAI-compatible /v1/embeddings 端点
|
|
// 替代方案: 等 ort 2.0 stable 后用 Rust 原生 ONNX
|
|
|
|
use reqwest::blocking::Client;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::time::Duration;
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct EmbedRequest {
|
|
input: EmbedInput,
|
|
model: String,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
#[serde(untagged)]
|
|
enum EmbedInput {
|
|
Single(String),
|
|
Batch(Vec<String>),
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct EmbedResponse {
|
|
data: Vec<EmbedData>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct EmbedData {
|
|
embedding: Vec<f32>,
|
|
}
|
|
|
|
pub struct BGEEncoder {
|
|
pub model_dir: String,
|
|
pub dim: usize,
|
|
endpoint: String,
|
|
client: Client,
|
|
}
|
|
|
|
impl BGEEncoder {
|
|
pub fn new(model_dir: &str) -> Result<Self, Box<dyn std::error::Error>> {
|
|
Ok(Self {
|
|
model_dir: model_dir.to_string(),
|
|
dim: 1024,
|
|
endpoint: "http://localhost:8000/v1/embeddings".to_string(),
|
|
client: Client::builder()
|
|
.timeout(Duration::from_secs(30))
|
|
.build()?,
|
|
})
|
|
}
|
|
|
|
pub fn encode(&self, text: &str) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
|
|
let results = self.encode_batch(&[text.to_string()])?;
|
|
results
|
|
.into_iter()
|
|
.next()
|
|
.ok_or_else(|| "BGE encode: empty response".into())
|
|
}
|
|
|
|
pub fn encode_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
|
|
if texts.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let req = EmbedRequest {
|
|
input: EmbedInput::Batch(texts.to_vec()),
|
|
model: "bge-m3".to_string(),
|
|
};
|
|
|
|
let resp = self
|
|
.client
|
|
.post(&self.endpoint)
|
|
.header("Content-Type", "application/json")
|
|
.timeout(Duration::from_secs(10))
|
|
.body(serde_json::to_string(&req)?)
|
|
.send()?;
|
|
|
|
if !resp.status().is_success() {
|
|
return Err(format!("BGE HTTP error: {} — {}", resp.status(), resp.text().unwrap_or_default()).into());
|
|
}
|
|
|
|
let body = resp.text()?;
|
|
let parsed: EmbedResponse = serde_json::from_str(&body)?;
|
|
|
|
let vectors: Vec<Vec<f32>> = parsed.data.into_iter().map(|d| d.embedding).collect();
|
|
|
|
if vectors.len() != texts.len() {
|
|
return Err(format!(
|
|
"BGE encode: expected {} vectors, got {}",
|
|
texts.len(),
|
|
vectors.len()
|
|
).into());
|
|
}
|
|
|
|
Ok(vectors)
|
|
}
|
|
}
|