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

//! Full V1 model: keystroke encoder + sentence transformer + character linear head.

use crate::tensor::Tensor;
use crate::v1::config::V1Config;
use crate::v1::encoder::SimpleConvTimeAgg;
use crate::v1::transformer::SentenceTransformer;
use crate::weights::{get, load_safetensors, param_to_tensor, WeightStore};

/// Brain2Qwerty V1 — conv encoder, ALiBi sentence transformer, and 29-class head.
pub struct Brain2QwertyV1 {
    pub config: V1Config,
    pub encoder: SimpleConvTimeAgg,
    pub transformer: SentenceTransformer,
    pub linear_w: Tensor,
    pub linear_b: Tensor,
}

impl Brain2QwertyV1 {
    pub fn from_config_and_weights(cfg: V1Config, weights_path: &str) -> anyhow::Result<Self> {
        let mut store = load_safetensors(weights_path)?;
        Self::from_config_and_store(cfg, &mut store)
    }

    pub fn from_tiny_weights(weights_path: &str) -> anyhow::Result<Self> {
        Self::from_config_and_weights(V1Config::tiny(), weights_path)
    }

    pub fn from_config_and_store(cfg: V1Config, store: &mut WeightStore) -> anyhow::Result<Self> {
        let hidden = cfg.brain_model_config.encoder_config.hidden;
        let encoder =
            SimpleConvTimeAgg::from_config_and_weights(&cfg.brain_model_config, store, "")?;
        let transformer = SentenceTransformer::from_config_and_weights(
            &cfg.transformer_config,
            hidden,
            store,
            "transformer.",
        )?;
        let linear_w = param_to_tensor(get(store, "linear.weight")?);
        let linear_b = param_to_tensor(get(store, "linear.bias")?);
        Ok(Self {
            config: cfg,
            encoder,
            transformer,
            linear_w,
            linear_b,
        })
    }

    /// Encode a batch of keystroke windows `(B, C, T)` → `(B, hidden)`.
    pub fn encode_keystrokes(
        &self,
        neuros: &Tensor,
        subject_ids: &[usize],
        chan_pos: Option<&Tensor>,
    ) -> Tensor {
        self.encoder.forward(neuros, subject_ids, chan_pos)
    }

    /// Sentence transformer on `(1, N, hidden)` with boolean mask length N.
    pub fn refine_sentence(&self, embeds: &Tensor, mask: &[bool]) -> Tensor {
        self.transformer.forward(embeds, mask)
    }

    /// Linear head on `(N, hidden)` → `(N, num_classes)`.
    pub fn classify(&self, hidden: &Tensor) -> Tensor {
        let n = hidden.shape[0];
        let d = hidden.shape[1];
        let c = self.linear_w.shape[0];
        let btd = hidden.reshape(&[1, n, d]);
        btd.linear(&self.linear_w, Some(&self.linear_b))
            .reshape(&[n, c])
    }

    /// Full forward for one sentence: `(N, C, T)` → logits `(N, num_classes)`.
    pub fn forward_sentence(
        &self,
        neuros: &Tensor,
        subject_ids: &[usize],
        chan_pos: Option<&Tensor>,
    ) -> Tensor {
        let n = neuros.shape[0];
        let embeds = self.encode_keystrokes(neuros, subject_ids, chan_pos);
        let seq = embeds.reshape(&[1, n, embeds.shape[1]]);
        let mask = vec![true; n];
        let refined = self.refine_sentence(&seq, &mask);
        self.classify(&refined.reshape(&[n, refined.shape[2]]))
    }
}