use anyhow::{Context, Result};
use ndarray::Array2;
use ort::session::{builder::GraphOptimizationLevel, Session, SessionInputValue};
use ort::value::Value;
use std::borrow::Cow;
use std::path::Path;
use tokenizers::Tokenizer;
const MAX_LENGTH: usize = 256;
const BATCH_SIZE: usize = 16;
#[derive(Debug, Clone)]
pub struct DenseSearchResult {
pub chunk_id: String,
pub score: f32,
}
pub struct DenseRetriever {
session: Session,
tokenizer: Tokenizer,
}
impl DenseRetriever {
pub fn new(model_path: &Path) -> Result<Self> {
let model_file = model_path.join("model.onnx");
let tokenizer_file = model_path.join("tokenizer.json");
let session = Session::builder()?
.with_optimization_level(GraphOptimizationLevel::Level3)?
.commit_from_file(&model_file)
.context("Failed to load dense model")?;
let tokenizer = Tokenizer::from_file(tokenizer_file)
.map_err(|e| anyhow::anyhow!("Failed to load dense tokenizer: {}", e))?;
Ok(Self { session, tokenizer })
}
pub fn embed_texts(&mut self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(Vec::new());
}
let mut vectors = Vec::with_capacity(texts.len());
for batch in texts.chunks(BATCH_SIZE) {
let refs: Vec<&str> = batch.iter().map(|s| s.as_str()).collect();
let encodings = self
.tokenizer
.encode_batch(refs, true)
.map_err(|e| anyhow::anyhow!("Dense tokenization failed: {}", e))?;
let batch_len = encodings.len();
let max_len = encodings
.iter()
.map(|e| e.len())
.max()
.unwrap_or(1)
.min(MAX_LENGTH)
.max(1);
let mut input_ids_vec = vec![0i64; batch_len * max_len];
let mut attention_mask_vec = vec![0i64; batch_len * max_len];
let token_type_ids_vec = vec![0i64; batch_len * max_len];
for (i, encoding) in encodings.iter().enumerate() {
let ids = encoding.get_ids();
let mask = encoding.get_attention_mask();
let len = ids.len().min(max_len);
for j in 0..len {
input_ids_vec[i * max_len + j] = ids[j] as i64;
attention_mask_vec[i * max_len + j] = mask[j] as i64;
}
}
let input_ids: Array2<i64> =
Array2::from_shape_vec((batch_len, max_len), input_ids_vec)?;
let attention_mask: Array2<i64> =
Array2::from_shape_vec((batch_len, max_len), attention_mask_vec.clone())?;
let token_type_ids: Array2<i64> =
Array2::from_shape_vec((batch_len, max_len), token_type_ids_vec)?;
let input_ids_value = Value::from_array(input_ids)?;
let attention_mask_value = Value::from_array(attention_mask)?;
let token_type_ids_value = Value::from_array(token_type_ids)?;
let inputs: Vec<(Cow<'_, str>, SessionInputValue<'_>)> = vec![
(
Cow::Borrowed("input_ids"),
SessionInputValue::from(&input_ids_value),
),
(
Cow::Borrowed("attention_mask"),
SessionInputValue::from(&attention_mask_value),
),
(
Cow::Borrowed("token_type_ids"),
SessionInputValue::from(&token_type_ids_value),
),
];
let outputs = self.session.run(inputs)?;
if let Some(embeddings) = outputs
.get("sentence_embedding")
.or_else(|| outputs.get("embeddings"))
{
let arr = embeddings.try_extract_array::<f32>()?;
let shape = arr.shape();
if shape.len() != 2 {
return Err(anyhow::anyhow!("Unexpected embedding shape: {:?}", shape));
}
for i in 0..batch_len {
let mut v = Vec::with_capacity(shape[1]);
for j in 0..shape[1] {
v.push(arr[[i, j]]);
}
l2_normalize(&mut v);
vectors.push(v);
}
continue;
}
let last_hidden = outputs
.get("last_hidden_state")
.or_else(|| outputs.get("token_embeddings"))
.context("Dense model output not found")?;
let token_embeddings = last_hidden.try_extract_array::<f32>()?;
let shape = token_embeddings.shape().to_vec();
if shape.len() != 3 {
return Err(anyhow::anyhow!(
"Unexpected last_hidden_state shape: {:?}",
shape
));
}
let hidden_dim = shape[2];
for i in 0..batch_len {
let mut sum = vec![0.0f32; hidden_dim];
let mut count = 0.0f32;
for t in 0..max_len {
let mask = attention_mask_vec[i * max_len + t];
if mask == 0 {
continue;
}
count += 1.0;
for h in 0..hidden_dim {
sum[h] += token_embeddings[[i, t, h]];
}
}
if count > 0.0 {
for value in &mut sum {
*value /= count;
}
}
l2_normalize(&mut sum);
vectors.push(sum);
}
}
Ok(vectors)
}
pub fn search_embeddings(
&mut self,
query: &str,
embeddings: &[(String, Vec<f32>)],
limit: usize,
) -> Result<Vec<DenseSearchResult>> {
if embeddings.is_empty() || limit == 0 {
return Ok(Vec::new());
}
let query_embedding = self
.embed_texts(&[query.to_string()])?
.into_iter()
.next()
.unwrap_or_default();
if query_embedding.is_empty() {
return Ok(Vec::new());
}
let mut scored = Vec::with_capacity(embeddings.len());
for (chunk_id, embedding) in embeddings {
if embedding.len() != query_embedding.len() {
continue;
}
let score = cosine_similarity(&query_embedding, embedding);
scored.push(DenseSearchResult {
chunk_id: chunk_id.clone(),
score,
});
}
scored.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
scored.truncate(limit);
Ok(scored)
}
pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
Ok(self
.embed_texts(&[query.to_string()])?
.into_iter()
.next()
.unwrap_or_default())
}
}
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
let mut dot = 0.0f32;
let mut norm_a = 0.0f32;
let mut norm_b = 0.0f32;
for (&x, &y) in a.iter().zip(b.iter()) {
dot += x * y;
norm_a += x * x;
norm_b += y * y;
}
let denom = (norm_a.sqrt() * norm_b.sqrt()).max(1e-12);
dot / denom
}
fn l2_normalize(v: &mut [f32]) {
let mut norm = 0.0f32;
for &x in v.iter() {
norm += x * x;
}
norm = norm.sqrt().max(1e-12);
for x in v.iter_mut() {
*x /= norm;
}
}