mindmap 0.1.2

Search your notes at the speed of thought
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

use std::path::PathBuf;

use acap::cos::cosine_distance;
use acap::knn::NearestNeighbors;
use acap::vp::VpTree;
use acap::{Distance, Proximity};
use anyhow::Result;

use crate::formatter::{self, OutputFormat};
use crate::{
    config::MindmapConfig,
    database::{self, EmbeddedSentence},
    embeddings::Model,
};

#[derive(Debug)]
pub struct SearchResult {
    pub path: PathBuf,
    pub start_line_no: usize,
    pub end_line_no: usize,
    pub distance: f32,
}

impl Proximity<EmbeddedSentence> for EmbeddedSentence {
    type Distance = f32;

    fn distance(&self, other: &Self) -> Self::Distance {
        cosine_distance(&self.embedding, &other.embedding)
    }
}

impl Proximity<EmbeddedSentence> for Vec<f32> {
    type Distance = f32;

    fn distance(&self, other: &EmbeddedSentence) -> Self::Distance {
        cosine_distance(&self, &other.embedding)
    }
}

pub struct EmbeddingTree<'a> {
    tree: VpTree<EmbeddedSentence>,
    model: Model,
    config: &'a MindmapConfig,
}

impl<'a> EmbeddingTree<'a> {
    pub fn new(corpus: Vec<EmbeddedSentence>, model: Model, config: &'a MindmapConfig) -> Self {
        Self {
            tree: VpTree::balanced(corpus),
            model,
            config,
        }
    }

    pub fn rebuild(&mut self, corpus: Vec<EmbeddedSentence>) {
        self.tree = VpTree::balanced(corpus);
    }

    pub fn search(&self, query: &str) -> Result<Vec<SearchResult>> {
        let emb = self.model.encode(query)?;
        let num_resuls = self.config.num_results;

        let results = self
            .tree
            .k_nearest(&emb, num_resuls)
            .iter()
            .map(|x| SearchResult {
                path: x.item.path.clone(),
                start_line_no: x.item.start_line_no,
                end_line_no: x.item.end_line_no,
                distance: x.distance.value(),
            })
            .collect();

        Ok(results)
    }
}

pub fn search(query: &str, config: &MindmapConfig, format: OutputFormat) -> Result<()> {
    let corpus = database::get_all(config)?;
    let model = Model::new(config).unwrap();

    let tree = EmbeddingTree::new(corpus, model, config);
    let results = tree.search(query)?;

    // Format response
    let formatted = formatter::format(&results, format);
    println!("{}", formatted);
    Ok(())
}