51 lines
1.6 KiB
Rust
51 lines
1.6 KiB
Rust
|
use std::time::Duration;
|
||
|
|
|
||
|
|
use stratum_embeddings::{EmbeddingOptions, EmbeddingService, FastEmbedProvider, MemoryCache};
|
||
|
|
use tracing::info;
|
||
|
|
|
||
|
|
#[tokio::main]
|
||
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||
|
|
tracing_subscriber::fmt::init();
|
||
|
|
|
||
|
|
info!("Initializing FastEmbed provider...");
|
||
|
|
let provider = FastEmbedProvider::small()?;
|
||
|
|
|
||
|
|
let cache = MemoryCache::new(1000, Duration::from_secs(300));
|
||
|
|
let service = EmbeddingService::new(provider).with_cache(cache);
|
||
|
|
|
||
|
|
info!("Service ready: {:?}", service.provider_info());
|
||
|
|
|
||
|
|
let options = EmbeddingOptions::default_with_cache();
|
||
|
|
|
||
|
|
info!("Embedding single text...");
|
||
|
|
let text = "Stratum embeddings is a unified embedding library";
|
||
|
|
let embedding = service.embed(text, &options).await?;
|
||
|
|
|
||
|
|
info!("Generated embedding with {} dimensions", embedding.len());
|
||
|
|
|
||
|
|
info!("Embedding same text again (should be cached)...");
|
||
|
|
let embedding2 = service.embed(text, &options).await?;
|
||
|
|
assert_eq!(embedding, embedding2);
|
||
|
|
info!("Cache hit confirmed!");
|
||
|
|
|
||
|
|
info!("Embedding batch of texts...");
|
||
|
|
let texts = vec![
|
||
|
|
"Rust is a systems programming language".to_string(),
|
||
|
|
"Knowledge graphs connect concepts".to_string(),
|
||
|
|
"Vector databases enable semantic search".to_string(),
|
||
|
|
];
|
||
|
|
|
||
|
|
let result = service.embed_batch(texts, &options).await?;
|
||
|
|
|
||
|
|
info!(
|
||
|
|
"Batch complete: {} embeddings generated",
|
||
|
|
result.embeddings.len()
|
||
|
|
);
|
||
|
|
info!("Model: {}, Dimensions: {}", result.model, result.dimensions);
|
||
|
|
info!("Cached count: {}", result.cached_count);
|
||
|
|
|
||
|
|
info!("Cache size: {}", service.cache_size());
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|