"""P064: input corruption teaches an unavailable Set representation."""
from collections.abc import Iterator
from functools import partial
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
from reporting import report
import relflow as rf
PROOF_ID = "P064"
BUDGET = 400
def records(*, rows: int, seed: int, novel: bool = False) -> Iterator[dict]:
rng = np.random.default_rng(seed)
for label, nonempty in zip(rng.integers(32, size=rows), rng.random(rows) < 0.5, strict=True):
yield {"tags": [f"{'novel' if novel else 'tag'}-{label}"] if nonempty else [], "nonempty": float(nonempty)}
def build(p_unavailable: float) -> rf.Model:
model = rf.Model(
d_model=32,
n_layers=1,
n_heads=4,
batch_size=64,
dropout=0.0,
tags=rf.Set(p_unavailable=p_unavailable),
nonempty=rf.Number(mask=True, objective="mse"),
)
model.optimizer = rf.adamw(learning_rate=0.002)
return model
def predict(model: rf.Model, rows: list[dict]) -> np.ndarray:
return np.asarray(
[row["nonempty"]["content"] for row in model.predict(pa.Table.from_pylist(rows))["predictions"].to_pylist()]
)
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
metrics, checks = {"constant_rmse": 0.5}, {}
for probability in (0.0, 0.5):
lit.seed_everything(seed, workers=True)
model = build(probability)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=8192, seed=seed + 1),
validate=partial(records, rows=512, seed=seed + 2),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
devices=1,
max_epochs=-1,
max_steps=BUDGET if steps is None else min(steps, BUDGET),
deterministic=True,
logger=False,
enable_checkpointing=False,
enable_progress_bar=False,
enable_model_summary=False,
num_sanity_val_steps=0,
)
trainer.fit(model, datamodule=data)
vocabulary = rf.Set.vocabulary(model, "/tags")
arm = "augmented" if probability else "control"
for novel in (False, True):
rows = list(records(rows=2048, seed=seed + 3, novel=novel))
expected = np.asarray([row["nonempty"] for row in rows])
values = predict(model, rows)
key = f"{arm}_{'unknown' if novel else 'known'}"
metrics[key] = float(np.sqrt(np.mean((values - expected) ** 2)))
if not novel or probability:
checks[f"{key}: RMSE below 0.10"] = metrics[key] < 0.10
if novel:
changed = [{**row, "tags": ["another", "another"] if row["tags"] else []} for row in rows]
drift = float(np.max(np.abs(values - predict(model, changed))))
metrics[f"{arm}_duplicate_and_identity_drift"] = drift
checks[f"{arm}: unknown spelling and duplicates are invariant"] = drift < 1e-5
checks[f"{arm}: unknown evaluation leaves vocabulary frozen"] = rf.Set.vocabulary(model, "/tags") == vocabulary
checks["Trained fallback improves unknown-input RMSE by more than 0.40"] = (
metrics["control_unknown"] - metrics["augmented_unknown"] > 0.40
)
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7641)
