boostr 0.1.0

ML framework built on numr - attention, quantization, model architectures
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
71
72
73
74
75
76
77
78

//! Hybrid layer configuration for mixed SSM/attention architectures.

use crate::error::{Error, Result};
use serde::{Deserialize, Serialize};

/// Hybrid layer configuration
///
/// Specifies which layers use SSM vs attention in hybrid architectures.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HybridConfig {
    /// Layer indices that use SSM (Mamba)
    pub ssm_layers: Vec<usize>,

    /// Layer indices that use attention (transformer)
    pub attention_layers: Vec<usize>,
}

impl HybridConfig {
    pub fn validate(&self, num_layers: usize) -> Result<()> {
        let mut assigned = vec![false; num_layers];

        for &layer in &self.ssm_layers {
            if layer >= num_layers {
                return Err(Error::ModelError {
                    reason: format!(
                        "ssm_layers contains invalid layer index {layer} (num_layers = {num_layers})"
                    ),
                });
            }
            if assigned[layer] {
                return Err(Error::ModelError {
                    reason: format!(
                        "Layer {layer} assigned to both ssm_layers and attention_layers"
                    ),
                });
            }
            assigned[layer] = true;
        }

        for &layer in &self.attention_layers {
            if layer >= num_layers {
                return Err(Error::ModelError {
                    reason: format!(
                        "attention_layers contains invalid layer index {layer} (num_layers = {num_layers})"
                    ),
                });
            }
            if assigned[layer] {
                return Err(Error::ModelError {
                    reason: format!(
                        "Layer {layer} assigned to both ssm_layers and attention_layers"
                    ),
                });
            }
            assigned[layer] = true;
        }

        for (i, &is_assigned) in assigned.iter().enumerate() {
            if !is_assigned {
                return Err(Error::ModelError {
                    reason: format!(
                        "Layer {i} is not assigned to either ssm_layers or attention_layers"
                    ),
                });
            }
        }

        Ok(())
    }

    pub fn is_ssm_layer(&self, layer_idx: usize) -> bool {
        self.ssm_layers.contains(&layer_idx)
    }

    pub fn is_attention_layer(&self, layer_idx: usize) -> bool {
        self.attention_layers.contains(&layer_idx)
    }
}