"""P069: learn a finite Enum relationship without changing declared class storage."""
from collections.abc import Iterator
from functools import partial
from pathlib import Path
from tempfile import TemporaryDirectory
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
from reporting import report
import relflow as rf
PROOF_ID = "P069"
BUDGET = 300
ROUTES = ("arch", "bay", "cove", "dune", "elm", "ford", "glen", "hill")
CHANNELS = ("web", "store")
LABELS = ("amber", "blue", "coral", "green")
def records(*, rows: int, seed: int, shuffled: bool = False) -> Iterator[dict]:
rng = np.random.default_rng(seed)
combinations = np.resize(np.arange(len(ROUTES) * len(CHANNELS)), rows)
rng.shuffle(combinations)
routes = combinations % len(ROUTES)
channels = combinations // len(ROUTES)
labels = (routes + 2 * channels) % len(LABELS)
context = rng.permutation(rows) if shuffled else np.arange(rows)
for row, source in enumerate(context):
yield {"route": ROUTES[routes[source]], "channel": CHANNELS[channels[source]], "label": LABELS[labels[row]]}
def build() -> rf.Model:
return rf.Model.xs(
batch_size=64,
route=rf.Enum(values=ROUTES, p_unavailable=0.0),
channel=rf.Enum(values=CHANNELS, p_unavailable=0.0),
label=rf.Enum(values=LABELS, p_unavailable=0.0, topk=[4], mask=True),
)
def storage(model: rf.Model) -> dict:
"""Inspect declared IDs and parameter/count row counts, excluding learned values."""
result = {}
for name in ("route", "channel", "label"):
address = rf.Address(name)
node = model.nodes[address]
result[name] = {
"labels": rf.Enum.vocabulary(model, address),
"embedding_rows": node.embedder.embeddings["content"].num_embeddings,
"count_rows": node.embedder.counters["content"].counts.numel(),
"decoder_rows": node.decoder.linears["content"].out_features if name == "label" else None,
}
return result
def prediction(model: rf.Model, rows: list[dict]) -> list[dict]:
output = model.predict(pa.Table.from_pylist(rows))["predictions"].to_pylist()
return [row["label"]["content"] for row in output]
def accuracy(output: list[dict], rows: list[dict]) -> float:
return float(np.mean([value["value"] == row["label"] for value, row in zip(output, rows, strict=True)]))
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
test = list(records(rows=2048, seed=seed + 3))
shuffled_test = list(records(rows=2048, seed=seed + 3, shuffled=True))
metrics = {"chance_accuracy": 0.25, "train_rows": 4096, "validation_rows": 512, "test_rows": 2048}
checks = {}
for shuffled in (False, True):
lit.seed_everything(seed, workers=True)
model = build()
initial = storage(model)
model.optimizer = rf.adamw(learning_rate=0.002)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=4096, seed=seed + 1, shuffled=shuffled),
validate=partial(records, rows=512, seed=seed + 2, shuffled=shuffled),
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)
model.eval()
key = "shuffled_control" if shuffled else "informative"
counts = {name: rf.Enum.counts(model, rf.Address(name)) for name in ("route", "channel", "label")}
output = prediction(model, test)
measured = accuracy(output, test)
permuted = accuracy(prediction(model, shuffled_test), shuffled_test)
metrics[key] = {
"steps": trainer.global_step,
"accuracy": measured,
"shuffled_test_accuracy": permuted,
"storage": storage(model),
"training_counts": counts,
}
checks[f"{key}: declaration order and storage remain fixed"] = storage(model) == initial
checks[f"{key}: prediction leaves all exposure counts unchanged"] = counts == {
name: rf.Enum.counts(model, rf.Address(name)) for name in counts
}
if not shuffled:
with TemporaryDirectory(prefix="relflow-enum-") as directory:
path = Path(directory) / "model.pt"
model.save(path)
loaded = rf.Model.load(path).to(model.device)
loaded.eval()
restored = prediction(loaded, test)
checks["Checkpoint preserves declared storage and counts"] = storage(loaded) == initial and counts == {
name: rf.Enum.counts(loaded, rf.Address(name)) for name in counts
}
checks["Checkpoint preserves held-out predictions"] = all(
before["value"] == after["value"] and abs(before["probability"] - after["probability"]) < 1e-5
for before, after in zip(output, restored, strict=True)
)
informative = metrics["informative"]["accuracy"]
control = metrics["shuffled_control"]["shuffled_test_accuracy"]
checks["Informative accuracy reaches 0.95"] = informative >= 0.95
checks["Shuffled-context control stays below 0.35 on independent shuffled test rows"] = control <= 0.35
checks["Shuffling test context removes learned accuracy"] = metrics["informative"]["shuffled_test_accuracy"] <= 0.35
checks["Informative accuracy exceeds the control by at least 0.55"] = informative - control >= 0.55
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=6901)
