Sequence modelling · NLP · on-device inference
Six emotions,
176,000 parameters,
zero servers.
A bidirectional LSTM trained on — tweets from the dair-ai/emotion corpus, quantised to int8 and shipped as a — KB binary. Everything below — the prediction, the token attributions, the nearest-neighbour search — is computed on your machine, in plain JavaScript, with no framework and no API call.
The model, live
Type anything. The network re-runs on every keystroke — a forward pass, one leave-one-out pass per token, and one pass per prefix — inside a web worker.
Warm tokens push the winning class up; cool tokens hold it down. Exact leave-one-out ablation — one full forward pass per token, no gradient approximation. Reported in log-odds because a saturated softmax flattens every raw Δp to zero; hover a token for both numbers.
Each step re-encodes the truncated prefix from scratch, so the curve shows what the model would have answered had the sentence stopped there.
Architecture
Four layers, faithful to the original Keras notebook — re-implemented in PyTorch so the initialisers, gate order and unmasked padding all match.
Why the second BiLSTM sees 40 channels
The first bidirectional layer returns sequences, concatenating a 20-unit forward pass with a 20-unit backward pass at every timestep. The second consumes that 40-channel sequence and emits only its endpoints: forward state at t=49, backward state at t=0.
Padding is not masked
The original model had no mask_zero, so all 50 timesteps — including padding — flow through the recurrence. The port keeps that behaviour rather than quietly "fixing" it, because fixing it would no longer be the same model.
The data, honestly
20,000 self-reported emotion tweets, heavily imbalanced. Everything the model can and cannot learn starts here.
Rank-frequency truncation at 10,000 types is the single largest information bottleneck in the pipeline — everything rarer collapses to one <UNK> vector.
Training run
Adam, batch 32, early stopping on validation accuracy with patience 2. One run, no cherry-picking — the curve below is the run that produced the weights you just used.
Evaluation
2,000 held-out tweets the model never saw. Accuracy alone would hide the interesting parts, so: per-class scores, a confusion matrix you can click, a reliability diagram and a risk–coverage curve.
Where it breaks
The failures are not random. Two structural confusions account for most of the lost accuracy, and both are visible in the label distribution.
Click any example to load it into the live model above.
What the embedding layer learned
16 dimensions, trained from scratch on 16k tweets — no GloVe, no pretraining. Projected to two principal components and coloured by the emotion the model predicts when the token is shown alone.
Cosine similarity over the raw 16-d embedding matrix, computed in the browser across all 10,000 rows.
Getting it into a browser
Training a model is half the job. The other half is making it small enough, fast enough and verifiable enough to hand to a stranger on a phone.
Symmetric per-tensor int8 for weight matrices, float32 for biases. The fidelity numbers are measured against the float32 checkpoint on all 2,000 test tweets — not estimated.
No TensorFlow.js (≈900 KB), no ONNX runtime, no WASM blob. The forward pass is 90 lines of typed-array arithmetic.
Wall time of a batch of single forward passes at sequence length 50, divided by the batch size and timed right now in your browser.
The Python reference implementation in ml/export.py emits golden probability vectors; the JavaScript engine reproduces them to the shown tolerance. Parity is asserted in CI, not assumed.
# 1 · fetch the official 16k/2k/2k split
python ml/fetch_data.py
# 2 · train (CPU, ~6 min) — writes ml/artifacts/
python ml/train.py --config baseline
# 3 · evaluate + quantise + emit docs/model/
python ml/export.py --config baseline
# 4 · assert the JS engine matches the Python reference
node ml/parity.mjsModel card
What it is for, and what it is not for.
- Demonstrating end-to-end sequence classification: data → training → evaluation → quantised deployment.
- Coarse emotional tone on short, first-person English text.
- Clinical, safety, hiring or moderation decisions about real people.
- Sarcasm, negation-heavy text, code-switching, or any language but English.
- Text unlike the training distribution — the corpus is self-reported "i feel …" statements, and the model leans hard on that framing.