#![cfg(feature = "bge")]
use ll_core::embed::EmbeddingProvider;
use ll_core::models::bge_small::{ensure_model, BgeSmallProvider};
use std::path::Path;
fn cosine(a: &[f32], b: &[f32]) -> f32 {
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
let n = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt();
dot / (n(a) * n(b))
}
#[test]
#[ignore = "needs ONNX Runtime and the bge-small model; set LL_MODELS_DIR to reuse a download"]
fn threads_option_matches_default() {
ll_core::dylib::ensure_dylib().unwrap();
let (model, tok) = ensure_model().unwrap();
let default = BgeSmallProvider::from_files(&model, &tok).unwrap();
let one = BgeSmallProvider::from_files_with_threads(&model, &tok, Some(1)).unwrap();
let none = BgeSmallProvider::from_files_with_threads(&model, &tok, None).unwrap();
let texts: Vec<String> = ["the shipper tails transcripts", "omit core seals text", "a"]
.map(String::from)
.to_vec();
let d = default.embed_documents(&texts).unwrap();
let o = one.embed_documents(&texts).unwrap();
let n = none.embed_documents(&texts).unwrap();
for i in 0..texts.len() {
assert!(cosine(&d[i], &n[i]) >= 0.99999);
assert!(cosine(&d[i], &o[i]) >= 0.99999);
}
}
#[test]
fn a_count_past_c_int_is_rejected() {
let res = BgeSmallProvider::from_files_with_threads(
Path::new("/nonexistent/model.onnx"),
Path::new("/nonexistent/tokenizer.json"),
Some(1usize << 32),
);
let err = res
.err()
.expect("a count ORT would truncate must be refused");
assert!(err
.to_string()
.contains("intra_threads must fit in a C int"));
}
#[test]
fn zero_threads_is_rejected() {
let res = BgeSmallProvider::from_files_with_threads(
Path::new("/nonexistent/model.onnx"),
Path::new("/nonexistent/tokenizer.json"),
Some(0),
);
let err = res.err().expect("Some(0) must be refused");
assert!(err.to_string().contains("intra_threads must be at least 1"));
}