seyk_core/engine/
image_scan.rs

1use super::*;
2use crate::cache::{ScanData, ScanVector};
3use seyk_embed::quant::QueryLut;
4
5/// Experimental image scan. Returned scores always use the retained exact f32 vectors.
6#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
7#[non_exhaustive]
8#[serde(rename_all = "snake_case")]
9pub enum ImageScan {
10    #[default]
11    Exact,
12    PhotoAnchorQ4,
13}
14
15impl std::str::FromStr for ImageScan {
16    type Err = Error;
17    fn from_str(value: &str) -> Result<Self> {
18        match value {
19            "exact" => Ok(Self::Exact),
20            "q4" => Ok(Self::PhotoAnchorQ4),
21            _ => Err(Error::Config("unknown image scan: use exact or q4".into())),
22        }
23    }
24}
25
26pub(super) fn search(
27    engine: &SearchEngine,
28    rows: &[ScanVector],
29    queries: &[Embedding],
30    take: usize,
31    min_score: f32,
32) -> Result<Vec<Vec<Scored>>> {
33    if take == 0 {
34        return Ok(vec![Vec::new(); queries.len()]);
35    }
36    let coded = rows
37        .iter()
38        .filter(|r| matches!(r.data, ScanData::PhotoAnchorQ4(_)))
39        .count();
40    let depth = coded.min(take.saturating_mul(4).max(50));
41    let mut shortlists = Vec::with_capacity(queries.len());
42    let mut hashes = HashSet::new();
43    for query in queries {
44        let lut = QueryLut::new(query).map_err(|e| Error::Config(e.to_string()))?;
45        let mut approx = Vec::with_capacity(coded);
46        for (id, row) in rows.iter().enumerate() {
47            if let ScanData::PhotoAnchorQ4(codes) = &row.data {
48                approx.push(Scored {
49                    id,
50                    score: lut.score(codes).map_err(|e| Error::Config(e.to_string()))?,
51                });
52            }
53        }
54        approx.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.id.cmp(&b.id)));
55        approx.truncate(depth);
56        for item in &approx {
57            hashes.insert(rows[item.id].segment.input_hash.clone());
58        }
59        shortlists.push(approx);
60    }
61    let exact = engine.with_cache(|c| c.selected_image_vectors(&engine.profile_id, &hashes))?;
62    queries
63        .iter()
64        .zip(shortlists)
65        .map(|(query, shortlist)| {
66            let mut scores = Vec::new();
67            for (id, row) in rows.iter().enumerate() {
68                if let ScanData::Exact(vector) = &row.data {
69                    scores.push(Scored {
70                        id,
71                        score: crate::embedding::dot(query.as_slice(), vector),
72                    });
73                }
74            }
75            for candidate in shortlist {
76                let vector = exact
77                    .get(&rows[candidate.id].segment.input_hash)
78                    .ok_or_else(|| Error::Config("selected image f32 vector disappeared".into()))?;
79                scores.push(Scored {
80                    id: candidate.id,
81                    score: crate::embedding::dot(query.as_slice(), vector),
82                });
83            }
84            scores.retain(|s| s.score.partial_cmp(&min_score) != Some(std::cmp::Ordering::Less));
85            scores.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.id.cmp(&b.id)));
86            scores.truncate(take);
87            Ok(scores)
88        })
89        .collect()
90}
seyk_core/engine/image_scan.rs

seyk_core/engine/
image_scan.rs
