"""P060: recover an imbalanced population prior with uninformative inputs."""
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
from relflow.helpers.resize import Resize
PROOF_ID = "P060"
BUDGET = 400
def records(*, rows: int, seed: int, novel: bool = False) -> Iterator[dict]:
rng = np.random.default_rng(seed)
for entity, positive in zip(rng.integers(0, 4096, rows), rng.random(rows) < 0.9, strict=True):
yield {"x": 0.0, "entity": f"{'new' if novel else 'id'}-{entity}", "label": "yes" if positive else "no"}
def build(size: int) -> rf.Model:
model = rf.Model.xs(
batch_size=128,
x=rf.Number,
entity=rf.Category(),
label=rf.Category(mask=True),
)
# Experimental storage control only: allocation is not a schema option.
node = model.nodes["/label"]
resize = Resize()
resize.embedding(node.embedder.embeddings["content"], size)
resize.linear(node.decoder.linears["content"], size)
node.embedder.counters["content"].resize(size, resize)
resize.attribute(node.embedder.vocab, "size", size)
resize.commit()
model.optimizer = rf.adamw(learning_rate=0.002)
return model
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
oracle_nll = float(-0.9 * np.log(0.9) - 0.1 * np.log(0.1))
metrics = {"prior": 0.9, "weighted_optimum": 0.75, "oracle_nll": oracle_nll}
checks = {}
for size in (2, 128):
lit.seed_everything(seed, workers=True)
model = build(size)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=32768, seed=seed + 1),
validate=partial(records, rows=1024, 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),
logger=False,
enable_checkpointing=False,
enable_progress_bar=False,
enable_model_summary=False,
deterministic=True,
num_sanity_val_steps=0,
)
trainer.fit(model, datamodule=data)
model.eval()
counts = rf.Category.counts(model, "/label")
vocabulary = rf.Category.vocabulary(model, "/entity")
for novel in (False, True):
rows = list(records(rows=8192, seed=seed + 3, novel=novel))
predictions = model.predict(pa.Table.from_pylist(rows))["predictions"].to_pylist()
values = [row["label"]["content"] for row in predictions]
p = np.asarray(
[item["probability"] if item["value"] == "yes" else 1 - item["probability"] for item in values]
)
p = p.clip(1e-7, 1 - 1e-7)
key = f"capacity_{size}_{'novel' if novel else 'known'}"
rmse = float(np.sqrt(np.mean((p - 0.9) ** 2)))
excess = float(np.mean(-0.9 * np.log(p) - 0.1 * np.log1p(-p)) - oracle_nll)
metrics[key] = {"mean_probability": float(p.mean()), "prior_rmse": rmse, "excess_nll": excess}
checks[f"{key}: probability RMSE below 0.06"] = rmse < 0.06
checks[f"{key}: excess expected NLL below 0.02"] = excess < 0.02
checks[f"capacity_{size}: evaluation preserves vocabulary and counts"] = vocabulary == rf.Category.vocabulary(
model, "/entity"
) and counts == rf.Category.counts(model, "/label")
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7601)
