brain2qwerty 0.0.1

Brain2Qwerty V1/V2 MEG neural decoding inference in Rust (parity-tested vs Python)
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! Safetensors loading and PyTorch key utilities.

use std::collections::HashMap;

use half::{bf16, f16};
use safetensors::SafeTensors;

use crate::tensor::Tensor;

#[derive(Clone, Debug)]
pub struct ParamBuf {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
}

pub type WeightStore = HashMap<String, ParamBuf>;

pub fn load_safetensors(path: &str) -> anyhow::Result<WeightStore> {
    let bytes = std::fs::read(path)?;
    let st = SafeTensors::deserialize(&bytes)?;
    let mut out = HashMap::with_capacity(st.len());
    for (raw_key, view) in st.tensors() {
        let key = raw_key
            .strip_prefix("model.")
            .or_else(|| raw_key.strip_prefix("network."))
            .unwrap_or(raw_key.as_str())
            .to_string();
        let shape: Vec<usize> = view.shape().to_vec();
        let raw = view.data();
        let data = match view.dtype() {
            safetensors::Dtype::BF16 => raw
                .chunks_exact(2)
                .map(|b| bf16::from_le_bytes([b[0], b[1]]).to_f32())
                .collect(),
            safetensors::Dtype::F16 => raw
                .chunks_exact(2)
                .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32())
                .collect(),
            safetensors::Dtype::F32 => raw
                .chunks_exact(4)
                .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
                .collect(),
            safetensors::Dtype::I64 | safetensors::Dtype::I32 | safetensors::Dtype::U8 => {
                continue;
            }
            other => anyhow::bail!("unsupported dtype {other:?} for key {key}"),
        };
        out.insert(key, ParamBuf { data, shape });
    }
    Ok(out)
}

pub fn param_to_tensor(p: &ParamBuf) -> Tensor {
    Tensor::from_vec(p.data.clone(), p.shape.clone())
}

pub fn take(store: &mut WeightStore, key: &str) -> anyhow::Result<ParamBuf> {
    store
        .remove(key)
        .ok_or_else(|| anyhow::anyhow!("missing weight key: {key}"))
}

pub fn get<'a>(store: &'a WeightStore, key: &str) -> anyhow::Result<&'a ParamBuf> {
    store
        .get(key)
        .ok_or_else(|| anyhow::anyhow!("missing weight key: {key}"))
}

/// PyTorch Linear weight (out, in) → row-major matmul (in, out) for our linear().
pub fn linear_weight_transposed(p: &ParamBuf) -> ParamBuf {
    assert_eq!(p.shape.len(), 2);
    let (out_d, in_d) = (p.shape[0], p.shape[1]);
    let mut data = vec![0.0f32; p.data.len()];
    for o in 0..out_d {
        for i in 0..in_d {
            data[i * out_d + o] = p.data[o * in_d + i];
        }
    }
    ParamBuf {
        data,
        shape: vec![in_d, out_d],
    }
}

/// PyTorch Linear weight (out, in) kept as (out, in) for mm on right side.
pub fn linear_weight_for_mm(p: &ParamBuf) -> ParamBuf {
    let (out_d, in_d) = (p.shape[0], p.shape[1]);
    let mut data = vec![0.0f32; p.data.len()];
    for o in 0..out_d {
        for i in 0..in_d {
            data[o * in_d + i] = p.data[o * in_d + i];
        }
    }
    ParamBuf {
        data,
        shape: vec![out_d, in_d],
    }
}

pub fn compute_output_lens(input_len: usize, kernel: usize, stride: usize) -> usize {
    if input_len < kernel {
        0
    } else {
        (input_len - kernel) / stride + 1
    }
}