def fit(
*,
dateparts: Sequence[str],
classes: int,
train: Callable[[], Iterator[dict]],
validate: Callable[[], Iterator[dict]],
epochs: int,
seed: int,
steps: int | None,
accelerator: str,
) -> rf.Model:
"""Train on the selected calendar coordinates with the remaining schema fixed."""
lit.seed_everything(seed, workers=True)
model = rf.Model(
name="calendar",
d_model=64,
n_layers=3,
n_heads=4,
batch_size=128,
observed_at=rf.DateParts(dateparts=list(dateparts)),
target=rf.Category(mask=True, size=classes, p_unavailable=0.0),
)
model.optimizer = lambda module: torch.optim.AdamW(module.parameters(), lr=3e-3)
data = rf.SyntheticDataModule(model=model, train=train, validate=validate, seed=seed)
trainer = lit.Trainer(
accelerator=accelerator,
max_epochs=epochs,
max_steps=steps if steps is not None else -1,
logger=False,
enable_progress_bar=False,
enable_model_summary=False,
enable_checkpointing=False,
deterministic=True,
num_sanity_val_steps=0,
)
trainer.fit(model=model, datamodule=data)
return model
def accuracy(model: rf.Model, records: Callable[[], Iterator[dict]], accelerator: str) -> float:
"""Evaluate held-out labels without updating the Category vocabulary."""
data = rf.SyntheticDataModule(model=model, test=records)
trainer = lit.Trainer(
accelerator=accelerator,
logger=False,
enable_progress_bar=False,
enable_model_summary=False,
enable_checkpointing=False,
deterministic=True,
)
metrics = trainer.test(model=model, datamodule=data, verbose=False)[0]
return float(metrics["calendar.target/test.accuracy.content"])
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
train = partial(records, years=tuple(range(1901, 1936)))
validate = partial(records, years=tuple(range(1951, 1986)))
test = partial(records, years=tuple(range(2001, 2036)))
day_only = fit(
dateparts=("day_of_year",),
classes=7,
train=train,
validate=validate,
epochs=14,
seed=seed,
steps=steps,
accelerator=accelerator,
)
ambiguous_accuracy = accuracy(day_only, test, accelerator)
identified = fit(
dateparts=("day_of_year", "day_of_week"),
classes=7,
train=train,
validate=validate,
epochs=14,
seed=seed,
steps=steps,
accelerator=accelerator,
)
identified_accuracy = accuracy(identified, test, accelerator)
gap = identified_accuracy - ambiguous_accuracy
return {
"day_only_accuracy": ambiguous_accuracy,
"with_weekday_accuracy": identified_accuracy,
"accuracy_gap": gap,
}, {
"Day-of-year accuracy is at most 0.20": ambiguous_accuracy <= 0.20,
"Visible weekday accuracy reaches 0.95": identified_accuracy >= 0.95,
"Visible weekday improves accuracy by at least 0.70": gap >= 0.70,
}