« Back to API reference

neuralbench.callbacks.TestFullRetrievalMetrics

class neuralbench.callbacks.TestFullRetrievalMetrics(event_type: Literal['Word', 'Image'] = 'Word', event_field: Literal['text', 'category'] = 'text', retrieval_item: Literal['trigger', 'window'] = 'trigger', window_decimals: int = 0, retrieval_set_sizes: tuple = (None, 250), save_outputs: bool = False, logger: Any = None, eval_val: bool = False)

Accumulate predictions on entire test set before evaluating metrics.

Requires the pl.LightningModule object this callback is attached to to define a torch.ModuleDict named retrieval_metrics containing the metrics to evaluate.

retrieval_item names what one candidate of the retrieval set is:

• "trigger": whatever the trigger is labelled with, its text or its filepath. Repeats of a label collapse into one candidate carrying the mean of their targets, which is what a word-level or image-level retrieval score is built on.

• "window": the file a continuous trigger (Audio, Video) plays, paired with the window’s time within that file (the trigger’s offset plus the window’s delay from the trigger), rounded to window_decimals decimals of a second. Only windows that saw the same moment of the same stimulus collapse, e.g. two subjects watching one scene. For a single event spanning the whole recording, "trigger" averages targets that differ, and the candidate it yields is a centroid rather than anything the model could have decoded.

on_test_batch_end(trainer: Trainer, pl_module: LightningModule, outputs, batch, batch_idx, dataloader_idx=0)

Called when the test batch ends.

on_test_epoch_end(trainer, pl_module) → None

Called when the test epoch ends.

on_test_epoch_start(trainer: Trainer, pl_module: LightningModule)

Called when the test epoch begins.

on_validation_batch_end(trainer: Trainer, pl_module: LightningModule, outputs, batch, batch_idx, dataloader_idx=0)

Called when the validation batch ends.

on_validation_epoch_end(trainer, pl_module) → None

Called when the val epoch ends.

on_validation_epoch_start(trainer: Trainer, pl_module: LightningModule)

Called when the val epoch begins.

setup(trainer: Trainer, pl_module: LightningModule, stage: str)

Called when fit, validate, test, predict, or tune begins.

« Back to API reference