brain2qwerty 0.0.1

# brain2qwerty-rs

Rust inference for [**Brain2Qwerty**](https://github.com/facebookresearch/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](https://www.nature.com/articles/s41593-026-02303-2) — keystroke-aligned windows + sentence transformer | [`V1Pipeline`](crates/brain2qwerty/src/v1/pipeline.rs) |
| **V2** | Sentence-level CTC + TinyLlama+LoRA beam search | [`Pipeline`](crates/brain2qwerty/src/pipeline.rs) |

## 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

```bash
cargo build --release -p brain2qwerty
```

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

```bash
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`](https://github.com/facebookresearch/brain2qwerty) repo on `PYTHONPATH`:

```bash
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:

```bash
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`](data/README.md) for the full data layout.

## Inference CLI (V2)

```bash
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

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

```rust
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

**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
```

| 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:

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