akuna-embed

Simple pure-rust text embedding models built on Burn.

No external model runtimes, and native hardware acceleration.
Super simple interface, give text and get embeddings.
Optionally, specify a model, and/or a Burn backend for different hardware execution.

Usage

use akuna_embed::TextEmbedding;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model = TextEmbedding::new(Default::default()).await?;

    let embedding = model.embed("Hello world")?;
    println!("Embedding has {} numbers", embedding.len());

    Ok(())
}

Embed Many Texts

use akuna_embed::TextEmbedding;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model = TextEmbedding::new(Default::default()).await?;

    let embeddings = model.embed_batch(&[
        "Hello world",
        "Rust embeddings",
        "Semantic search",
    ], None)?;

    println!("Created {} embeddings", embeddings.len());

    Ok(())
}

Search Usage

When building search, embed stored content with embed or embed_batch.

Embed user search text with embed_query or embed_query_batch. Some models, like BGE, use a special query prompt. This crate applies that prompt automatically.

use akuna_embed::TextEmbedding;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model = TextEmbedding::new(Default::default()).await?;

    let document = model.embed("Burn is a deep learning framework for Rust")?;
    let query = model.embed_query("Rust machine learning")?;

    assert_eq!(document.len(), query.len());

    Ok(())
}

Choose A Model

EmbeddingModel::MiniLmL12 is the default.

Available models:

EmbeddingModel::MiniLmL12
EmbeddingModel::MiniLmL6
EmbeddingModel::BgeSmallEnV15
EmbeddingModel::BgeBaseEnV15

use akuna_embed::{EmbeddingModel, TextEmbedding, TextEmbeddingOptions};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model = TextEmbedding::new(TextEmbeddingOptions {
        model: EmbeddingModel::BgeSmallEnV15,
        ..Default::default()
    })
    .await?;

    let embedding = model.embed("Hello world")?;
    assert!(!embedding.is_empty());

    Ok(())
}

Development

This project uses a Nix development shell.

If you use nix-direnv, it should activate automatically. To enter it manually:

nix develop

Run all checks with:

./scripts/check.sh

Run tests only with:

cargo nextest run