Add record_datasets_as_runs setting to trainer

configs/language/tests/offline/pubmedqa_combine.yaml

@@ -0,0 +1,90 @@
+---
+trainer:
+  class_path: eva.Trainer
+  init_args:
+    n_runs: &N_RUNS ${oc.env:N_RUNS, 1}
+    record_datasets_as_runs: true
+    default_root_dir: ${oc.env:OUTPUT_ROOT, logs/test/offline/pubmedqa_combine}
+    checkpoint_type: null
+    accelerator: ${oc.env:ACCELERATOR, auto}
+    devices: ${oc.env:NUM_DEVICES, 1}
+    callbacks:
+      - class_path: eva.callbacks.ConfigurationLogger
+      - class_path: eva.language.callbacks.writers.TextPredictionWriter
+        init_args:
+          output_dir: &PREDICTIONS_OUTPUT_DIR ${oc.env:PREDICTIONS_OUTPUT_DIR, ./predictions/test/pubmedqa_combine}
+          dataloader_idx_map:
+            0: val
+          save_format: &PREDICTIONS_SAVE_FORMAT ${oc.env:PREDICTIONS_SAVE_FORMAT, jsonl}
+          model:
+            class_path: eva.language.models.wrappers.ModelFromRegistry
+            init_args:
+              model_name: ${oc.env:MODEL_NAME, anthropic/claude-3-7-sonnet-20250219}
+              model_extra_kwargs: ${oc.env:MODEL_EXTRA_KWARGS, null}
+          overwrite: false
+model:
+  class_path: eva.language.models.OfflineLanguageModule
+  init_args:
+    metrics:
+      common:
+        - class_path: eva.metrics.MulticlassClassificationMetrics
+          init_args:
+            num_classes: 3
+            input_type: "discrete"
+    postprocess:
+      predictions_transforms:
+        - class_path: eva.language.models.postprocess.ExtractDiscreteAnswer
+          init_args:
+            answer_format: ${oc.env:ANSWER_FORMAT, json}
+            extract_kwargs:
+              mapping: {"no": 0, "yes": 1, "maybe": 2}
+              missing_limit: ${oc.env:MISSING_LIMIT, 5}
+              missing_answer: ${oc.env:MISSING_ANSWER, 2}
+data:
+  class_path: eva.DataModule
+  init_args:
+    datasets:
+      val:
+        - class_path: eva.language.datasets.TextPredictionDataset
+          init_args: &VAL_DATASET_ARGS
+            path: ${oc.env:PREDICTIONS_OUTPUT_DIR, ./predictions/test/pubmedqa_combine}/manifest.${oc.env:PREDICTIONS_SAVE_FORMAT, jsonl}
+            split: val
+        - class_path: eva.language.datasets.TextPredictionDataset
+          init_args:
+            <<: *VAL_DATASET_ARGS
+      test:
+        class_path: eva.language.datasets.TextPredictionDataset
+        init_args:
+          <<: *VAL_DATASET_ARGS
+      predict:
+        - class_path: eva.language.datasets.PubMedQA
+          init_args: &PREDICT_DATASET_ARGS
+            root: ${oc.env:DATA_ROOT, ./data/pubmedqa}
+            split: val
+            download: ${oc.env:DOWNLOAD_DATA, false}
+            prompt_template:
+              class_path: eva.language.prompts.templates.MultipleChoicePromptTemplate
+              init_args:
+                answer_format: ${oc.env:ANSWER_FORMAT, json}
+        - class_path: eva.language.datasets.PubMedQA
+          init_args:
+            <<: *PREDICT_DATASET_ARGS
+            split: test
+    dataloaders:
+      val: &DATALOADER_ARGS
+        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 256}
+        num_workers: &N_DATA_WORKERS ${oc.env:N_DATA_WORKERS, 0}
+        shuffle: false
+        collate_fn: eva.language.data.dataloaders.prediction_collate
+        pin_memory: false
+        persistent_workers: false
+        prefetch_factor: null
+      test:
+        <<: *DATALOADER_ARGS
+      predict:
+        batch_size: &PREDICT_BATCH_SIZE ${oc.env:PREDICT_BATCH_SIZE, 16}
+        num_workers: *N_DATA_WORKERS
+        collate_fn: eva.language.data.dataloaders.text_collate
+        pin_memory: false
+        persistent_workers: false
+        prefetch_factor: null

src/eva/core/interface/interface.py

@@ -41,6 +41,7 @@ def fit(
             stages=["fit", "validate", "test"],
             n_runs=trainer.n_runs,
             verbose=trainer.n_runs > 1,
+            record_datasets_as_runs=trainer.record_datasets_as_runs,
         )
 
     def predict(
@@ -116,6 +117,7 @@ def validate(
             stages=["validate"],
             n_runs=trainer.n_runs,
             verbose=trainer.n_runs > 1,
+            record_datasets_as_runs=trainer.record_datasets_as_runs,
         )
 
     def test(
@@ -135,6 +137,7 @@ def test(
             stages=["test"],
             n_runs=trainer.n_runs,
             verbose=trainer.n_runs > 1,
+            record_datasets_as_runs=trainer.record_datasets_as_runs,
         )
 
     def validate_test(
@@ -156,4 +159,5 @@ def validate_test(
             stages=["validate", "test"],
             n_runs=trainer.n_runs,
             verbose=trainer.n_runs > 1,
+            record_datasets_as_runs=trainer.record_datasets_as_runs,
         )

src/eva/core/trainers/functional.py

@@ -1,5 +1,6 @@
 """Fit session related functions."""
 
+from itertools import zip_longest
 from typing import List, Literal, Tuple
 
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
@@ -17,6 +18,7 @@ def run_evaluation_session(
     *,
     n_runs: int = 1,
     stages: List[Literal["fit", "validate", "test"]] | None = None,
+    record_datasets_as_runs: bool = False,
     verbose: bool = True,
 ) -> None:
     """Runs a downstream evaluation session out-of-place.
@@ -32,6 +34,9 @@ def run_evaluation_session(
         datamodule: The data module.
         n_runs: The number of runs to perform.
         stages: List of stages to execute. Options: "fit", "validate", "test".
+        record_datasets_as_runs: If True, when multiple validation and/or test datasets are
+            configured, each dataset output is recorded as a separate run in the session
+            summary. Otherwise, dataset outputs are logged separately within the same run.
         verbose: Whether to verbose the session metrics instead of
             those of each individual run and vice-versa.
     """
@@ -45,10 +50,20 @@ def run_evaluation_session(
             datamodule,
             run_id=run_index,
             stages=stages,
+            record_datasets_as_runs=record_datasets_as_runs,
             verbose=not verbose,
         )
         if validation_scores or test_scores:
-            recorder.update(validation_scores, test_scores)
+            if record_datasets_as_runs:
+                for val_result, test_result in zip_longest(
+                    validation_scores or [], test_scores or []
+                ):
+                    recorder.update(
+                        [val_result] if val_result is not None else None,
+                        [test_result] if test_result is not None else None,
+                    )
+            else:
+                recorder.update(validation_scores, test_scores)
     recorder.save()
 
 
@@ -59,6 +74,7 @@ def run_evaluation(
     *,
     run_id: int | None = None,
     stages: List[Literal["fit", "validate", "test"]] | None = None,
+    record_datasets_as_runs: bool = False,
     verbose: bool = True,
 ) -> Tuple[_EVALUATE_OUTPUT | None, _EVALUATE_OUTPUT | None]:
     """Runs the specified evaluation stages out-of-place.
@@ -70,6 +86,8 @@ def run_evaluation(
         run_id: The run id to be appended to the output log directory.
             If `None`, it will use the log directory of the trainer as is.
         stages: List of stages to execute. Options: "fit", "validate", "test".
+        record_datasets_as_runs: If True, evaluate each dataloader separately
+            so that metrics are computed independently per dataset.
         verbose: Whether to print the validation and test metrics
             in the end of the training.
 
@@ -92,18 +110,22 @@ def run_evaluation(
     if "fit" in stages:
         trainer.fit(model, datamodule=datamodule)
     if "validate" in stages and getattr(datamodule.datasets, "val", None) is not None:
-        validation_scores = trainer.validate(
-            model=model,
-            datamodule=datamodule,
+        validation_scores = _evaluate_stage(
+            trainer,
+            model,
+            datamodule,
+            stage="validate",
+            record_datasets_as_runs=record_datasets_as_runs,
             verbose=verbose,
-            ckpt_path=trainer.checkpoint_type,
         )
     if "test" in stages and getattr(datamodule.datasets, "test", None) is not None:
-        test_scores = trainer.test(
-            model=model,
-            datamodule=datamodule,
+        test_scores = _evaluate_stage(
+            trainer,
+            model,
+            datamodule,
+            stage="test",
+            record_datasets_as_runs=record_datasets_as_runs,
             verbose=verbose,
-            ckpt_path=trainer.checkpoint_type,
         )
     trainer.finish_logger_run(run_id)
     return validation_scores, test_scores
@@ -132,3 +154,37 @@ def infer_model(
         datamodule=datamodule,
         return_predictions=return_predictions,
     )
+
+
+def _evaluate_stage(
+    trainer: eva_trainer.Trainer,
+    model: modules.ModelModule,
+    datamodule: datamodules.DataModule,
+    *,
+    stage: Literal["validate", "test"],
+    record_datasets_as_runs: bool = False,
+    verbose: bool = True,
+) -> _EVALUATE_OUTPUT:
+    """Evaluates a validation or test stage.
+
+    When ``record_datasets_as_runs`` is enabled and multiple dataloaders exist,
+    each dataloader is evaluated separately so that metrics are computed
+    independently per dataset instead of being aggregated across all dataloaders.
+    """
+    evaluate_fn = trainer.validate if stage == "validate" else trainer.test
+    ckpt_path = trainer.checkpoint_type
+
+    if record_datasets_as_runs:
+        datamodule.setup(stage)
+        get_dls = datamodule.val_dataloader if stage == "validate" else datamodule.test_dataloader
+        dls = get_dls()
+        if isinstance(dls, list) and len(dls) > 1:
+            return [
+                metrics
+                for dl in dls
+                for metrics in evaluate_fn(
+                    model=model, dataloaders=[dl], verbose=verbose, ckpt_path=ckpt_path
+                )
+            ]
+
+    return evaluate_fn(model=model, datamodule=datamodule, verbose=verbose, ckpt_path=ckpt_path)

src/eva/core/trainers/trainer.py

@@ -33,6 +33,7 @@ def __init__(
         checkpoint_type: Literal["best", "last"] = "best",
         accelerator: str = "auto",
         devices: int = 1,
+        record_datasets_as_runs: bool = False,
         **kwargs: Any,
     ) -> None:
         """Initializes the trainer.
@@ -49,6 +50,9 @@ def __init__(
                 callback for evaluations on validation & test sets.
             accelerator: The accelerator to use for training (e.g. "cpu", "gpu").
             devices: The number of devices (GPUs) to use for training.
+            record_datasets_as_runs: If True, when multiple validation and/or test datasets are
+                configured, each dataset output is recorded as a separate run in the session
+                summary. Otherwise, dataset outputs are logged separately within the same run.
             kwargs: Kew-word arguments of ::class::`lightning.pytorch.Trainer`.
         """
         super().__init__(
@@ -61,6 +65,7 @@ def __init__(
 
         self.checkpoint_type = checkpoint_type
         self.n_runs = n_runs
+        self.record_datasets_as_runs = record_datasets_as_runs
 
         self._session_id: str = _logging.generate_session_id()
         self._log_dir: str = self.default_log_dir
@@ -150,4 +155,5 @@ def run_evaluation_session(
             datamodule=datamodule,
             n_runs=self.n_runs,
             verbose=self.n_runs > 1,
+            record_datasets_as_runs=self.record_datasets_as_runs,
         )