brain2qwerty-rs

Rust inference for Brain2Qwerty: MEG/EEG neural decoding to text. Implements both published pipelines with numeric parity against the official Python reference (Pearson r ≥ 0.999999 on continuous tensors; exact match on discrete outputs).

brain2qwerty-rs/
├── crates/brain2qwerty/   # library + brain2qwerty-infer CLI
├── scripts/               # parity refs, checkpoint conversion, test runner
├── data/                  # configs (committed), weights & refs (generated)
└── bench/                 # backend benchmarks

cargo build --release -p brain2qwerty

Default features enable the RLX Conformer tail (rlx-encoder, rlx-cpu). For reference-only CPU inference:

cargo build --release -p brain2qwerty --no-default-features --features pure-rust

Apple Silicon backends: --features apple-silicon (Metal, MLX, CoreML/ANE, wgpu).
NVIDIA: --features nvidia-gpu.

Requires the Python brain2qwerty repo on PYTHONPATH:

export PYTHONPATH=/path/to/brain2qwerty

# Generate all reference tensors + tiny weights (V1, V2 tiny, V2 prod)
python3 scripts/generate_parity_refs.py --profile all
python3 scripts/generate_v1_parity_refs.py

# Run every parity test
bash scripts/run_all_parity.sh

Individual suites:

Convert a Lightning V2 checkpoint to safetensors:

python3 scripts/convert_checkpoint.py \
  --ckpt /path/to/checkpoint.ckpt \
  --output data/

Produces data/encoder.safetensors (encoder + segmenter + word adapter) and optionally data/llm/ (TinyLlama + LoRA).

See data/README.md for the full data layout.

cargo run --release -p brain2qwerty --features rlx-encoder,rlx-metal --bin brain2qwerty-infer -- \
  --config data/config.yaml \
  --encoder-weights data/encoder.safetensors \
  --llm-weights data/llm \
  --input neuro.bin \
  --chan-pos chan_pos.bin \
  --subject 0 \
  --output predictions.json

Flags:

• --backend rlx (default) or --backend rust — RLX vs pure-Rust reference encoder
• --device auto|cpu|metal|mlx|cuda|… — RLX device when using RLX backend

Without --llm-weights, output is CTC text only.

use brain2qwerty::pipeline::{InferenceInput, Pipeline, PipelineOptions};

let mut pipeline = Pipeline::from_paths_with_options(
    "data/config.yaml",
    "data/encoder.safetensors",
    Some("data/llm"),
    PipelineOptions::default(),
)?;
let out = pipeline.run(&InferenceInput {
    neuros,           // (B, T, C)
    subject_ids: vec![0],
    chan_pos: Some(chan_pos),
})?;
println!("{}", out.pred_text);

Input layout is (N, channels, time) per keystroke window (different from V2).

use brain2qwerty::{V1Input, V1Pipeline};

let pipeline = V1Pipeline::from_weights("data/v1_config.yaml", "data/encoder_v1.safetensors")?;
let out = pipeline.run(&V1Input {
    neuros,           // (N, C, T)
    subject_ids: vec![0; n],
    chan_pos: Some(chan_pos),
})?;
println!("{}", out.pred_text);

V2

MEG (B,T,C) → ConvConformer → temporal downsample → Conformer → CTC
  → CTCSpaceSegmenter → word adapter → TinyLlama prefix + beam search → pred_text

V1

MEG window (N,C,T) per keystroke → SimpleConvTimeAgg + Bahdanau attention
  → SentenceTransformer (ALiBi) → linear → greedy char decode

Included: V1 and V2 inference (encoder, decode, LLM for V2), safetensors loading, RLX backends, parity tests.

Not included: training, dataloaders, augmentations, V1 n-gram LM post-processing (see Python brain2qwerty_v1/scripts/ngram_decoding.py).

V2 model registration requires a side-effect import in Python:

import brain2qwerty_v2.models  # registers ConvConformer

Parity ref generators: scripts/generate_parity_refs.py, scripts/generate_v1_parity_refs.py.

Apache-2.0 (see upstream Brain2Qwerty license).