spark-bert 0.1.1

Hybrid vector search using an inverted index and BERT embeddings
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

use tantivy::postings::{Postings, SegmentPostings};
use tantivy::query::{EmptyScorer, EnableScoring, Explanation, Query, Scorer, Weight};
use tantivy::schema::IndexRecordOption;
use tantivy::{DocId, DocSet, Score, SegmentReader, Term};

// TF Scorer
struct TfScorer {
    postings: SegmentPostings,
}

impl DocSet for TfScorer {
    fn advance(&mut self) -> DocId {
        self.postings.advance();
        self.doc()
    }

    fn doc(&self) -> DocId {
        self.postings.doc()
    }

    fn size_hint(&self) -> u32 {
        unimplemented!()
    }
}

impl Scorer for TfScorer {
    fn score(&mut self) -> Score {
        self.postings.term_freq() as Score
    }
}

#[derive(Debug, Clone)]
pub struct TfTermQuery {
    term: Term,
}

impl TfTermQuery {
    pub fn new(term: Term) -> Self {
        Self { term }
    }
}

struct TfWeight {
    term: Term,
}

impl Weight for TfWeight {
    fn scorer(&self, reader: &SegmentReader, boost: Score) -> tantivy::Result<Box<dyn Scorer>> {
        let inv = reader.inverted_index(self.term.field())?;
        if let Some(info) = inv.get_term_info(&self.term)? {
            let postings = inv.read_postings_from_terminfo(&info, IndexRecordOption::WithFreqs)?;
            Ok(Box::new(TfScorer { postings }))
        } else {
            Ok(Box::new(EmptyScorer))
        }
    }

    fn explain(&self, reader: &SegmentReader, doc: DocId) -> tantivy::Result<Explanation> {
        unimplemented!()
    }
}

impl Query for TfTermQuery {
    fn weight(&self, enable_scoring: EnableScoring<'_>) -> tantivy::Result<Box<dyn Weight>> {
        let weight = TfWeight {
            term: self.term.clone(),
        };
        tantivy::Result::Ok(Box::new(weight))
    }
}