libmir-metal 0.3.0

Metal inference backend for libmir
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

use std::collections::HashMap;

use super::PrefixEntry;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(super) struct PrefixKey(pub(super) [u8; 32]);

pub(super) fn indexed_prefixes(
    model: &str,
    tokens: &[u32],
    block_size: Option<usize>,
) -> Vec<(PrefixKey, usize)> {
    let mut hasher = prefix_hasher(model);
    let mut indexed = Vec::new();
    for (index, token) in tokens.iter().enumerate() {
        hasher.update(&token.to_le_bytes());
        let position = index + 1;
        let efficient = position.saturating_mul(2) >= tokens.len();
        if position == tokens.len()
            || efficient && block_size.is_some_and(|size| position % size == 0)
        {
            indexed.push((PrefixKey(*hasher.finalize().as_bytes()), position));
        }
    }
    indexed
}

pub(super) fn longest_indexed_prefix(
    model: &str,
    tokens: &[u32],
    entries: &HashMap<PrefixKey, PrefixEntry>,
) -> Option<(PrefixKey, PrefixEntry)> {
    let mut hasher = prefix_hasher(model);
    let mut longest = None;
    for token in tokens {
        hasher.update(&token.to_le_bytes());
        let key = PrefixKey(*hasher.finalize().as_bytes());
        if let Some(entry) = entries.get(&key) {
            longest = Some((key, *entry));
        }
    }
    longest
}

fn prefix_hasher(model: &str) -> blake3::Hasher {
    let mut hasher = blake3::Hasher::new();
    hasher.update(model.as_bytes());
    hasher.update(&[0]);
    hasher
}