Units-Surprisal

Code for On the Proper Treatment of Units in Surprisal Theory.

Different unit inventories (acontextual whitespace words, Penn-Treebank tokens, GPT-2 subword tokens, characters) are defined as finite-state transducers (FSTs) built with pynini. Per-unit surprisal is computed by marginalizing a language model (GPT-2 Small) over all tokenizations consistent with a unit boundary, using the framework from Transducing Language Models.

Contents

• Repository layout
• Installation
• Data
• Reproducing the results
• Computing surprisal for your own units
• Citation

Repository layout

src/units_surprisal/
  main.py                     # CLI: runs the data-preparation pipeline stages
  transducers/                # FST builders: whitespace, ptb, character, model_tokens
  processors/                 # units_processor, unit_aggregator (fixations -> units)
  processing_scripts/         # one stage per pipeline step (extract/transduce/aggregate/...)
  unigram/                    # unigram-surprisal baseline estimators
  preprocessing/              # MECO fixation extraction, OCR alignment
  analysis/                   # run_gamm.R  (the paper's GAMM analysis)
  utils/                      # load_transducer, LM adapters
  tests/                      # pytest suite (`make test`)
  data/                       # downloaded from Hugging Face (see "Data" below)
scripts/
  run_all_gamm.sh             # fit GAMMs for every (unit, reading-time measure)
  plot_gamm_results.py        # render the paper's figures
  download_data.py            # fetch the MECO-derived data from Hugging Face

Installation

conda create -n units-surprisal python=3.12 && conda activate units-surprisal
conda install -c conda-forge pynini   # pynini needs OpenFST (easiest via conda)

pip install -e .            # core: transduced_lm (from source) + genlm, torch, ...

• The analysis additionally needs R with the mgcv and argparse packages: install.packages(c("mgcv", "argparse")).

Key from-source dependency: transduced_lm provides VectorizedFST, TransducedLM, Config, and GenLMRealpha for marginalized surprisal (and transitively pulls pynini, numpy, pandas, scipy, matplotlib, and genlm-bytes). After pip install -e . the pipeline is available as the units-surprisal command (or python -m units_surprisal). The data directory is resolved from the package location, or override it with the UNITS_SURPRISAL_DATA environment variable.

Data

The MECO-derived data is not stored in this repository; it is hosted on Hugging Face and downloaded into src/units_surprisal/data/:

python scripts/download_data.py units-surprisal   # English data + intermediates (paper)
python scripts/download_data.py meco-char          # full multilingual MECO (all 13 langs)

Dataset	Contents
samuki-hf/units-surprisal	English fixations, stimuli, and every pipeline intermediate needed to reproduce the results.
samuki-hf/meco-char	Full character-level MECO data for all 13 languages.

The underlying raw eye-tracking data is from MECO (Siegelman et al., 2022). The processing of the English MECO dataset is from Re et al., 2025.

Reproducing the results

The dataframes for English (subject_units, surprisal, unigram_probs) are on samuki-hf/units-surprisal. To reproduce the results you can fetch the dataframes and re-fit the GAMMs with run_gamm.R.

This requires installing R with the mgcv and argparse packages:

make data # fetch the dataframes from Hugging Face

# Fit GAMMs for every (unit inventory x reading-time measure) -> outputs/gamm_results
bash scripts/run_all_gamm.sh

# Render the paper figures and result tables
python scripts/plot_gamm_results.py --results_dir outputs/gamm_results
# -> outputs/gamm_figures/{gamm_dll_main.pdf, gamm_coefficients_appendix.pdf, *.tex}

A single cell can be reproduced directly, e.g.:

Rscript src/units_surprisal/analysis/run_gamm.R \
    --transducer whitespace --mode leading --dv total_dur_ms --spillover 2 \
    --out_dir outputs/gamm_results

Rebuilding intermediates from scratch

A Makefile wraps the whole flow (run make help for the targets):

make data        # download English data + intermediates from Hugging Face
make prepare     # CPU: extract-text + transduce + aggregate, all inventories
make surprisal   # GPU: per-unit LM surprisal, all inventories
make gamm        # fit GAMMs (needs R + mgcv) and render figures/tables
make all         # prepare + surprisal + gamm

To run a single stage by hand, use the CLI subcommands (per unit inventory and reading mode):

units-surprisal extract-text --language en
units-surprisal transduce      --transducer <whitespace|ptb|character|model_tokens> --mode <leading|trailing> --language en
units-surprisal subject-units  --transducer <...> --mode <...> --language en
units-surprisal subject-means  --transducer <...> --mode <...> --language en
units-surprisal surprisal      --transducer <...> --mode <...> --language en   # needs a GPU
units-surprisal unigram        --transducer <...> --mode <...> --language en   # unigram baseline (needs a GPU)
# or stages 2-4 for one inventory at once:
units-surprisal pipeline       --transducer <...> --mode <...> --language en

Unit inventory	--transducer	Description
Acontextual words	whitespace	SEP at whitespace boundaries (leading or trailing)
Contextual (PTB)	ptb	Penn-Treebank tokenizer FST
Subword tokens	model_tokens	GPT-2 BPE tokens
Characters	character	Character-level units (byte marginalization)

Surprisal computation requires a GPU.

Computing surprisal for your own units

Define a new unit and compute surprisal for it. The library API takes the text, a unit inventory, and a base LM:

import units_surprisal as us

# a built-in inventory by name
rows = us.surprisal_for("The cat sat.", units="whitespace", lm="gpt2", mode="leading")
for r in rows:
    print(r.unit, r.surprisal_bits) # -> UnitSurprisal(unit, logprob_nat, surprisal_bits, ...)

# your own units: any pynini FST that emits SEP at unit boundaries
from units_surprisal.transducers.whitespace import construct_whitespace_bytes
my_fst = construct_whitespace_bytes(mode="leading", delims=(32, 9, 10))  # split on space/tab/newline
rows = us.surprisal_for(text, units=my_fst, lm="gpt2")

# reuse one engine across a corpus (loads the LM + builds the beam once)
engine = us.build_engine(units="ptb", lm="gpt2", mode="leading")
for doc in corpus:
    rows = engine.score(doc)

The contract. A unit inventory is defined via finite-state transducer over the byte-symbol convention (transducers/symbols.py: bytes "0"–"255", EOS="256", EPS="257", SEP="258") that emits SEP at your unit boundaries. The LM is marginalized through that FST, so each unit's surprisal includes its SEP boundary mass. us.validate_byte_fst(fst) checks a custom FST against the convention.

Register a named inventory with one TransducerSpec and use aggregate=whitespace_units so units are recovered by splitting on SEP:

from units_surprisal.transducers.registry import TransducerSpec
from units_surprisal.transducers.aggregation import whitespace_units

us.register_transducer("my_units", TransducerSpec(
    name="my_units", fst_type="byte", uses_sep=True, special_syms=("258", "256"),
    build=lambda llm, mode: my_fst, aggregate=lambda toks, llm, seps: whitespace_units(toks, seps),
    config=us.get_spec("whitespace").config, strip_default=False, drop_space_default=False,
))
rows = us.surprisal_for(text, units="my_units", lm="gpt2")

Walkthrough: notebooks/getting_started.ipynb is a tutorial that builds the whitespace inventory end to end. whitespace.py is the reference FST builder; ptb_ported.py shows a cdrewrite/rule-based one. To take a custom inventory all the way to a GAMM, run its stages (transduce -> subject-units -> subject-means -> surprisal) and add its name to scripts/run_all_gamm.sh.

Citation

@inproceedings{kiegeland-etal-2026-proper,
    title = "On the Proper Treatment of Units in Surprisal Theory",
    author = "Kiegeland, Samuel  and
      Sn{\ae}bjarnarson, V{\'e}steinn  and
      Vieira, Tim  and
      Cotterell, Ryan",
    editor = "Liakata, Maria  and
      Moreira, Viviane P.  and
      Zhang, Jiajun  and
      Jurgens, David",
    booktitle = "Proceedings of the 64th Annual Meeting of the {A}ssociation for {C}omputational {L}inguistics (Volume 1: Long Papers)",
    month = jul,
    year = "2026",
    address = "San Diego, California, United States",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2026.acl-long.1485/",
    doi = "10.18653/v1/2026.acl-long.1485",
    pages = "32202--32224",
    ISBN = "979-8-89176-390-6",
}