Skip to content

mbullington/arctic_ane

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

arctic_ane

arctic_ane is a Rust library for running Snowflake's Arctic Embed retrieval model with Core ML and the Apple Neural Engine (ANE).

The current target model is snowflake-arctic-embed-s, a compact English embedding model with 384-dimensional output embeddings.

Requirements

  • macOS on Apple Silicon, or an iOS target with Core ML support
  • Rust nightly, because the crate currently uses portable_simd
  • A compiled Core ML model at runtime (.mlmodelc)

The tokenizer is embedded in the crate. The Core ML model artifact is not committed to the repo.

Build the model

Build the Core ML package and compile it before running examples or tests:

./model/build_model.sh

This downloads Snowflake/snowflake-arctic-embed-s and writes the generated model files under model/, including:

  • model/SnowflakeArcticEmbed.mlpackage
  • model/SnowflakeArcticEmbed.mlmodelc

Usage

Add the crate to your Rust project. Until it is published to crates.io, use a path or Git dependency:

[dependencies]
arctic_ane = { path = "../arctic_ane" }

Load the compiled Core ML model and embed text:

use std::path::PathBuf;

use arctic_ane::SnowflakeEmbedModel;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model_path = PathBuf::from("model/SnowflakeArcticEmbed.mlmodelc").canonicalize()?;
    let model = SnowflakeEmbedModel::load(model_path)?;

    let text = "Represent this sentence for searching relevant passages: best battery settings for macbook air";
    let output = unsafe { model.predict_with_text(text)? };

    println!("embedding dimensions: {}", output.embedding.len());
    println!("first values: {:?}", &output.embedding[..5]);

    Ok(())
}

For retrieval queries, use the Snowflake query prefix shown above:

Represent this sentence for searching relevant passages: <query>

Development

Run tests after the model has been built:

cargo test

Run the local benchmark:

cargo run --bin bench --release

Model score

From Snowflake's Arctic Embed documentation:

Model name MTEB Retrieval Score (NDCG @ 10)
snowflake-arctic-embed-s 51.98
bge-small-en-v1.5 51.68
Cohere-embed-english-light-v3.0 51.34
text-embedding-3-small 51.08
e5-small-v2 49.04

About

Rust library for running Snowflake's Arctic Embed retrieval model with Core ML and the Apple Neural Engine (ANE)

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors