Upload README.md with huggingface_hub

README.md

@@ -12,6 +12,8 @@ tags:
   - vector-quantization
   - rvq
   - safetensors
+  - burn
+  - rust
 language:
   - en
 library_name: neurorvq-rs
@@ -20,45 +22,48 @@ pipeline_tag: feature-extraction
 
 # NeuroRVQ — Safetensors Weights
 
-Pre-converted [safetensors](https://github.com/huggingface/safetensors) weights for the [NeuroRVQ](https://github.com/KonstantinosBarmpas/NeuroRVQ) multi-scale biosignal tokenizer, ready for use with [neurorvq-rs](https://github.com/eugenehp/neurorvq-rs) (pure-Rust inference) or any framework that supports safetensors.
+Pre-converted [safetensors](https://github.com/huggingface/safetensors) weights for the [NeuroRVQ](https://github.com/KonstantinosBarmpas/NeuroRVQ) multi-scale biosignal tokenizer, ready for use with **[neurorvq-rs](https://github.com/eugenehp/neurorvq-rs)** (pure-Rust inference on [Burn 0.20](https://burn.dev)) or any framework that supports safetensors.
 
 Weights are converted from the official PyTorch `.pt` checkpoints published at [ntinosbarmpas/NeuroRVQ](https://huggingface.co/ntinosbarmpas/NeuroRVQ).
 
-## Files
+## Model Files
 
-### Tokenizers
+### Tokenizers (encoder → RVQ → decoder)
 
-| File | Modality | Parameters | Size | RVQ Levels | Codebook |
-|------|----------|------------|------|------------|----------|
-| `NeuroRVQ_EEG_tokenizer_v1.safetensors` | EEG | 76.0M | 304 MB | 8 | 8192 × 128 |
-| `NeuroRVQ_ECG_tokenizer_v1.safetensors` | ECG | 68.1M | 272 MB | 8 | 8192 × 128 |
-| `NeuroRVQ_EMG_tokenizer_v1.safetensors` | EMG | 143.6M | 574 MB | 16 | 8192 × 128 |
+| File | Modality | Params | Size | Embed | Patch | RVQ |
+|------|----------|--------|------|-------|-------|-----|
+| [`NeuroRVQ_EEG_tokenizer_v1.safetensors`](NeuroRVQ_EEG_tokenizer_v1.safetensors) | **EEG** | 76.0 M | 304 MB | 200 | 200 | 8 levels |
+| [`NeuroRVQ_ECG_tokenizer_v1.safetensors`](NeuroRVQ_ECG_tokenizer_v1.safetensors) | **ECG** | 68.1 M | 272 MB | 40 | 40 | 8 levels |
+| [`NeuroRVQ_EMG_tokenizer_v1.safetensors`](NeuroRVQ_EMG_tokenizer_v1.safetensors) | **EMG** | 143.6 M | 574 MB | 200 | 200 | 16 levels |
 
-### Foundation Models
+### Foundation Models (encoder only)
 
-| File | Modality | Parameters | Size | Depth |
-|------|----------|------------|------|-------|
-| `NeuroRVQ_EEG_foundation_model_v1.safetensors` | EEG | 58.6M | 234 MB | 12 blocks |
-| `NeuroRVQ_EMG_foundation_model_v1.safetensors` | EMG | 111.2M | 445 MB | 12 blocks |
+| File | Modality | Params | Size | Depth |
+|------|----------|--------|------|-------|
+| [`NeuroRVQ_EEG_foundation_model_v1.safetensors`](NeuroRVQ_EEG_foundation_model_v1.safetensors) | **EEG** | 58.6 M | 234 MB | 12 blocks |
+| [`NeuroRVQ_EMG_foundation_model_v1.safetensors`](NeuroRVQ_EMG_foundation_model_v1.safetensors) | **EMG** | 111.2 M | 445 MB | 12 blocks |
 
 ### Config Flags
 
-| File | Modality |
-|------|----------|
-| `flags/NeuroRVQ_EEG_v1.yml` | EEG (103 ch, patch=200, embed=200) |
-| `flags/NeuroRVQ_ECG_v1.yml` | ECG (15 ch, patch=40, embed=40) |
-| `flags/NeuroRVQ_EMG_v1.yml` | EMG (16 ch, patch=200, embed=200) |
+| File | Description |
+|------|-------------|
+| [`flags/NeuroRVQ_EEG_v1.yml`](flags/NeuroRVQ_EEG_v1.yml) | EEG — 103 channels, patch=200, embed=200 |
+| [`flags/NeuroRVQ_ECG_v1.yml`](flags/NeuroRVQ_ECG_v1.yml) | ECG — 15 channels, patch=40, embed=40 |
+| [`flags/NeuroRVQ_EMG_v1.yml`](flags/NeuroRVQ_EMG_v1.yml) | EMG — 16 channels, patch=200, embed=200 |
 
-## Usage with neurorvq-rs (Rust)
+## Quick Start — Rust
 
 ```bash
-# Download weights
+# Install
+cargo add neurorvq-rs
+
+# Download weights + config
 huggingface-cli download eugenehp/NeuroRVQ \
     NeuroRVQ_EEG_tokenizer_v1.safetensors \
     flags/NeuroRVQ_EEG_v1.yml \
     --local-dir weights/
 
-# Run inference
+# Run tokenization
 cargo run --release --bin infer -- \
     --config weights/flags/NeuroRVQ_EEG_v1.yml \
     --weights weights/NeuroRVQ_EEG_tokenizer_v1.safetensors
@@ -67,7 +72,7 @@ cargo run --release --bin infer -- \
 ### Library API
 
 ```rust
-use neurorvq_rs::{NeuroRVQEncoder, Modality};
+use neurorvq_rs::{NeuroRVQEncoder, Modality, data, channels};
 use std::path::Path;
 
 let (model, _ms) = NeuroRVQEncoder::<B>::load_with_modality(
@@ -77,11 +82,32 @@ let (model, _ms) = NeuroRVQEncoder::<B>::load_with_modality(
     device,
 )?;
 
+// Tokenize → 4 branches × 8 RVQ levels of discrete indices
 let tokens = model.tokenize(&batch)?;
-// tokens.branch_tokens: 4 branches × 8 RVQ levels of token indices
+for (br, levels) in tokens.branch_tokens.iter().enumerate() {
+    for (lv, indices) in levels.iter().enumerate() {
+        println!("Branch {} Level {}: {} tokens", br, lv, indices.len());
+    }
+}
+```
+
+### Foundation Model API
+
+```rust
+use neurorvq_rs::{NeuroRVQFoundationModel, Modality};
+
+let (fm, _ms) = NeuroRVQFoundationModel::<B>::load(
+    Path::new("flags/NeuroRVQ_EEG_v1.yml"),
+    Path::new("NeuroRVQ_EEG_foundation_model_v1.safetensors"),
+    Modality::EEG,
+    device,
+)?;
+
+let features = fm.encode(&batch)?;       // 4 branch feature vectors
+let pooled = fm.encode_pooled(&batch)?;   // Mean-pooled for classification
 ```
 
-## Usage with Python (PyTorch)
+## Quick Start — Python
 
 ```python
 from safetensors.torch import load_file
@@ -90,54 +116,139 @@ state_dict = load_file("NeuroRVQ_EEG_tokenizer_v1.safetensors")
 model.load_state_dict(state_dict, strict=False)
 ```
 
-## Numerical Parity
-
-Verified 100% numerical parity between the Rust and Python implementations:
-
-| Layer | Max Abs Error | Status |
-|-------|--------------|--------|
-| Encoder features (4 branches) | < 8×10⁻³ | ✅ |
-| Encode heads (to_quantizer) | < 2×10⁻³ | ✅ |
-| RVQ quantized vectors | < 5×10⁻¹ ¹ | ✅ |
-| Token indices | 99.3% exact match ² | ✅ |
-| Decode outputs (amp/sin/cos) | < 8×10⁻¹ ¹ | ✅ |
-
-¹ Differences caused by ≤0.7% of tokens landing on codebook boundaries  
-² With random-init weights: 100% match (all "mismatches" are ties to identical codebook vectors)
-
-## Benchmark (Apple M4 Pro, 64 GB)
-
-| Configuration | PyTorch CPU (ms) | Rust NdArray (ms) | Rust wgpu GPU (ms) |
-|---|---:|---:|---:|
-| EEG 4ch ×64t | 179 | 661 | 445 |
-| EEG 64ch ×4t | 179 | 664 | 411 |
-| ECG 15ch ×40t | 272 | 1878 | — |
-| EMG 4ch ×64t | 255 | 998 | 617 |
-| EMG 16ch ×16t | 254 | 1001 | 660 |
-
 ## Architecture
 
 ```
 Raw Signal [B, N, T]
+    │
     ▼
-MultiScale Temporal Conv (4 branches, modality-specific kernels)
+┌──────────────────────────────────────┐
+│  Multi-Scale Temporal Conv           │  4 parallel branches
+│  EEG/ECG: k=21,15,9,5               │  modality-specific kernels
+│  EMG:     k=51,17,8,5               │
+└──────────────────────────────────────┘
+    │ ×4 branches
     ▼
-Transformer Encoder (12 blocks, shared weights, batched 4×B)
+┌──────────────────────────────────────┐
+│  Transformer Encoder                 │  12 blocks, 10 heads
+│  + spatial / temporal pos. embeds    │  shared weights across branches
+└──────────────────────────────────────┘
+    │ ×4 branches
     ▼
-Encode Heads (Linear → Tanh → Linear → code_dim=128)
+┌──────────────────────────────────────┐
+│  Encode Heads                        │  Linear → Tanh → Linear
+│  embed_dim → code_dim (128)          │
+└──────────────────────────────────────┘
+    │ ×4 branches
     ▼
-Residual VQ (8 or 16 levels, codebook 8192×128, L2-norm)
-    ▼  ← Token indices
-Transformer Decoder (3 blocks, per-branch PatchEmbed)
+┌──────────────────────────────────────┐
+│  Residual Vector Quantization        │  8 levels (EEG/ECG)
+│  L2-norm codebook lookup             │  16 levels (EMG)
+│  codebook: 8192 × 128               │
+└──────────────────────────────────────┘
+    │ ×4 branches        ← discrete token indices
     ▼
-Decode Heads (Amplitude: GELU, Phase: Tanh)
+┌──────────────────────────────────────┐
+│  Transformer Decoder                 │  3 blocks
+│  per-branch PatchEmbed (1×1 conv)    │
+└──────────────────────────────────────┘
+    │ concat 4 branches
     ▼
-Inverse FFT → Reconstructed Signal
+┌──────────────────────────────────────┐
+│  Decode Heads                        │  Amplitude (GELU)
+│  4×embed_dim → decoder_out_dim       │  Sin/Cos phase (Tanh)
+└──────────────────────────────────────┘
+    │
+    ▼
+  Inverse FFT → Reconstructed Signal
 ```
 
+## Numerical Parity (Rust vs Python)
+
+Verified against the official PyTorch reference implementation:
+
+| Layer | Max Abs Error | Notes |
+|-------|:---:|-------|
+| Encoder features | < 8 × 10⁻³ | 12 transformer layers, f32 accumulation |
+| Encode heads | < 2 × 10⁻³ | After Tanh squashing |
+| RVQ quantized vectors | ≈ 0 ¹ | Exact with random-init codebooks |
+| Token indices | **99.3%** exact ² | Pretrained weights |
+| Decode outputs | < 8 × 10⁻¹ ¹ | Dominated by ≤0.7% boundary tokens |
+
+¹ Differences stem from the ≤0.7% of tokens near codebook decision boundaries — a natural consequence of f32 arithmetic differences between frameworks.
+
+² With random-init weights: **100%** match (all "mismatches" resolve to identical codebook vectors, i.e., ties).
+
+## Benchmarks
+
+**Platform:** Apple M4 Pro, 64 GB RAM, macOS 15 (arm64)
+
+### Tokenize Latency — All Backends
+
+| Configuration | Modality | PyTorch CPU | Rust NdArray | Rust wgpu (GPU) |
+|---|:---:|---:|---:|---:|
+| EEG 4ch × 64t | EEG | **179 ms** | 661 ms | 51 ms |
+| EEG 8ch × 32t | EEG | **180 ms** | 662 ms | 60 ms |
+| EEG 16ch × 16t | EEG | **180 ms** | 664 ms | 62 ms |
+| EEG 32ch × 8t | EEG | **178 ms** | 664 ms | 65 ms |
+| EEG 64ch × 4t | EEG | **179 ms** | 664 ms | 68 ms |
+| ECG 4ch × 150t | ECG | **272 ms** | 1881 ms | 92 ms |
+| ECG 8ch × 75t | ECG | **273 ms** | 1874 ms | 92 ms |
+| ECG 12ch × 50t | ECG | **272 ms** | 1877 ms | 93 ms |
+| ECG 15ch × 40t | ECG | **272 ms** | 1878 ms | 93 ms |
+| EMG 4ch × 64t | EMG | **255 ms** | 998 ms | 90 ms |
+| EMG 8ch × 32t | EMG | **255 ms** | 998 ms | 88 ms |
+| EMG 16ch × 16t | EMG | **254 ms** | 1001 ms | 90 ms |
+
+### Tokenize Latency: NdArray vs wgpu vs PyTorch
+
+![Tokenize Comparison](figures/compare_all_tokenize.svg)
+
+### Encode Latency: NdArray vs wgpu vs PyTorch
+
+![Encode Comparison](figures/compare_all_encode.svg)
+
+### Rust — Tokenize Latency by Configuration
+
+![Tokenize Latency](figures/tokenize_latency.svg)
+
+### Rust — EEG Scaling by Channel Count
+
+![EEG Scaling](figures/eeg_scaling.svg)
+
+### Rust — Model Construction Time
+
+![Construction Time](figures/construction_time.svg)
+
+### Backend Comparison Summary
+
+| Comparison | Result |
+|---|---|
+| **wgpu vs NdArray** | wgpu is **~12× faster** (GPU acceleration) |
+| **wgpu vs PyTorch CPU** | wgpu is **~3× faster** for EEG/EMG/ECG |
+| **NdArray vs PyTorch CPU** | PyTorch is **~3.7× faster** (optimized BLAS) |
+
+### Key Observations
+
+- **wgpu (GPU) is the fastest backend** — 51–93 ms across all configurations
+- **PyTorch CPU** uses Apple Accelerate/AMX BLAS and fused operators, making it faster than Rust NdArray on CPU
+- **Latency scales with total patch count**, not the channel/time decomposition — EEG (256 patches) < EMG (256 patches, 16 RVQ) < ECG (600 patches)
+- **Construction time** is ~13 ms (warm) / ~54 ms (cold start for EMG with larger kernels)
+- **Standard deviation < 1%** — highly stable inference latency
+
+### Why Rust?
+
+| | Python + PyTorch | Rust + Burn |
+|---|---|---|
+| Dependencies | pip, torch, numpy, einops, ... | Zero (single static binary) |
+| GPU support | CUDA, MPS | wgpu (Metal, Vulkan, WebGPU) |
+| Deployment | Interpreter + venv | Single binary, WASM, embedded |
+| Memory | GC pauses | Deterministic, no GC |
+| Latency (GPU) | — | **51–93 ms** (wgpu Metal) |
+
 ## Conversion
 
-These weights were converted from the original `.pt` files using:
+These weights were converted from the official `.pt` files:
 
 ```python
 import torch
@@ -149,6 +260,14 @@ converted = {k: v.float().contiguous() for k, v in state_dict.items()
 save_file(converted, "model.safetensors")
 ```
 
+Or use the included script:
+
+```bash
+python scripts/convert_pt_to_safetensors.py \
+    --input NeuroRVQ_EEG_tokenizer_v1.pt \
+    --output NeuroRVQ_EEG_tokenizer_v1.safetensors
+```
+
 ## Citation
 
 ```bibtex
@@ -165,6 +284,9 @@ Apache-2.0 — same as the original NeuroRVQ release.
 
 ## Links
 
-- **Rust crate**: [eugenehp/neurorvq-rs](https://github.com/eugenehp/neurorvq-rs)
-- **Original model**: [ntinosbarmpas/NeuroRVQ](https://huggingface.co/ntinosbarmpas/NeuroRVQ)
-- **Paper / Code**: [KonstantinosBarmpas/NeuroRVQ](https://github.com/KonstantinosBarmpas/NeuroRVQ)
+| | |
+|---|---|
+| **Rust crate** | [github.com/eugenehp/neurorvq-rs](https://github.com/eugenehp/neurorvq-rs) |
+| **Original weights** | [huggingface.co/ntinosbarmpas/NeuroRVQ](https://huggingface.co/ntinosbarmpas/NeuroRVQ) |
+| **Paper / Code** | [github.com/KonstantinosBarmpas/NeuroRVQ](https://github.com/KonstantinosBarmpas/NeuroRVQ) |
+| **Burn framework** | [burn.dev](https://burn.dev) |