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).

Version	Paper / use case	Rust API
V1	Nature Neuroscience 2026 — keystroke-aligned windows + sentence transformer	`V1Pipeline`
V2	Sentence-level CTC + TinyLlama+LoRA beam search	`Pipeline`

Layout

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

Install

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.

Quick start — parity (CI)

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:

Test	Command
V2 tiny encoder / decode / e2e / LLM	`cargo test -p brain2qwerty --tests`
V2 production (dim=1024)	`cargo test -p brain2qwerty --test v2_parity`
V1 keystroke pipeline	`cargo test -p brain2qwerty --test v1_parity`
RLX Conformer tail	`cargo test -p brain2qwerty --features rlx-encoder,rlx-cpu,rlx-mlx,rlx-metal --test rlx_encoder_parity`

Weight conversion

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.

Inference CLI (V2)

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.

Library usage

V2

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);

V1

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);

Architecture

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

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

Module	Role
`src/model/`	Pure-Rust V2 reference encoder (oracle vs Python)
`src/model_rlx/`	RLX-accelerated Conformer tail (default encoder)
`src/decode/`	CTC greedy, space segmenter, intra-word pooler
`src/llm/`	TinyLlama CPU decoder, LoRA merge, beam search
`src/v1/`	Full V1 keystroke pipeline

Scope

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).

Python reference

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.

License

Apache-2.0 (see upstream Brain2Qwerty license).

brain2qwerty 0.0.1