Struct CachedEmbeddingModel

pub struct CachedEmbeddingModel<M: Embedder, S = DefaultHasher> { /* private fields */ }

Embedding models can be expensive to run. This struct wraps an embedding model with a cache that stores embeddings that have been computed before.

Example

use kalosm::language::*;
use std::num::NonZeroUsize;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let bert = Bert::builder()
        .build()
        .await?
        // You can call the `.cached` method to cache the results of the Bert embedding in a LRU cache with the given capacity.
        .cached(NonZeroUsize::new(1000).unwrap());

    // Try to load the cache from the filesystem
    if let Ok(cache) = std::fs::read("cache.bin") {
        let cache: Vec<(EmbeddingInput, Vec<f32>)> = postcard::from_bytes(&cache)?;
        bert.load_cache(cache);
    }

    let start_time = std::time::Instant::now();
    let sentences = [
        "Cats are cool",
        "The geopolitical situation is dire",
        "Pets are great",
        "Napoleon was a tyrant",
        "Napoleon was a great general",
    ];
    // When you embed a new sentence, the cache will store the embedding for that sentence.
    let embeddings = bert.embed_batch(sentences).await?;
    println!("{:?}", embeddings);
    println!("embedding uncached took {:?}", start_time.elapsed());

    let start_time = std::time::Instant::now();
    // If you embed the same sentences again, the cache will be used.
    let embeddings = bert.embed_batch(sentences).await?;
    println!("{:?}", embeddings);
    println!("embedding cached took {:?}", start_time.elapsed());

    let sentences = [
        "Cats are cool",
        "The geopolitical situation is dire",
        "Pets are great",
        "Napoleon was a king",
        "Napoleon was kind of not a great guy",
    ];
    // When you embed a new sentence, the cache will store the embedding for that sentence.
    let embeddings = bert.embed_batch(sentences).await?;
    println!("{:?}", embeddings);
    println!("embedding partially cached took {:?}", start_time.elapsed());

    // Save the cache to the filesystem for future use
    let cache = bert.export_cache();
    let file = std::fs::File::create("cache.bin")?;
    let mut writer = std::io::BufWriter::new(file);
    postcard::to_io(&cache, &mut writer)?;

    Ok(())
}

Implementations

impl<M: Embedder> CachedEmbeddingModel<M>

pub fn new(model: M, cache_size: NonZeroUsize) -> Self

Create a new cached embedding model.

impl<M: Embedder, S> CachedEmbeddingModel<M, S>

pub fn get_embedder(&self) -> &M

Get a reference to the underlying embedder.

pub fn get_embedder_mut(&mut self) -> &mut M

Get a mutable reference to the underlying embedder.

impl<M: Embedder, S: BuildHasher> CachedEmbeddingModel<M, S>

pub fn new_with_hasher(model: M, cache_size: NonZeroUsize, hasher: S) -> Self

Create a new cached embedding model with a custom hasher.

pub fn export_cache(&self) -> Vec<(EmbeddingInput, Box<[f32]>)>

Return a serializable cache of the embeddings for future use. You can load the cache from the file with Self::load_cache.

Example

let bert = Bert::builder()
    .build()
    .await?
    // You can call the `.cached` method to cache the results of the Bert embedding in a LRU cache with the given capacity.
    .cached(NonZeroUsize::new(1000).unwrap());
let sentences = [
    "Cats are cool",
    "The geopolitical situation is dire",
    "Pets are great",
    "Napoleon was a tyrant",
    "Napoleon was a great general",
];
// When you embed a new sentence, the cache will store the embedding for that sentence.
let embeddings = bert.embed_batch(sentences).await?;
println!("{:?}", embeddings);
// Save the cache to the filesystem for future use
let cache = bert.export_cache();
let file = std::fs::File::create("cache.bin")?;
let mut writer = std::io::BufWriter::new(file);
postcard::to_io(&cache, &mut writer)?;

pub fn load_cache(&self, cached_items: Vec<(EmbeddingInput, Vec<f32>)>)

Load the cache from a file.

Example

let bert = Bert::builder()
    .build()
    .await?
    // You can call the `.cached` method to cache the results of the Bert embedding in a LRU cache with the given capacity.
    .cached(NonZeroUsize::new(1000).unwrap());

// Try to load the cache from the filesystem
let cache = std::fs::read("cache.bin")?;
let cache: Vec<(EmbeddingInput, Vec<f32>)> = postcard::from_bytes(&cache)?;
let _ = bert.load_cache(cache);

let sentences = [
    "Cats are cool",
    "The geopolitical situation is dire",
    "Pets are great",
    "Napoleon was a tyrant",
    "Napoleon was a great general",
];
// If the sentence is already in the cache, it will be returned from the cache instead of running the model
let embeddings = bert.embed_batch(sentences).await?;
println!("{:?}", embeddings);

Trait Implementations

impl<M: Embedder> Embedder for CachedEmbeddingModel<M>

where M::Error: Error,

type Error = <M as Embedder>::Error

The error type that can occur when embedding a string.

fn embed_for( &self, input: EmbeddingInput, ) -> impl Future<Output = Result<Embedding, Self::Error>> + Send

Embed a single string.

fn embed_vec_for( &self, inputs: Vec<EmbeddingInput>, ) -> impl Future<Output = Result<Vec<Embedding>, Self::Error>> + Send

Embed a batch of strings.

fn embed_string( &self, input: String, ) -> impl Future<Output = Result<Embedding, Self::Error>> + Send

Embed some text into a vector space.

fn embed_vec( &self, inputs: Vec<String>, ) -> impl Future<Output = Result<Vec<Embedding>, Self::Error>> + Send

Embed a Vec<String> into a vector space. Returns a list of embeddings in the same order as the inputs.