Metadata-Version: 2.4
Name: pio-arch
Version: 0.2.1
Summary: Perceiver IO-style permutation-invariant set models in PyTorch.
Author-email: Benjamin Cross <btcross26@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: attention,deep-learning,deepsets,multi-task,perceiver,perceiver-io,permutation-invariant,pytorch,set-models,set-transformer
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.12
Requires-Dist: numpy>=1.24
Requires-Dist: torch<2.3,>=2.0; platform_system == 'Darwin' and platform_machine == 'x86_64'
Requires-Dist: torch>=2.0; platform_system != 'Darwin' or platform_machine != 'x86_64'
Provides-Extra: model-dev
Requires-Dist: lightgbm>=4.0; extra == 'model-dev'
Requires-Dist: matplotlib>=3.7; extra == 'model-dev'
Requires-Dist: optuna>=3.5; extra == 'model-dev'
Requires-Dist: pandas>=2.0; extra == 'model-dev'
Requires-Dist: polars>=1.0; extra == 'model-dev'
Requires-Dist: pyarrow>=15.0; extra == 'model-dev'
Requires-Dist: scikit-learn>=1.3; extra == 'model-dev'
Requires-Dist: seaborn>=0.12; extra == 'model-dev'
Requires-Dist: sentence-transformers<4.0,>=3.0; extra == 'model-dev'
Requires-Dist: shap>=0.44; extra == 'model-dev'
Requires-Dist: tqdm>=4.65; extra == 'model-dev'
Requires-Dist: transformers<5.0,>=4.41; extra == 'model-dev'
Requires-Dist: xgboost>=2.0; extra == 'model-dev'
Description-Content-Type: text/markdown

# pio-arch

[![tests](https://github.com/btcross26/pio-arch/actions/workflows/ci.yml/badge.svg)](https://github.com/btcross26/pio-arch/actions/workflows/ci.yml)

[Perceiver IO](https://arxiv.org/abs/2107.14795)-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):

```bash
pip install pio-arch
```

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

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

Python 3.12+ required.

---

## Quick start

```python
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),
    TaskSpec("time_to_fraud", "survival_discrete", output_dim=10),  # 10 day bins
]

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], "time_to_fraud": [4, 10]}

from pio_arch.models.survival import build_discrete_hazard_targets
fraud_targets = torch.rand(4)
amount_targets = torch.randn(4)
ttf_targets, ttf_mask = build_discrete_hazard_targets(
    torch.tensor([2, 5, 0, 9]), torch.tensor([True, False, True, False]), num_bins=10
)
targets = torch.cat([fraud_targets.unsqueeze(1), amount_targets.unsqueeze(1), ttf_targets], dim=1)
target_mask = torch.cat([torch.ones(4, 2, dtype=torch.bool), ttf_mask], dim=1)
loss, logs = masked_multitask_loss(predictions, targets, target_mask, tasks)
loss.backward()
```

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

```python
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.models.contrastive` | `ContrastiveTaskSpec`, `ContrastiveHead`, `contrastive_infonce_loss` — SimCLR-style InfoNCE auxiliary on pooled latents (v0.2.0). |
| `pio_arch.models.survival` | `build_discrete_hazard_targets`, `survival_from_hazards`, `weibull_nll`, `weibull_survival`, and friends — discrete-time hazard and Weibull AFT primitives backing the `"survival_discrete"` and `"survival_weibull"` task kinds (v0.2.1). |
| `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/`](notebooks/):

- **[`sentence_only_pio.ipynb`](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`](notebooks/full_pio_walkthrough.ipynb)** —
  full mixed-input pipeline with numeric, categorical, list-of-codes, and
  list-of-reason-sentences features.
- **[`optuna_sweep.ipynb`](notebooks/optuna_sweep.ipynb)** — Optuna
  hyperparameter sweep over the PIO architecture, with typical
  Perceiver IO dim guidance.
- **[`contrastive_multitask_walkthrough.ipynb`](notebooks/contrastive_multitask_walkthrough.ipynb)** —
  v0.2.0 walkthrough for the contrastive auxiliary task, multi-target
  heads, Kendall σ trajectory plot, and three freeze helpers for fine-
  tuning workflows.
- **[`discrete_hazard_walkthrough.ipynb`](notebooks/discrete_hazard_walkthrough.ipynb)** —
  v0.2.1 walkthrough for the `"survival_discrete"` task kind: discrete-
  time hazard math, target/mask construction, plotting `S(t)` curves, and
  joint contrastive + survival training.
- **[`weibull_aft_walkthrough.ipynb`](notebooks/weibull_aft_walkthrough.ipynb)** —
  v0.2.1 walkthrough for the `"survival_weibull"` task kind: parametric
  AFT with right-censoring and continuous-time inference grids.

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`](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

```bash
# 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:

```bash
.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:

```bash
.venv/bin/python scripts/smoke_test_notebooks.py
```

---

## License

MIT. See [`LICENSE`](LICENSE).
