use fastembed::{FlagEmbedding, InitOptions, EmbeddingModel, EmbeddingBase};

// With default InitOptions
let model: FlagEmbedding = FlagEmbedding::try_new(Default::default())?;

// With custom InitOptions
let model: FlagEmbedding = FlagEmbedding::try_new(InitOptions {
    model_name: EmbeddingModel::BGEBaseEN,
    show_download_message: true,
    ..Default::default()
})?;

let documents = vec![
    "passage: Hello, World!",
    "query: Hello, World!",
    "passage: This is an example passage.",
    // You can leave out the prefix but it's recommended
    "fastembed-rs is licensed under MIT"
    ];

 // Generate embeddings with the default batch size, 256
 let embeddings = model.embed(documents, None)?;

 println!("Embeddings length: {}", embeddings.len()); // -> Embeddings length: 4
 println!("Embedding dimension: {}", embeddings[0].len()); // -> Embedding dimension: 768

 // Generate embeddings for the passages
 // The texts are prefixed with "passage" for better results
 // The batch size is set to 1 for demonstration purposes
 let passages = vec![
     "This is the first passage. It contains provides more context for retrieval.",
     "Here's the second passage, which is longer than the first one. It includes additional information.",
     "And this is the third passage, the longest of all. It contains several sentences and is meant for more extensive testing."
    ];

 let embeddings = model.passage_embed(passages, Some(1))?;

 println!("Passage embeddings length: {}", embeddings.len()); // -> Embeddings length: 3
 println!("Passage embedding dimension: {}", embeddings[0].len()); // -> Passage embedding dimension: 768

 // Generate embeddings for the query
 // The text is prefixed with "query" for better retrieval
 let query = "What is the answer to this generic question?";

 let query_embedding = model.query_embed(query)?;

 println!("Query embedding dimension: {}", query_embedding.len()); // -> Query embedding dimension: 768