libmir-metal 0.1.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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MetalProgressStage {
    LoadWeights,
    PrefillTokens,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MetalProgressUnit {
    Byte,
    Token,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MetalProgressEvent {
    pub stage: MetalProgressStage,
    pub current: u64,
    pub total: u64,
    pub unit: MetalProgressUnit,
    pub detail: String,
}

impl MetalProgressEvent {
    #[must_use]
    pub fn load_weights(current: u64, total: u64, detail: String) -> Self {
        Self {
            stage: MetalProgressStage::LoadWeights,
            current,
            total,
            unit: MetalProgressUnit::Byte,
            detail,
        }
    }

    #[must_use]
    pub fn prefill_tokens(current: usize, total: usize) -> Self {
        Self {
            stage: MetalProgressStage::PrefillTokens,
            current: current as u64,
            total: total as u64,
            unit: MetalProgressUnit::Token,
            detail: format!("token {current}/{total}"),
        }
    }
}