use tysm::embeddings::EmbeddingsClient;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let client = EmbeddingsClient::from_env("text-embedding-3-small")?;
let documents = vec![
"Artificial intelligence is transforming how we interact with technology.".to_string(),
"Machine learning models can recognize patterns in large datasets.".to_string(),
"Natural language processing helps computers understand human language.".to_string(),
];
println!("Embedding multiple documents:");
let embeddings = client.embed(&documents).await?;
for (document, embedding) in embeddings.iter() {
println!(
"\"{}\" -> Vector with {} dimensions (showing first 5: {:?}...)",
document,
embedding.dimension(),
&embedding.elements[..5.min(embedding.dimension())]
);
}
println!("\nEmbedding a single document:");
let single_document =
"Vector databases store and query high-dimensional vectors efficiently.".to_string();
let single_embedding = client.embed_single(single_document.clone()).await?;
println!(
"Single document: \"{}\" -> Vector with {} dimensions (showing first 5: {:?}...)",
single_document,
single_embedding.dimension(),
&single_embedding.elements[..5.min(single_embedding.dimension())]
);
if !embeddings.is_empty() && !single_embedding.elements.is_empty() {
println!("\nCalculating cosine similarity between embeddings:");
for (i, (_, doc_embedding)) in embeddings.iter().enumerate() {
let similarity = doc_embedding.cosine_similarity(&single_embedding);
println!(
"Similarity between document {} and single document: {:.4}",
i + 1,
similarity
);
}
}
Ok(())
}