diskann-tools 0.55.0

DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases.
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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */
use std::time::Instant;

use clap::Parser;
use diskann_providers::storage::FileStorageProvider;
use diskann_tools::utils::{
    compute_multivec_ground_truth_from_datafiles, init_subscriber, CMDResult, DataType,
    MultivecAggregationMethod,
};
use diskann_vector::distance::Metric;
use diskann_vector::Half;

fn main() -> CMDResult<()> {
    init_subscriber();
    let timer = Instant::now();

    let args = ComputeMultivecGroundTruthArgs::parse();

    tracing::info!("Computing ground truth file");

    let storage_provider = FileStorageProvider;

    let err = match args.data_type {
        DataType::Float => {
            compute_multivec_ground_truth_from_datafiles::<f32, FileStorageProvider>(
                &storage_provider,
                args.distance_function,
                args.aggregation,
                &args.base_file,
                &args.query_file,
                &args.ground_truth_file,
                args.recall_at,
                args.base_file_labels.as_deref(),
                args.query_file_labels.as_deref(),
            )
        }
        DataType::Fp16 => {
            compute_multivec_ground_truth_from_datafiles::<Half, FileStorageProvider>(
                &storage_provider,
                args.distance_function,
                args.aggregation,
                &args.base_file,
                &args.query_file,
                &args.ground_truth_file,
                args.recall_at,
                args.base_file_labels.as_deref(),
                args.query_file_labels.as_deref(),
            )
        }
        DataType::Uint8 => compute_multivec_ground_truth_from_datafiles::<u8, FileStorageProvider>(
            &storage_provider,
            args.distance_function,
            args.aggregation,
            &args.base_file,
            &args.query_file,
            &args.ground_truth_file,
            args.recall_at,
            args.base_file_labels.as_deref(),
            args.query_file_labels.as_deref(),
        ),
        DataType::Int8 => compute_multivec_ground_truth_from_datafiles::<i8, FileStorageProvider>(
            &storage_provider,
            args.distance_function,
            args.aggregation,
            &args.base_file,
            &args.query_file,
            &args.ground_truth_file,
            args.recall_at,
            args.base_file_labels.as_deref(),
            args.query_file_labels.as_deref(),
        ),
    };

    match err {
        Ok(_) => {
            tracing::info!(
                "Compute ground-truth completed successfully in {:?}",
                timer.elapsed()
            );
            Ok(())
        }
        Err(err) => {
            tracing::error!("Error: {:?}", err);
            Err(err)
        }
    }
}

#[derive(Debug, Parser)]
struct ComputeMultivecGroundTruthArgs {
    /// data type <int8/uint8/float / fp16> (required)
    #[arg(long = "data-type", default_value = "float")]
    pub data_type: DataType,

    /// Distance function to use.
    #[arg(long = "dist-fn", default_value = "l2")]
    pub distance_function: Metric,

    /// Whether to use average or min aggregation
    #[arg(long = "aggregation", default_value = "average")]
    pub aggregation: MultivecAggregationMethod,

    /// File containing the base vectors in binary format
    #[arg(long = "base-file", short, required = true)]
    pub base_file: String,

    #[arg(long = "base-file-labels", default_value = None)]
    pub base_file_labels: Option<String>,

    /// File containing the query vectors in binary format
    #[arg(long = "query-file", short, required = true)]
    pub query_file: String,

    #[arg(long = "query-file-labels")]
    pub query_file_labels: Option<String>,

    /// Path of the file to write the ground truth to in binary format.  Please don't append .bin at the end if no filter_label or filter_label_file is provided.  It will save the file with '.bin' at the end.  Otherwise it will save the file as filename_label.bin.
    #[arg(long = "gt-file", short, required = true)]
    pub ground_truth_file: String,

    /// Number of ground truth nearest neighbors to compute
    #[arg(long = "recall-at", short = 'K', default_value = "10")]
    pub recall_at: u32,
}