xz-embed 0.1.0

Text embedding and vector storage abstraction layer
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

use std::collections::HashMap;
use xz_embed::*;

#[tokio::test]
async fn test_mock_embedder() {
    let mut mock = MockEmbedder::new(384, 32);
    mock.set_output(vec![vec![0.1; 384], vec![0.2; 384]]);

    let vectors = mock.embed(&["hello", "world"]).await.unwrap();
    assert_eq!(vectors.len(), 2);
    assert_eq!(vectors[0].len(), 384);
    assert_eq!(mock.dimensions(), 384);
}

#[tokio::test]
async fn test_batch_embed_with_expectation() {
    let mut mock = MockEmbedder::new(4, 32);
    mock.expect_embed(
        vec!["hello", "world"],
        vec![vec![0.1; 4], vec![0.2; 4]],
    );

    let vectors = mock.embed(&["hello", "world"]).await.unwrap();
    assert_eq!(vectors.len(), 2);
    assert_eq!(vectors[0].len(), 4);
}

#[tokio::test]
async fn test_embed_error_empty_batch() {
    let mock = MockEmbedder::new(4, 32);
    let result = mock.embed(&[]).await;
    assert!(matches!(result.unwrap_err(), EmbedError::EmptyBatch));
}