freyja-readout

Linear and nonlinear readouts between column statistical profiles (Φ, 30-d) and column embeddings (Ψ, 384-d MiniLM by default).

This repo answers a single empirical question: given a column's distributional profile (cardinality, entropy, frequency stats, length stats, …), how well can we predict its sentence-encoder embedding, and vice versa? It ships:

• a drop-in package (freyja_readout) exposing three model classes with identical fit / predict / evaluate / save / load API,
• six pre-trained checkpoints (linear, MLP, PCA-bottleneck) for the FREYJA and SANTOS column corpora,
• the profiler + encoder pipeline that produced them,
• reproduction scripts for every figure and table in results/.

Headline numbers

Inverse direction (Φ → Ψ), held-out test set, pooled R² with the training-mean baseline. Higher is better; 0 means "as good as predicting the mean".

corpus	n_train / n_test	Linear ridge	MLP 30→128→D	PCA bottleneck k=50
freyja_minilm_v1	487 / 122	+0.129	+0.193	+0.341
santos_minilm_v1	3756 / 939	+0.141	+0.532	+0.472

The PCA bottleneck wins on the small corpus (regularization), the MLP wins on the large corpus (capacity). Both confirm the underlying signal is real and substantially nonlinear.

Aside on a number you may have seen elsewhere. A naïve closed-form ridge fit without an intercept on raw embedding targets reports R² ≈ −2.4. That is an artifact of forgetting to re-center predictions by the train mean when targets have non-zero mean (e.g. L2-normalised embeddings have ‖Ȳ_train‖ ≈ 0.85). All linear models in this repo use sklearn.linear_model.Ridge(fit_intercept=True), which fixes it. See scripts/01_inverse_verification.py for the full diagnosis.

Install

git clone https://github.com/dtim-upc/freyja-readout.git
cd freyja-readout
python -m venv .venv && source .venv/bin/activate
pip install -e .

Tested on macOS 14+ / Linux with Python 3.10–3.12, PyTorch 2.x, MPS or CPU.

Quickstart

from freyja_readout import load_readout
import numpy as np

# Load any of the six shipped checkpoints
model = load_readout("artifacts/santos_minilm_v1/nn")
print(model)           # NNReadout(inverse, 30->384)
print(model.train_metrics_["pooled_r2"])   # 0.5320...

# Predict an embedding from a 30-d profile (z-scored with the corpus's mu/sigma)
phi = np.zeros((1, 30))     # zero = "average column" in this corpus
psi_hat = model.predict(phi)
print(psi_hat.shape)        # (1, 384)

From a folder of CSVs to predictions in three calls

from freyja_readout import ColumnEncoder, load_readout

enc   = ColumnEncoder()                                  # off-the-shelf MiniLM-L6
model = load_readout("artifacts/santos_minilm_v1/nn")

Phi_z, Psi, meta, info = enc.featurize("path/to/csv_folder/")
print(meta.head())                                       # dataset_name, attribute_name
psi_pred = model.predict(Phi_z)                          # predicted embeddings
report   = model.evaluate(Phi_z, Psi)                    # how well does it explain Psi?

See notebooks/01_quickstart.ipynb for a runnable version.

Public API

Models — freyja_readout.{Linear,NN,PCABottleneck}Readout

All three implement the same Readout interface:

class Readout:
    direction: "forward" | "inverse"          # Ψ→Φ or Φ→Ψ
    in_dim_: int
    out_dim_: int
    train_baseline_mean_: np.ndarray          # for pooled R^2

    def fit(self, X, Y, X_val=None, Y_val=None, **kwargs) -> Self
    def predict(self, X) -> np.ndarray         # (n, out_dim_)
    def evaluate(self, X, Y) -> dict           # pooled_r2, per_dim_r2, cosine, ...
    def save(self, path) -> Path               # writes a directory
    @classmethod
    def load(cls, path) -> "Readout"           # returns the right subclass

class	purpose	when to use
LinearReadout	ridge regression with intercept	fastest baseline, near-ceiling on small n
NNReadout	30→128→D MLP, weight-decay grid	best on larger n, captures nonlinearity
PCABottleneckReadout	inner readout predicts top-k PCs of Ψ, then projects back	best when n is small and Ψ is high-dim — regularizes by construction

load_readout(path) is a convenience that picks the right subclass from the saved config.json.

Encoder — freyja_readout.ColumnEncoder

ColumnEncoder(sentence_encoder="sentence-transformers/all-MiniLM-L6-v2",
              device=None,                  # auto: mps / cuda / cpu
              profile_mu=None, profile_sigma=None,
              inv_cardinality=True,
              max_workers=8)
    .profile(csv_folder)        -> pd.DataFrame                 # FREYJA profiler
    .embed(csv_folder)          -> (np.ndarray, pd.DataFrame)   # (Psi, meta)
    .featurize(csv_folder)      -> (Phi_z, Psi, meta, info)     # both, aligned

Pass profile_mu / profile_sigma from a checkpoint's profile_stats.npz if you want SANTOS-trained statistics applied to a new corpus (so the predictions land in the same z-space the model was trained in).

Metrics — freyja_readout.metrics

pooled_r2(y_true, y_pred, baseline_mean) -> float
per_dim_r2(y_true, y_pred)               -> np.ndarray
uniform_average_r2(y_true, y_pred)       -> float        # sklearn default
row_cosine_similarity(A, B)              -> np.ndarray
evaluate(y_true, y_pred, baseline_mean)  -> dict          # one-call summary

The pooled R² is the primary metric used throughout. Per-column R² is brittle on low-variance dimensions (a few of the 384 BERT dims have very small variance and produce huge negative R² when predictions are slightly biased), which is why the headline numbers use the pooled form.

Repo layout

freyja-readout/
├── freyja/                        # vendored upstream FREYJA tree (untouched)
│   ├── app/                       # profiler, distances, model, API
│   └── ground_truths/             # reference ground-truth CSVs
│
├── freyja_readout/                # our package
│   ├── metrics.py                 # pooled R², per-dim R², cosine
│   ├── encoder.py                 # ColumnEncoder (profile + embed)
│   └── readout/                   # Readout base + Linear / NN / PCA bottleneck
│
├── artifacts/                     # six pre-trained checkpoints
│   ├── freyja_minilm_v1/{linear,nn,pca_bottleneck_k50}/
│   └── santos_minilm_v1/{linear,nn,pca_bottleneck_k50}/
│
├── data/
│   ├── filtered_data.pkl          # FREYJA stored Φ / Ψ_M / Ψ_B (5.8 MB)
│   ├── all_results.pkl            # legacy ridge fits (used by script 01 for diagnosis)
│   └── santos/                    # pre-derived SANTOS profiles + embeddings (13 MB)
│
├── scripts/
│   ├── train_checkpoints.py       # rebuilds artifacts/ from data/ in ~2 min
│   ├── download_santos.py         # one-shot SANTOS source CSV download
│   ├── 01_inverse_verification.py
│   ├── 02_effective_dimensionality.py
│   ├── 03_minilm_vs_bert.py
│   ├── 04_pca_bottleneck.py
│   ├── 05_nn_inverse_check.py
│   ├── 06_santos_quicktest.py
│   └── 07_cross_benchmark.py
│
├── results/                       # outputs from the seven scripts above (CSVs, PNGs, txt reports)
├── notebooks/01_quickstart.ipynb
├── tests/test_smoke.py            # 11 tests
├── pyproject.toml
└── requirements.txt

Reproduction

Every CSV/PNG already in results/ is shipped pre-computed. To regenerate any of them:

python scripts/01_inverse_verification.py    # ~5  s
python scripts/02_effective_dimensionality.py # ~10 s
python scripts/03_minilm_vs_bert.py          # ~10 s
python scripts/04_pca_bottleneck.py          # ~40 s
python scripts/05_nn_inverse_check.py        # ~15 s
python scripts/06_santos_quicktest.py        # ~15 s   (needs data/datalakes/santos_small/, see below)
python scripts/07_cross_benchmark.py         # ~2  min (needs data/datalakes/santos_small/)

python scripts/train_checkpoints.py          # ~2  min — rebuilds artifacts/

Numbers reproduce within ±1e-3 across BLAS / PyTorch versions.

Pulling the SANTOS source CSVs

Scripts 06 and 07 process the raw SANTOS small benchmark CSVs (~544 MB extracted), which we don't redistribute in this repo. To fetch them:

python scripts/download_santos.py

This downloads from the SANTOS authors' Zenodo record and unpacks under data/datalakes/santos_small/. If the download fails you can place the zip manually at data/datalakes/santos_benchmark.zip and re-run.

Use it on a new corpus

from freyja_readout import ColumnEncoder, NNReadout

# 1) profile + embed your CSVs
enc = ColumnEncoder()
Phi_z, Psi, meta, info = enc.featurize("path/to/your/csvs/")

# 2) split, fit, evaluate
import numpy as np
n = Phi_z.shape[0]; rng = np.random.default_rng(0); idx = np.arange(n); rng.shuffle(idx)
ntr = int(0.8 * n); tr, te = idx[:ntr], idx[ntr:]

model = NNReadout(direction="inverse").fit(
    Phi_z[tr], Psi[tr], X_val=Phi_z[te], Y_val=Psi[te]
)
print(model.evaluate(Phi_z[te], Psi[te]))

# 3) save for later
model.save("artifacts/my_corpus_v1/nn")

To predict embeddings later from new profiles only (no ground-truth Ψ needed), load the checkpoint and call model.predict(Phi_z). Make sure the new profiles are z-scored with the same (mu, sigma) the model was trained with — they live in profile_stats.npz next to the weights.

Tests

pip install pytest
pytest tests/

Eleven smoke tests cover: package imports, the two key edge cases of the metric (perfect prediction → 1, constant prediction → 0), loading each of the six shipped checkpoints, and Linear/NN save↔load roundtrips.

Notes on what is and isn't shipped

• freyja/app/ is vendored from the upstream FREYJA repository unchanged for in-place use; updates can be pulled in by re-vendoring.
• data/filtered_data.pkl is a derived artifact (Φ + Ψ for 609 columns from a curated joinable subset). The original CSVs that produced it are not redistributed here.
• data/santos/ contains pre-derived SANTOS profiles + off-the-shelf MiniLM-L6 embeddings (a ~13 MB convenience). The 544 MB of raw SANTOS CSVs is not in the repo; use download_santos.py.
• The freyja_minilm_v1 checkpoints were trained on stored embeddings produced by an upstream fine-tuned encoder. They will only give the reported numbers when applied to columns embedded by that encoder; for arbitrary new corpora use santos_minilm_v1 (off-the-shelf encoder) or train a fresh checkpoint with train_checkpoints.py.