toolhub_recommender/
embed.rs

1use std::path::PathBuf;
2
3use anyhow::Context;
4use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
5
6pub struct Embedder {
7    model: TextEmbedding,
8}
9
10impl Embedder {
11    pub fn new() -> anyhow::Result<Self> {
12        let cache = cache_dir();
13        std::fs::create_dir_all(&cache).ok();
14        let opts = InitOptions::new(EmbeddingModel::BGESmallENV15)
15            .with_cache_dir(cache)
16            .with_show_download_progress(false);
17        let model = TextEmbedding::try_new(opts).context("init bge-small-en-v1.5 model")?;
18        Ok(Self { model })
19    }
20
21    pub fn embed_one(&self, text: &str) -> anyhow::Result<Vec<f32>> {
22        let mut out = self.model.embed(vec![text.to_string()], None)?;
23        out.pop()
24            .ok_or_else(|| anyhow::anyhow!("embed() returned no vectors"))
25    }
26
27    pub fn embed_batch(&self, texts: Vec<String>) -> anyhow::Result<Vec<Vec<f32>>> {
28        self.model.embed(texts, None)
29    }
30}
31
32fn cache_dir() -> PathBuf {
33    std::env::var("XDG_CACHE_HOME")
34        .map(PathBuf::from)
35        .ok()
36        .or_else(|| {
37            std::env::var("HOME")
38                .ok()
39                .map(|h| PathBuf::from(h).join(".cache"))
40        })
41        .unwrap_or_else(|| PathBuf::from("."))
42        .join("toolhub")
43        .join("models")
44}
toolhub_recommender/embed.rs

toolhub_recommender/
embed.rs
