pio-arch 0.2.0

Perceiver IO-style permutation-invariant set models in PyTorch.

pio-arch

Perceiver IO-style permutation-invariant set models in PyTorch, for tabular, mixed-row, and sentence/set-structured data.

pio-arch provides a small set of composable building blocks for problems where each sample is a variable-length bag of heterogeneous features (numeric, categorical, free text) and the model must produce one or more predictions per sample.

The flagship model — PIOAttentionModel — is structurally Perceiver IO:

• learnable latent queries cross-attend to the variable-length context
• optional self-attention on the latent stack
• learnable task queries cross-attend to the latents
• per-task or shared MLP heads produce one prediction per task

All stage dimensions (context, latent, task query) are independently configurable. Permutation invariance over the input set is preserved end-to-end and is enforced by tests.

---

Install

Base install (PyTorch + numpy only):

pip install pio-arch

With the data-science / notebook extras (LightGBM, sentence-transformers, Optuna, polars, etc.):

pip install "pio-arch[model-dev]"

Python 3.12+ required.

---

Quick start

import torch
from pio_arch.models.pio_attention import PIOAttentionModel
from pio_arch.models.sentence_encoder import TaskSpec, masked_multitask_loss

tasks = [TaskSpec("fraud", "binary"), TaskSpec("amount", "regression", weight=0.5)]

model = PIOAttentionModel(
    input_dim=384,          # raw context-row dim (e.g. sentence-transformer output)
    tasks=tasks,
    model_dim=64,           # context stage
    latent_dim=64,          # latent bottleneck (defaults to model_dim)
    task_dim=32,            # task query stage (defaults to latent_dim)
    num_latents=8,
    num_context_self_attn_blocks=1,
    num_latent_self_attn_blocks=1,
    num_heads=4,
)

context = torch.randn(4, 12, 384)               # [batch, n_rows, input_dim]
padding_mask = torch.zeros(4, 12, dtype=torch.bool)  # True at padded rows

predictions = model(context, padding_mask)
# predictions = {"fraud": [4, 1], "amount": [4, 1]}

targets = torch.stack([torch.rand(4), torch.randn(4)], dim=1)
target_mask = torch.ones(4, 2, dtype=torch.bool)
loss, logs = masked_multitask_loss(predictions, targets, target_mask, tasks)
loss.backward()

A native PyTorch training loop is available in pio_arch.utils.train:

from pio_arch.utils.train import Trainer, make_masked_multitask_loss_fn

trainer = Trainer(
    model=model,
    optimizer=torch.optim.AdamW(model.parameters(), lr=3e-4),
    loss_fn=make_masked_multitask_loss_fn(tasks),
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
)
trainer.fit(train_loader, val_loader, num_epochs=5)

---

What's in the box

Module	What it provides
pio_arch.models.embedder	UniversalFeatureTransformer — type-aware raw-feature embedder (RFF for numeric, learned embedding for discrete).
pio_arch.models.context_encoder	ContextRowEncoder — mixed numeric / discrete / projected-text row encoder with feature-type embeddings.
pio_arch.models.sentence_encoder	SentenceSetEncoder (masked-mean baseline), SentenceSetMultiTaskModel, MultiTaskHead, TaskSpec, masked_multitask_loss.
pio_arch.models.pio_attention	PIOAttentionModel (Perceiver IO with task queries), SelfAttentionBlock, CrossAttentionBlock.
pio_arch.utils.rff	Shared Random Fourier Feature helpers.
pio_arch.utils.collate	sentence_set_collate — tensor-only collate.
pio_arch.utils.data	SentenceSetDataset + collate_sentence_set_batch for trainer-ready dict batches.
pio_arch.utils.train	Trainer, train_one_epoch, evaluate, make_masked_multitask_loss_fn.
pio_arch.utils.text_dropout	TextRowDropout — training-time row dropout for robustness.
pio_arch.utils.sentence_pool	pool_sentence_embeddings — DeepSets-style aggregator collapsing [B, N, D] → [B, D].

---

Notebooks

End-to-end walkthroughs live under notebooks/:

• sentence_only_pio.ipynb — sentence-only PIO walkthrough using a frozen MiniLM sentence-transformer. Compares the masked-mean baseline, PIOAttentionModel, and a DeepSets-style pooled baseline.
• full_pio_walkthrough.ipynb — full mixed-input pipeline with numeric, categorical, list-of-codes, and list-of-reason-sentences features.
• optuna_sweep.ipynb — Optuna hyperparameter sweep over the PIO architecture, with typical Perceiver IO dim guidance.

A synthetic 50k-row dataset (notebooks/pio_synthetic_50k.parquet) ships with the repo so every notebook can run end-to-end out of the box.

---

Reading the code

LEARNING_GUIDE.md is a staged tour of the source — which files to read in what order, what to note in each, and self-check exercises. Designed for ~4 hours start-to-finish with the exercises, or ~1 hour as a skim.

---

Design notes

• context.md — reference material on the broader permutation-invariant architecture design space (DeepSets, Set Transformer, Perceiver, Slot Attention). Historical / background only — not all of it is implemented in this repo.
• notebooks/pio_encoder_design.md — current encoder + attention design rationale.
• notebooks/pio_text_embedding_plan.md — sentence-transformer integration, list-valued field semantics, deployment considerations.
• AGENT_GUIDE.shared.md — tool-neutral project rules for any AI coding agent or human reading the repo for the first time.

---

Development

# clone, then
python3.12 -m venv .venv
.venv/bin/pip install -e ".[model-dev]" --group dev
.venv/bin/pre-commit install

# run tests + coverage
.venv/bin/pytest

# lint + format
.venv/bin/ruff check pio_arch tests
.venv/bin/ruff format --check pio_arch tests

# run the full pre-commit suite
.venv/bin/pre-commit run --all-files

The package currently has 162 unit + integration tests at 100% line coverage on the runtime code.

To rebuild the walkthrough notebooks after editing them:

.venv/bin/python scripts/build_walkthrough_notebooks.py

To smoke-test the notebook code paths on a 2k-row subsample without running the full epochs:

.venv/bin/python scripts/smoke_test_notebooks.py

---

License

MIT. See LICENSE.