"""P070: separate a declared Enum class from learning its relationship to context."""
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 = "P070"
BUDGET = 300
LABELS = ("amber", "blue", "coral", "green")
CENTERS = (-3.0, -1.0, 1.0, 3.0)
def records(*, rows: int, seed: int, phase: str = "initial") -> Iterator[dict]:
rng = np.random.default_rng(seed)
indices = {"initial": (0, 1, 2), "withheld": (3,), "mixed": (0, 1, 2, 3)}[phase]
classes = np.resize(np.asarray(indices), rows)
rng.shuffle(classes)
noise = rng.uniform(-0.15, 0.15, rows)
for index, deviation in zip(classes, noise, strict=True):
yield {"x": float(CENTERS[index] + deviation), "label": LABELS[index]}
def build() -> rf.Model:
return rf.Model.xs(
batch_size=64,
x=rf.Number,
label=rf.Enum(values=LABELS, p_unavailable=0.0, topk=[4], mask=True),
)
def storage(model: rf.Model) -> dict:
"""Read the fixed mapping and row counts independently from exposure values."""
node = model.nodes[rf.Address("label")]
return {
"labels": rf.Enum.vocabulary(model, rf.Address("label")),
"embedding_rows": node.embedder.embeddings["content"].num_embeddings,
"count_rows": node.embedder.counters["content"].counts.numel(),
"decoder_rows": node.decoder.linears["content"].out_features,
}
def fit(model: rf.Model, phase: str, seed: int, budget: int, accelerator: str) -> int:
model.optimizer = rf.adamw(learning_rate=0.002)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=4096, seed=seed + 1, phase=phase),
validate=partial(records, rows=512, seed=seed + 2, phase=phase),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
devices=1,
max_epochs=-1,
max_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()
return trainer.global_step
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 measure(output: list[dict], rows: list[dict]) -> dict:
probabilities = [{candidate["value"]: candidate["probability"] for candidate in row["topk"]} for row in output]
return {
"accuracy": float(np.mean([value["value"] == row["label"] for value, row in zip(output, rows, strict=True)])),
"mean_withheld_probability": float(np.mean([row[LABELS[-1]] for row in probabilities])),
"full_declared_support": all(set(row) == set(LABELS) for row in probabilities),
"probability_sum_error": float(max(abs(sum(row.values()) - 1.0) for row in probabilities)),
}
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
lit.seed_everything(seed, workers=True)
budget = BUDGET if steps is None else min(steps, BUDGET)
model = build()
initial_storage = storage(model)
seen = list(records(rows=2048, seed=seed + 3))
withheld = list(records(rows=2048, seed=seed + 4, phase="withheld"))
cold = measure(prediction(model, withheld), withheld)
checks = {
"Cold prediction exposes every declared candidate": cold["full_declared_support"],
"Cold prediction leaves exposure counts at zero": all(
value == 0 for value in rf.Enum.counts(model, "/label").values()
),
}
initial_steps = fit(model, "initial", seed, budget, accelerator)
initial_counts = rf.Enum.counts(model, "/label")
initial_seen = measure(prediction(model, seen), seen)
initial_withheld = measure(prediction(model, withheld), withheld)
checks["Initial training reaches 0.95 seen-class accuracy"] = initial_seen["accuracy"] >= 0.95
checks["Withheld class has zero initial training exposure"] = initial_counts[LABELS[-1]] == 0
checks["Initial training retains the complete fixed declaration"] = storage(model) == initial_storage
checks["Prediction of withheld records cannot teach the class"] = rf.Enum.counts(model, "/label") == initial_counts
checks["Untrained class retains output support"] = initial_withheld["full_declared_support"]
continuation_steps = fit(model, "mixed", seed + 20, budget, accelerator)
final_counts = rf.Enum.counts(model, "/label")
final_seen = measure(prediction(model, seen), seen)
final_output = prediction(model, withheld)
final_withheld = measure(final_output, withheld)
checks["Continuation reaches 0.95 withheld-class accuracy"] = final_withheld["accuracy"] >= 0.95
checks["Rehearsal retains 0.95 seen-class accuracy"] = final_seen["accuracy"] >= 0.95
checks["Continuation adds exposure to every declared class"] = all(
final_counts[label] > initial_counts[label] for label in LABELS
)
checks["Continuation leaves class IDs and all row counts unchanged"] = storage(model) == initial_storage
checks["Final evaluation preserves counts"] = rf.Enum.counts(model, "/label") == final_counts
with TemporaryDirectory(prefix="relflow-enum-continuation-") as directory:
path = Path(directory) / "continued.pt"
model.save(path)
restored = rf.Model.load(path).to(model.device)
restored.eval()
restored_output = prediction(restored, withheld)
checks["Checkpoint preserves fixed class storage and exposures"] = (
storage(restored) == initial_storage and rf.Enum.counts(restored, "/label") == final_counts
)
checks["Checkpoint preserves withheld predictions"] = all(
before["value"] == after["value"] and abs(before["probability"] - after["probability"]) < 1e-5
for before, after in zip(final_output, restored_output, strict=True)
)
lit.seed_everything(seed, workers=True)
control = build()
control_steps = fit(control, "mixed", seed, budget, accelerator)
control_seen = measure(prediction(control, seen), seen)
control_withheld = measure(prediction(control, withheld), withheld)
checks["All-classes-trained control reaches 0.95 on both populations"] = (
min(control_seen["accuracy"], control_withheld["accuracy"]) >= 0.95
)
checks["Written class probabilities sum to one"] = (
max(
value["probability_sum_error"]
for value in (
cold,
initial_seen,
initial_withheld,
final_seen,
final_withheld,
control_seen,
control_withheld,
)
)
< 1e-5
)
return {
"train_rows_per_stage": 4096,
"validation_rows_per_stage": 512,
"test_rows_per_population": 2048,
"initial_steps": initial_steps,
"continuation_steps": continuation_steps,
"control_steps": control_steps,
"uniform_four_class_accuracy": 0.25,
"initial_seen_majority_accuracy": 683 / 2048,
"cold_withheld": cold,
"initial_seen": initial_seen,
"initial_withheld": initial_withheld,
"final_seen": final_seen,
"final_withheld": final_withheld,
"control_seen": control_seen,
"control_withheld": control_withheld,
"initial_counts": initial_counts,
"final_counts": final_counts,
"storage": storage(model),
}, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7001)
