wehub-resource-sync
a789495a98
FreeBSD Smoke / FreeBSD Smoke (x86_64) (push) Has been cancelled
CI / Quality Guardrails (push) Has been cancelled
CI / Build & Test (macos-latest) (push) Has been cancelled
CI / Build & Test (ubuntu-latest) (push) Has been cancelled
CI / Build & Test (windows-latest) (push) Has been cancelled
CI / Format (push) Has been cancelled
CI / PowerShell Syntax (push) Has been cancelled
CI / Windows Cross-Target Check (Linux) (push) Has been cancelled
583 lines
20 KiB
Rust
583 lines
20 KiB
Rust
use anyhow::{Context, Result};
|
|
use std::cmp::Reverse;
|
|
use std::collections::BinaryHeap;
|
|
use std::io::Write;
|
|
use std::path::Path;
|
|
use tokenizers::Tokenizer;
|
|
use tract_hir::infer::Factoid as _;
|
|
use tract_hir::prelude::*;
|
|
|
|
pub const MODEL_NAME: &str = "all-MiniLM-L6-v2";
|
|
type RunnableEmbeddingModel =
|
|
SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>;
|
|
|
|
#[derive(Debug)]
|
|
struct TopKItem<T> {
|
|
score: f32,
|
|
ordinal: usize,
|
|
value: T,
|
|
}
|
|
|
|
impl<T> PartialEq for TopKItem<T> {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.score.to_bits() == other.score.to_bits() && self.ordinal == other.ordinal
|
|
}
|
|
}
|
|
|
|
impl<T> Eq for TopKItem<T> {}
|
|
|
|
impl<T> PartialOrd for TopKItem<T> {
|
|
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
|
Some(self.cmp(other))
|
|
}
|
|
}
|
|
|
|
impl<T> Ord for TopKItem<T> {
|
|
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
|
|
self.score
|
|
.total_cmp(&other.score)
|
|
.then_with(|| self.ordinal.cmp(&other.ordinal))
|
|
}
|
|
}
|
|
|
|
fn top_k_scored<T, I>(items: I, limit: usize) -> Vec<(T, f32)>
|
|
where
|
|
I: IntoIterator<Item = (T, f32)>,
|
|
{
|
|
if limit == 0 {
|
|
return Vec::new();
|
|
}
|
|
|
|
let mut heap: BinaryHeap<Reverse<TopKItem<T>>> = BinaryHeap::new();
|
|
for (ordinal, (value, score)) in items.into_iter().enumerate() {
|
|
let candidate = Reverse(TopKItem {
|
|
score,
|
|
ordinal,
|
|
value,
|
|
});
|
|
|
|
if heap.len() < limit {
|
|
heap.push(candidate);
|
|
continue;
|
|
}
|
|
|
|
let replace = heap
|
|
.peek()
|
|
.map(|smallest| score > smallest.0.score)
|
|
.unwrap_or(false);
|
|
if replace {
|
|
heap.pop();
|
|
heap.push(candidate);
|
|
}
|
|
}
|
|
|
|
let mut results: Vec<_> = heap
|
|
.into_iter()
|
|
.map(|Reverse(item)| (item.value, item.score, item.ordinal))
|
|
.collect();
|
|
results.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.2.cmp(&b.2)));
|
|
results
|
|
.into_iter()
|
|
.map(|(value, score, _)| (value, score))
|
|
.collect()
|
|
}
|
|
const EMBEDDING_DIM: usize = 384;
|
|
const MAX_SEQ_LENGTH: usize = 256;
|
|
|
|
const MODEL_URL: &str =
|
|
"https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model.onnx";
|
|
const TOKENIZER_URL: &str =
|
|
"https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json";
|
|
|
|
pub type EmbeddingVec = Vec<f32>;
|
|
|
|
pub struct Embedder {
|
|
model: RunnableEmbeddingModel,
|
|
tokenizer: Tokenizer,
|
|
/// Per-input binding plan: (role, dtype) in the model's DECLARED input order.
|
|
/// Exporters differ in both input ORDER (MiniLM puts input_ids first; e5/bge
|
|
/// put attention_mask first) and DTYPE (f32 vs i64), so we bind by name and
|
|
/// feed each input its model-declared dtype instead of assuming a position.
|
|
input_plan: Vec<(InputRole, DatumType)>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
enum InputRole {
|
|
InputIds,
|
|
AttentionMask,
|
|
TokenTypeIds,
|
|
}
|
|
|
|
fn classify_input(name: &str) -> InputRole {
|
|
let n = name.to_ascii_lowercase();
|
|
if n.contains("attention") || n.contains("mask") {
|
|
InputRole::AttentionMask
|
|
} else if n.contains("token_type") || n.contains("segment") {
|
|
InputRole::TokenTypeIds
|
|
} else {
|
|
InputRole::InputIds
|
|
}
|
|
}
|
|
|
|
impl Embedder {
|
|
pub fn load_from_dir(model_dir: &Path) -> Result<Self> {
|
|
let model_path = model_dir.join("model.onnx");
|
|
let tokenizer_path = model_dir.join("tokenizer.json");
|
|
|
|
if !model_path.exists() || !tokenizer_path.exists() {
|
|
download_model(model_dir)?;
|
|
}
|
|
|
|
let tokenizer = Tokenizer::from_file(&tokenizer_path)
|
|
.map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
|
|
|
|
let raw = tract_onnx::onnx()
|
|
.model_for_path(&model_path)
|
|
.context("Failed to load ONNX model")?;
|
|
|
|
// Determine each input's role (by name) and dtype (declared, else i64).
|
|
let input_outlets = raw
|
|
.input_outlets()
|
|
.context("Failed to read model input outlets")?
|
|
.to_vec();
|
|
let mut input_plan: Vec<(InputRole, DatumType)> = Vec::with_capacity(input_outlets.len());
|
|
for (ix, outlet) in input_outlets.iter().enumerate() {
|
|
let role = classify_input(&raw.node(outlet.node).name);
|
|
let dt = raw
|
|
.input_fact(ix)
|
|
.ok()
|
|
.and_then(|f| f.datum_type.concretize())
|
|
.unwrap_or(DatumType::I64);
|
|
input_plan.push((role, dt));
|
|
}
|
|
|
|
let mut model = raw;
|
|
for (ix, (_, dt)) in input_plan.iter().enumerate() {
|
|
model = model.with_input_fact(ix, InferenceFact::dt_shape(*dt, [1, MAX_SEQ_LENGTH]))?;
|
|
}
|
|
let model = model
|
|
.into_optimized()
|
|
.context("Failed to optimize model")?
|
|
.into_runnable()
|
|
.context("Failed to make model runnable")?;
|
|
|
|
Ok(Self {
|
|
model,
|
|
tokenizer,
|
|
input_plan,
|
|
})
|
|
}
|
|
|
|
pub fn embed(&self, text: &str) -> Result<EmbeddingVec> {
|
|
let encoding = self
|
|
.tokenizer
|
|
.encode(text, true)
|
|
.map_err(|e| anyhow::anyhow!("Tokenization failed: {}", e))?;
|
|
|
|
let mut input_ids = vec![0i64; MAX_SEQ_LENGTH];
|
|
let mut attention_mask = vec![0i64; MAX_SEQ_LENGTH];
|
|
let token_type_ids = vec![0i64; MAX_SEQ_LENGTH];
|
|
|
|
let ids = encoding.get_ids();
|
|
let len = ids.len().min(MAX_SEQ_LENGTH);
|
|
|
|
for i in 0..len {
|
|
input_ids[i] = ids[i] as i64;
|
|
attention_mask[i] = 1;
|
|
}
|
|
|
|
// Build each input tensor by role, cast to the model's declared dtype.
|
|
let make = |data: &[i64], dt: DatumType| -> Result<Tensor> {
|
|
let t: Tensor =
|
|
tract_ndarray::Array2::from_shape_vec((1, MAX_SEQ_LENGTH), data.to_vec())?.into();
|
|
Ok(t.cast_to_dt(dt)?.into_owned())
|
|
};
|
|
|
|
let mut inputs: TVec<TValue> = tvec![];
|
|
for (role, dt) in &self.input_plan {
|
|
let data: &[i64] = match role {
|
|
InputRole::InputIds => &input_ids,
|
|
InputRole::AttentionMask => &attention_mask,
|
|
InputRole::TokenTypeIds => &token_type_ids,
|
|
};
|
|
inputs.push(make(data, *dt)?.into());
|
|
}
|
|
|
|
let outputs = self.model.run(inputs)?;
|
|
|
|
let output = outputs[0].to_array_view::<f32>()?.to_owned();
|
|
|
|
let shape = output.shape();
|
|
if shape.len() == 3 {
|
|
let seq_len = shape[1];
|
|
let hidden_dim = shape[2];
|
|
let mut embedding = vec![0f32; hidden_dim];
|
|
|
|
let valid_tokens = len.min(seq_len);
|
|
|
|
for i in 0..valid_tokens {
|
|
for j in 0..hidden_dim {
|
|
embedding[j] += output[[0, i, j]];
|
|
}
|
|
}
|
|
|
|
for val in &mut embedding {
|
|
*val /= valid_tokens.max(1) as f32;
|
|
}
|
|
|
|
let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
if norm > 0.0 {
|
|
for val in &mut embedding {
|
|
*val /= norm;
|
|
}
|
|
}
|
|
|
|
Ok(embedding)
|
|
} else {
|
|
anyhow::bail!("Unexpected output shape: {:?}", shape);
|
|
}
|
|
}
|
|
|
|
pub fn embed_batch(&self, texts: &[&str]) -> Result<Vec<EmbeddingVec>> {
|
|
texts.iter().map(|t| self.embed(t)).collect()
|
|
}
|
|
}
|
|
|
|
/// Build the run() inputs for a BERT-style model from token ids, honoring the
|
|
/// model's declared input order/role/dtype. Shared by Embedder and CrossEncoder.
|
|
fn build_bert_inputs(
|
|
input_plan: &[(InputRole, DatumType)],
|
|
ids: &[u32],
|
|
type_ids: &[u32],
|
|
) -> Result<TVec<TValue>> {
|
|
let mut input_ids = vec![0i64; MAX_SEQ_LENGTH];
|
|
let mut attention_mask = vec![0i64; MAX_SEQ_LENGTH];
|
|
let mut token_type_ids = vec![0i64; MAX_SEQ_LENGTH];
|
|
let len = ids.len().min(MAX_SEQ_LENGTH);
|
|
for i in 0..len {
|
|
input_ids[i] = ids[i] as i64;
|
|
attention_mask[i] = 1;
|
|
if i < type_ids.len() {
|
|
token_type_ids[i] = type_ids[i] as i64;
|
|
}
|
|
}
|
|
let make = |data: &[i64], dt: DatumType| -> Result<Tensor> {
|
|
let t: Tensor =
|
|
tract_ndarray::Array2::from_shape_vec((1, MAX_SEQ_LENGTH), data.to_vec())?.into();
|
|
Ok(t.cast_to_dt(dt)?.into_owned())
|
|
};
|
|
let mut inputs: TVec<TValue> = tvec![];
|
|
for (role, dt) in input_plan {
|
|
let data: &[i64] = match role {
|
|
InputRole::InputIds => &input_ids,
|
|
InputRole::AttentionMask => &attention_mask,
|
|
InputRole::TokenTypeIds => &token_type_ids,
|
|
};
|
|
inputs.push(make(data, *dt)?.into());
|
|
}
|
|
Ok(inputs)
|
|
}
|
|
|
|
/// A cross-encoder reranker: scores a (query, passage) pair jointly and returns
|
|
/// a single relevance logit. Used to reorder a candidate set after first-stage
|
|
/// retrieval (recall-5). Higher score = more relevant.
|
|
pub struct CrossEncoder {
|
|
model: RunnableEmbeddingModel,
|
|
tokenizer: Tokenizer,
|
|
input_plan: Vec<(InputRole, DatumType)>,
|
|
}
|
|
|
|
impl CrossEncoder {
|
|
pub fn load_from_dir(model_dir: &Path) -> Result<Self> {
|
|
let model_path = model_dir.join("model.onnx");
|
|
let tokenizer_path = model_dir.join("tokenizer.json");
|
|
let tokenizer = Tokenizer::from_file(&tokenizer_path)
|
|
.map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
|
|
|
|
let raw = tract_onnx::onnx()
|
|
.model_for_path(&model_path)
|
|
.context("Failed to load cross-encoder ONNX model")?;
|
|
let input_outlets = raw.input_outlets().context("read input outlets")?.to_vec();
|
|
let mut input_plan: Vec<(InputRole, DatumType)> = Vec::with_capacity(input_outlets.len());
|
|
for (ix, outlet) in input_outlets.iter().enumerate() {
|
|
let role = classify_input(&raw.node(outlet.node).name);
|
|
let dt = raw
|
|
.input_fact(ix)
|
|
.ok()
|
|
.and_then(|f| f.datum_type.concretize())
|
|
.unwrap_or(DatumType::I64);
|
|
input_plan.push((role, dt));
|
|
}
|
|
let mut model = raw;
|
|
for (ix, (_, dt)) in input_plan.iter().enumerate() {
|
|
model = model.with_input_fact(ix, InferenceFact::dt_shape(*dt, [1, MAX_SEQ_LENGTH]))?;
|
|
}
|
|
let model = model
|
|
.into_optimized()
|
|
.context("optimize cross-encoder")?
|
|
.into_runnable()
|
|
.context("make cross-encoder runnable")?;
|
|
Ok(Self {
|
|
model,
|
|
tokenizer,
|
|
input_plan,
|
|
})
|
|
}
|
|
|
|
/// Relevance score for a (query, passage) pair. Higher = more relevant.
|
|
pub fn score(&self, query: &str, passage: &str) -> Result<f32> {
|
|
let encoding = self
|
|
.tokenizer
|
|
.encode((query, passage), true)
|
|
.map_err(|e| anyhow::anyhow!("Tokenization failed: {}", e))?;
|
|
let inputs = build_bert_inputs(
|
|
&self.input_plan,
|
|
encoding.get_ids(),
|
|
encoding.get_type_ids(),
|
|
)?;
|
|
let outputs = self.model.run(inputs)?;
|
|
let view = outputs[0].to_array_view::<f32>()?;
|
|
// logits shape is [1, 1] (relevance) or [1, N]; take the first/primary.
|
|
view.iter()
|
|
.next()
|
|
.copied()
|
|
.ok_or_else(|| anyhow::anyhow!("empty cross-encoder output"))
|
|
}
|
|
|
|
/// Rerank `(id, text)` candidates by cross-encoder score against `query`.
|
|
/// Returns ids sorted by descending relevance.
|
|
pub fn rerank(
|
|
&self,
|
|
query: &str,
|
|
candidates: &[(String, String)],
|
|
) -> Result<Vec<(String, f32)>> {
|
|
let mut scored: Vec<(String, f32)> = Vec::with_capacity(candidates.len());
|
|
for (id, text) in candidates {
|
|
let s = self.score(query, text)?;
|
|
scored.push((id.clone(), s));
|
|
}
|
|
scored.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
Ok(scored)
|
|
}
|
|
}
|
|
|
|
pub const fn embedding_dim() -> usize {
|
|
EMBEDDING_DIM
|
|
}
|
|
|
|
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
|
if a.len() != b.len() || a.is_empty() {
|
|
return 0.0;
|
|
}
|
|
|
|
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
|
|
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
|
|
if norm_a == 0.0 || norm_b == 0.0 {
|
|
return 0.0;
|
|
}
|
|
|
|
dot / (norm_a * norm_b)
|
|
}
|
|
|
|
pub fn batch_cosine_similarity(query: &[f32], candidates: &[&[f32]]) -> Vec<f32> {
|
|
let dim = query.len();
|
|
if dim == 0 || candidates.is_empty() {
|
|
return vec![0.0; candidates.len()];
|
|
}
|
|
|
|
candidates
|
|
.iter()
|
|
.map(|c| {
|
|
if c.len() != dim {
|
|
0.0
|
|
} else {
|
|
c.iter().zip(query.iter()).map(|(a, b)| a * b).sum()
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
pub fn find_similar(
|
|
query: &[f32],
|
|
candidates: &[EmbeddingVec],
|
|
threshold: f32,
|
|
top_k: usize,
|
|
) -> Vec<(usize, f32)> {
|
|
let refs: Vec<&[f32]> = candidates.iter().map(|v| v.as_slice()).collect();
|
|
let scores = batch_cosine_similarity(query, &refs);
|
|
|
|
top_k_scored(
|
|
scores
|
|
.into_iter()
|
|
.enumerate()
|
|
.filter(|(_, score)| *score >= threshold),
|
|
top_k,
|
|
)
|
|
}
|
|
|
|
pub fn is_model_available(model_dir: &Path) -> bool {
|
|
model_dir.join("model.onnx").exists() && model_dir.join("tokenizer.json").exists()
|
|
}
|
|
|
|
fn download_model(model_dir: &Path) -> Result<()> {
|
|
let model_dir = model_dir.to_path_buf();
|
|
match std::thread::spawn(move || download_model_blocking(&model_dir)).join() {
|
|
Ok(result) => result,
|
|
Err(panic) => {
|
|
let panic_msg = if let Some(msg) = panic.downcast_ref::<&str>() {
|
|
(*msg).to_string()
|
|
} else if let Some(msg) = panic.downcast_ref::<String>() {
|
|
msg.clone()
|
|
} else {
|
|
"unknown panic payload".to_string()
|
|
};
|
|
anyhow::bail!("Embedding model download thread panicked: {}", panic_msg);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn download_model_blocking(model_dir: &Path) -> Result<()> {
|
|
let client = reqwest::blocking::Client::builder()
|
|
.user_agent(concat!("jcode-embedding/", env!("CARGO_PKG_VERSION")))
|
|
.timeout(std::time::Duration::from_secs(300))
|
|
.build()?;
|
|
|
|
std::fs::create_dir_all(model_dir)?;
|
|
|
|
let model_path = model_dir.join("model.onnx");
|
|
if !model_path.exists() {
|
|
let response = client.get(MODEL_URL).send()?;
|
|
if !response.status().is_success() {
|
|
anyhow::bail!("Failed to download model: {}", response.status());
|
|
}
|
|
let bytes = response.bytes()?;
|
|
let mut file = std::fs::File::create(&model_path)?;
|
|
file.write_all(&bytes)?;
|
|
}
|
|
|
|
let tokenizer_path = model_dir.join("tokenizer.json");
|
|
if !tokenizer_path.exists() {
|
|
let response = client.get(TOKENIZER_URL).send()?;
|
|
if !response.status().is_success() {
|
|
anyhow::bail!("Failed to download tokenizer: {}", response.status());
|
|
}
|
|
let bytes = response.bytes()?;
|
|
let mut file = std::fs::File::create(&tokenizer_path)?;
|
|
file.write_all(&bytes)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn cosine_similarity_handles_basic_cases() {
|
|
let a = vec![1.0, 0.0, 0.0];
|
|
let b = vec![1.0, 0.0, 0.0];
|
|
let c = vec![0.0, 1.0, 0.0];
|
|
let d = vec![-1.0, 0.0, 0.0];
|
|
|
|
assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);
|
|
assert!((cosine_similarity(&a, &c) - 0.0).abs() < 0.001);
|
|
assert!((cosine_similarity(&a, &d) - (-1.0)).abs() < 0.001);
|
|
}
|
|
|
|
#[test]
|
|
fn find_similar_returns_only_top_k_sorted_hits() {
|
|
let query = vec![1.0, 0.0, 0.0];
|
|
let candidates = vec![
|
|
vec![0.2, 0.0, 0.0],
|
|
vec![0.9, 0.0, 0.0],
|
|
vec![0.7, 0.0, 0.0],
|
|
vec![0.8, 0.0, 0.0],
|
|
];
|
|
|
|
let hits = find_similar(&query, &candidates, 0.1, 2);
|
|
|
|
assert_eq!(hits, vec![(1, 0.9), (3, 0.8)]);
|
|
}
|
|
|
|
fn related_beats_unrelated(model_dir: &Path) {
|
|
let e = Embedder::load_from_dir(model_dir).expect("load model");
|
|
let q = e.embed("how do I set the cargo build profile").unwrap();
|
|
let related = e.embed("The build uses the selfdev cargo profile").unwrap();
|
|
let unrelated = e.embed("Bees pollinate flowers in spring").unwrap();
|
|
let sim_rel = cosine_similarity(&q, &related);
|
|
let sim_unrel = cosine_similarity(&q, &unrelated);
|
|
assert!(
|
|
sim_rel > sim_unrel + 0.05,
|
|
"expected related ({sim_rel:.3}) >> unrelated ({sim_unrel:.3}) for {}",
|
|
model_dir.display()
|
|
);
|
|
}
|
|
|
|
/// Regression test for the input-binding fix: the real MiniLM model must
|
|
/// produce meaningfully higher similarity for related vs unrelated text.
|
|
/// Skipped automatically if the model isn't present locally.
|
|
#[test]
|
|
fn minilm_related_beats_unrelated_if_present() {
|
|
let dir = std::env::var_os("HOME")
|
|
.map(|h| std::path::PathBuf::from(h).join(".jcode/models/all-MiniLM-L6-v2"))
|
|
.filter(|d| is_model_available(d));
|
|
match dir {
|
|
Some(d) => related_beats_unrelated(&d),
|
|
None => eprintln!("skip: MiniLM model not present locally"),
|
|
}
|
|
}
|
|
|
|
/// Exploratory check for an alternate model with a DIFFERENT input
|
|
/// order/dtype (e5-small-v2: attention_mask declared first). Reports the
|
|
/// related vs unrelated gap but does NOT hard-fail: some models need
|
|
/// model-specific pooling/normalization (e.g. CLS pooling) that the shared
|
|
/// mean-pool path does not yet provide. Skipped if not present locally.
|
|
#[test]
|
|
fn alt_model_related_beats_unrelated_if_present() {
|
|
let dir = std::env::var_os("HOME")
|
|
.map(|h| std::path::PathBuf::from(h).join("jcode-memory-bench/models/e5-small-v2"))
|
|
.filter(|d| is_model_available(d));
|
|
let Some(d) = dir else {
|
|
eprintln!("skip: e5-small-v2 model not present locally");
|
|
return;
|
|
};
|
|
let e = Embedder::load_from_dir(&d).expect("load model");
|
|
let q = e.embed("how do I set the cargo build profile").unwrap();
|
|
let related = e.embed("The build uses the selfdev cargo profile").unwrap();
|
|
let unrelated = e.embed("Bees pollinate flowers in spring").unwrap();
|
|
let sim_rel = cosine_similarity(&q, &related);
|
|
let sim_unrel = cosine_similarity(&q, &unrelated);
|
|
eprintln!(
|
|
"e5-small-v2: related={sim_rel:.4} unrelated={sim_unrel:.4} gap={:.4} (informational; mean-pool may need CLS for this family)",
|
|
sim_rel - sim_unrel
|
|
);
|
|
}
|
|
|
|
/// Cross-encoder reranker must score a relevant (query, passage) pair higher
|
|
/// than an irrelevant one. Skipped if the model isn't present locally.
|
|
#[test]
|
|
fn cross_encoder_scores_relevant_higher_if_present() {
|
|
let dir = std::env::var_os("HOME")
|
|
.map(|h| std::path::PathBuf::from(h).join("jcode-memory-bench/models/ce-minilm-l6"))
|
|
.filter(|d| d.join("model.onnx").exists() && d.join("tokenizer.json").exists());
|
|
let Some(d) = dir else {
|
|
eprintln!("skip: cross-encoder model not present locally");
|
|
return;
|
|
};
|
|
let ce = CrossEncoder::load_from_dir(&d).expect("load cross-encoder");
|
|
let q = "how do I set the cargo build profile";
|
|
let rel = ce
|
|
.score(q, "The build uses the selfdev cargo profile")
|
|
.unwrap();
|
|
let unrel = ce.score(q, "Bees pollinate flowers in spring").unwrap();
|
|
eprintln!("cross-encoder: relevant={rel:.3} irrelevant={unrel:.3}");
|
|
assert!(
|
|
rel > unrel,
|
|
"cross-encoder must score relevant ({rel:.3}) > irrelevant ({unrel:.3})"
|
|
);
|
|
}
|
|
}
|