Fastembed

cargo add anyhow fastembed serde tokio surrealdb --features surrealdb/kv-mem

use anyhow::Error;
use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
use serde::Serialize;
use surrealdb::{
    Surreal, Value,
    engine::any::{Any, connect},
};

const SCHEMA: &str = "DEFINE TABLE document;
        DEFINE FIELD text ON document TYPE string;
        DEFINE FIELD embedding ON document TYPE array<float>;
        // Uncomment this to use HNSW index, ensure that number after DIMENSION matches size of embedding
        // DEFINE INDEX hnsw_embed ON document FIELDS embedding HNSW DIMENSION 384 DIST COSINE";

const INSERT_QUERY: &str = "INSERT INTO document $docs";

const VECTOR_QUERY: &str = "SELECT 
        text, 
        vector::distance::knn() AS distance 
        FROM document
        WHERE embedding <|3,COSINE|> $embeds
        ORDER BY distance";

#[derive(Serialize)]
struct DocumentInput {
    text: String,
    embedding: Vec<f32>,
}

async fn store_docs(
    input: Vec<&str>,
    db: &Surreal<Any>,
    model: &mut TextEmbedding,
) -> Result<(), Error> {
    let docs = model
        .embed(input.clone(), None)?
        .into_iter()
        .zip(input.into_iter())
        .map(|(embedding, text)| DocumentInput {
            text: text.to_string(),
            embedding,
        })
        .collect::<Vec<DocumentInput>>();

    db.query(INSERT_QUERY).bind(("docs", docs)).await?;
    Ok(())
}

async fn test_embed(
    input: &str,
    db: &Surreal<Any>,
    model: &mut TextEmbedding,
) -> Result<(), Error> {
    let Some(embeds) = model.embed(vec![input], None)?.into_iter().next() else {
        return Err(anyhow::anyhow!("Nothing found at index 0"));
    };

    let val = db
        .query(VECTOR_QUERY)
        .bind(("embeds", embeds.clone()))
        .await?
        .take::<Value>(0)?;
    println!("{val}\n");
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    // Default model
    let mut model = TextEmbedding::try_new(InitOptions::new(EmbeddingModel::BGESmallENV15))?;

    let db = connect("memory").await?;

    db.use_ns("ns").use_db("db").await?;

    db.query(SCHEMA).await?;

    let input = vec![
        // Cities
        "Calgary is a city in the Canadian province of Alberta.",
        "Ljubljana is the capital and largest city of Slovenia.",
        // Historical / mythological figures
        "Xenophon of Athens was a Greek military leader, philosopher, and historian.",
        "King Arthur was a mythical king in the mythology of Great Britain.",
        // Planets
        "Venus is the second planet from the Sun.",
        "Ceres is a dwarf planet in the middle main asteroid belt between the orbits of Mars and Jupiter.",
        // Languages
        "Manx is a Gaelic language of the insular Celtic branch of the Celtic language family",
        "Interlingue, originally Occidental, is an international auxiliary language created in 1922.",
        // Sea animals
        "Octopuses have a complex nervous system and are among the most intelligent and behaviourally diverse invertebrates.",
        "Clams have no central nervous system at all and are near to plants in intelligence.",
    ];
    store_docs(input, &db, &mut model).await?;

    println!("Edmonton is closest to:");
    test_embed("Edmonton", &db, &mut model).await?;

    println!("Merlin is closest to:");
    test_embed("Merlin", &db, &mut model).await?;

    println!("Earth is closest to:");
    test_embed("Earth", &db, &mut model).await?;

    println!("Irish is closest to:");
    test_embed("Irish language", &db, &mut model).await?;

    println!("Squid are closest to:");
    test_embed("Squid", &db, &mut model).await?;

    Ok(())
}

Edmonton is closest to:
[
    { distance: 0.2596421358215669f, text: 'Calgary is a city in the Canadian province of Alberta.' },
    { distance: 0.5010449624435647f, text: 'Ljubljana is the capital and largest city of Slovenia.' },
    { distance: 0.5242241576926254f, text: 'Interlingue, originally Occidental, is an international auxiliary language created in 1922.' }
]

Merlin is closest to:
[
    { distance: 0.3653307924860497f, text: 'King Arthur was a mythical king in the mythology of Great Britain.' },
    { distance: 0.4515194174120666f, text: 'Manx is a Gaelic language of the insular Celtic branch of the Celtic language family' },
    { distance: 0.5317039966149415f, text: 'Calgary is a city in the Canadian province of Alberta.' }
]

Earth is closest to:
[
    { distance: 0.3380429615054925f, text: 'Venus is the second planet from the Sun.' },
    { distance: 0.3764237673020161f, text: 'Ceres is a dwarf planet in the middle main asteroid belt between the orbits of Mars and Jupiter.' },
    { distance: 0.444087039462282f, text: 'Calgary is a city in the Canadian province of Alberta.' }
]

Irish is closest to:
[
    { distance: 0.27517683002655635f, text: 'Manx is a Gaelic language of the insular Celtic branch of the Celtic language family' },
    { distance: 0.34080671701374754f, text: 'Interlingue, originally Occidental, is an international auxiliary language created in 1922.' },
    { distance: 0.5113325799682362f, text: 'King Arthur was a mythical king in the mythology of Great Britain.' }
]

Squid are closest to:
[
    { distance: 0.3439891425642231f, text: 'Octopuses have a complex nervous system and are among the most intelligent and behaviourally diverse invertebrates.' },
    { distance: 0.4707156750207915f, text: 'Manx is a Gaelic language of the insular Celtic branch of the Celtic language family' },
    { distance: 0.517311424260043f, text: 'Clams have no central nervous system at all and are near to plants in intelligence.' }
]