Install, verify data, train, evaluate

Requirements

Python 3.9 or later, and PyTorch 2.0 or later. The rest of the stack comes from requirements.txt: transformers for the BART decoder, mne for EEG preprocessing, plus numpy, scipy, pandas, lightning, h5py, jiwer, and python-Levenshtein for evaluation.

A CPU or an Apple Silicon Mac with MPS is enough for --quick-test and the CTC model. The BART model is trained with --fp16, which requires an NVIDIA GPU with CUDA; NEST v2 BART has not been trained to completion yet because that GPU time has not been available.

Install

Clone the repository, create a virtual environment, and install the dependencies.

git clone https://github.com/wazder/nest
cd nest
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

There is no PyPI package. pip install nest does not exist, and pip install -e . does not work yet: pyproject.toml has no [project] table, so there is nothing for pip to build. Run scripts directly from a cloned checkout, as shown throughout this page.

Data

NEST trains on the ZuCo corpus. Request access at osf.io/q3zws and place the extracted pickles so the layout matches what the loader expects, one pickle file per task:

ZuCo_Dataset/ZuCo/task1-SR/pickle/task1-SR-dataset.pickle
ZuCo_Dataset/ZuCo/task2-NR/pickle/task2-NR-dataset.pickle
ZuCo_Dataset/ZuCo/task3-TSR/pickle/task3-TSR-dataset.pickle

Verify the dataset loads and reports the expected number of sentence-subject pairs:

python src/data/zuco_pickle_dataset.py ZuCo_Dataset/ZuCo

Quick test

Before a full run, confirm the pipeline works end to end. --quick-test trains for 2 epochs on 50 samples and skips the held-out test evaluation.

python scripts/train_nest_v2.py --quick-test --model ctc

Train

The training entry point is scripts/train_nest_v2.py. It supports two model types, ctc and bart, selected with --model.

CTC baseline (linear CTC head, runs on CPU, MPS or GPU):

python scripts/train_nest_v2.py --model ctc --epochs 200

BART decoder (best quality, needs an NVIDIA GPU for --fp16):

python scripts/train_nest_v2.py \
    --model bart \
    --epochs 200 \
    --batch-size 16 \
    --fp16 \
    --d-model 768 \
    --num-layers 6 \
    --tasks task1-SR task2-NR task3-TSR

Flags that exist on train_nest_v2.py, verified against its argument parser:

Evaluate

Evaluation is not a separate script. Pass a checkpoint to --resume, point --output-dir at that checkpoint's directory so the final evaluation step reloads it, and set --epochs 0 to skip training:

python scripts/train_nest_v2.py \
    --resume results/nest_v2_bart_TIMESTAMP/best_model.pt \
    --output-dir results/nest_v2_bart_TIMESTAMP \
    --epochs 0

Results are written to results.json with a test_wer and a test_cer field: word error rate and character error rate, both computed as Levenshtein edit distance against the reference text. The BART path decodes with model.generate(), free autoregressive decoding with no teacher forcing (greedy in the current implementation). The CTC path decodes greedily. No checkpoint exists yet, so no results.json has been produced by this path.

Cloud

Local training needs an NVIDIA GPU for the BART run. notebooks/NEST_CloudTraining.ipynb is a ready-to-run Colab notebook for that case. On an A100, a full 200-epoch BART run is estimated at 8 to 12 hours; that estimate has not been measured end to end because the run has not completed.

Repository layout

src/          preprocessing, models (nest.py, nest_bart.py, nest_v2.py),
              data loaders, training and evaluation modules
scripts/      train_nest_v2.py, train_nest_bart.py, upload_to_huggingface.py
notebooks/    NEST_CloudTraining.ipynb (Colab)
papers/       NEST_manuscript.md (technical report draft)
configs/      YAML model configurations
tests/        80 test functions across 5 files
website/      this site (Cloudflare Pages)

Known issues

Track these and file new ones at github.com/wazder/nest/issues.

Evaluation protocol commitment

No word error rate is published on this site until it is measured with free decoding on held-out subjects, alongside a noise-input control. The full commitment and the reasoning behind it are on the research page.