diff --git a/configs/language/tests/offline/pubmedqa_combine.yaml b/configs/language/tests/offline/pubmedqa_combine.yaml new file mode 100644 index 000000000..65142ab1b --- /dev/null +++ b/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 diff --git a/src/eva/core/interface/interface.py b/src/eva/core/interface/interface.py index 74abba349..6c7c62d50 100644 --- a/src/eva/core/interface/interface.py +++ b/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, ) diff --git a/src/eva/core/trainers/functional.py b/src/eva/core/trainers/functional.py index b37163eee..d1be2b65c 100644 --- a/src/eva/core/trainers/functional.py +++ b/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) diff --git a/src/eva/core/trainers/trainer.py b/src/eva/core/trainers/trainer.py index 060f1927c..afc561267 100644 --- a/src/eva/core/trainers/trainer.py +++ b/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, ) diff --git a/tests/eva/core/trainers/test_functional.py b/tests/eva/core/trainers/test_functional.py new file mode 100644 index 000000000..4a80e84b7 --- /dev/null +++ b/tests/eva/core/trainers/test_functional.py @@ -0,0 +1,191 @@ +"""Tests for trainer session functions.""" + +from unittest import mock + +from eva.core.trainers import functional + + +def test_run_evaluation_session_records_each_dataset_pair_as_run() -> None: + """Tests that each validation/test dataset pair is recorded as its own run.""" + trainer = mock.Mock(default_log_dir="logs") + model = mock.Mock() + datamodule = mock.Mock() + recorder = mock.Mock() + validation_scores = [ + {"val/MulticlassAccuracy": 0.1}, + {"val/MulticlassAccuracy": 0.2}, + ] + test_scores = [ + {"test/MulticlassAccuracy": 0.3}, + {"test/MulticlassAccuracy": 0.4}, + ] + + with ( + mock.patch.object(functional._recorder, "SessionRecorder", return_value=recorder), + mock.patch.object( + functional, + "run_evaluation", + return_value=(validation_scores, test_scores), + ), + ): + functional.run_evaluation_session( + base_trainer=trainer, + base_model=model, + datamodule=datamodule, + record_datasets_as_runs=True, + verbose=False, + ) + + assert recorder.update.call_args_list == [ + mock.call([validation_scores[0]], [test_scores[0]]), + mock.call([validation_scores[1]], [test_scores[1]]), + ] + recorder.save.assert_called_once_with() + + +def test_run_evaluation_session_records_validate_only_datasets_as_runs() -> None: + """Tests that validate-only dataset outputs are recorded without placeholder test results.""" + trainer = mock.Mock(default_log_dir="logs") + model = mock.Mock() + datamodule = mock.Mock() + recorder = mock.Mock() + validation_scores = [ + {"val/MulticlassAccuracy": 0.1}, + {"val/MulticlassAccuracy": 0.2}, + ] + + with ( + mock.patch.object(functional._recorder, "SessionRecorder", return_value=recorder), + mock.patch.object( + functional, + "run_evaluation", + return_value=(validation_scores, None), + ), + ): + functional.run_evaluation_session( + base_trainer=trainer, + base_model=model, + datamodule=datamodule, + record_datasets_as_runs=True, + verbose=False, + ) + + assert recorder.update.call_args_list == [ + mock.call([validation_scores[0]], None), + mock.call([validation_scores[1]], None), + ] + assert mock.call([validation_scores[0]], [None]) not in recorder.update.call_args_list + assert mock.call([validation_scores[1]], [None]) not in recorder.update.call_args_list + recorder.save.assert_called_once_with() + + +def test_evaluate_stage_calls_validate_per_dataloader_when_record_datasets_as_runs() -> None: + """Tests that each dataloader is evaluated separately when record_datasets_as_runs is True.""" + trainer = mock.Mock() + trainer.validate.side_effect = [ + [{"val/Accuracy": 0.8}], + [{"val/Accuracy": 0.6}], + ] + trainer.checkpoint_type = "best" + model = mock.Mock() + datamodule = mock.Mock() + dl_1, dl_2 = mock.Mock(), mock.Mock() + datamodule.val_dataloader.return_value = [dl_1, dl_2] + + result = functional._evaluate_stage( + trainer, + model, + datamodule, + stage="validate", + record_datasets_as_runs=True, + ) + + assert result == [{"val/Accuracy": 0.8}, {"val/Accuracy": 0.6}] + assert trainer.validate.call_count == 2 + trainer.validate.assert_any_call( + model=model, dataloaders=[dl_1], verbose=True, ckpt_path="best" + ) + trainer.validate.assert_any_call( + model=model, dataloaders=[dl_2], verbose=True, ckpt_path="best" + ) + datamodule.setup.assert_called_once_with("validate") + + +def test_evaluate_stage_uses_datamodule_when_single_dataloader() -> None: + """Tests that a single dataloader falls back to the standard datamodule-based call.""" + trainer = mock.Mock() + trainer.validate.return_value = [{"val/Accuracy": 0.9}] + trainer.checkpoint_type = "best" + model = mock.Mock() + datamodule = mock.Mock() + datamodule.val_dataloader.return_value = [mock.Mock()] + + result = functional._evaluate_stage( + trainer, + model, + datamodule, + stage="validate", + record_datasets_as_runs=True, + ) + + assert result == [{"val/Accuracy": 0.9}] + trainer.validate.assert_called_once_with( + model=model, datamodule=datamodule, verbose=True, ckpt_path="best" + ) + + +def test_evaluate_stage_uses_datamodule_when_record_disabled() -> None: + """Tests that record_datasets_as_runs=False uses the standard datamodule-based call.""" + trainer = mock.Mock() + trainer.validate.return_value = [{"val/Accuracy": 0.7}, {"val/Accuracy": 0.7}] + trainer.checkpoint_type = "best" + model = mock.Mock() + datamodule = mock.Mock() + + result = functional._evaluate_stage( + trainer, + model, + datamodule, + stage="validate", + record_datasets_as_runs=False, + ) + + assert result == [{"val/Accuracy": 0.7}, {"val/Accuracy": 0.7}] + trainer.validate.assert_called_once_with( + model=model, datamodule=datamodule, verbose=True, ckpt_path="best" + ) + + +def test_run_evaluation_session_keeps_dataset_results_grouped_when_recording_disabled() -> None: + """Tests that full result lists stay grouped when dataset outputs are not recorded as runs.""" + trainer = mock.Mock(default_log_dir="logs") + model = mock.Mock() + datamodule = mock.Mock() + recorder = mock.Mock() + validation_scores = [ + {"val/MulticlassAccuracy": 0.1}, + {"val/MulticlassAccuracy": 0.2}, + ] + test_scores = [ + {"test/MulticlassAccuracy": 0.3}, + {"test/MulticlassAccuracy": 0.4}, + ] + + with ( + mock.patch.object(functional._recorder, "SessionRecorder", return_value=recorder), + mock.patch.object( + functional, + "run_evaluation", + return_value=(validation_scores, test_scores), + ), + ): + functional.run_evaluation_session( + base_trainer=trainer, + base_model=model, + datamodule=datamodule, + record_datasets_as_runs=False, + verbose=False, + ) + + recorder.update.assert_called_once_with(validation_scores, test_scores) + recorder.save.assert_called_once_with() diff --git a/tests/eva/language/test_language_cli.py b/tests/eva/language/test_language_cli.py index e6d7756bb..fcd87b4b6 100644 --- a/tests/eva/language/test_language_cli.py +++ b/tests/eva/language/test_language_cli.py @@ -1,6 +1,8 @@ """Tests regarding eva's CLI commands on language datasets.""" +import json import os +import statistics import tempfile from typing import Any, Dict, List from unittest import mock @@ -18,6 +20,7 @@ [ "configs/language/pathology/online/multiple_choice/pubmedqa.yaml", "configs/language/pathology/offline/multiple_choice/pubmedqa.yaml", + "configs/language/tests/offline/pubmedqa_combine.yaml", ], ) def test_configuration_initialization(configuration_file: str, lib_path: str) -> None: @@ -84,6 +87,50 @@ def test_predict_validate_from_configuration(configuration_file: str, lib_path: ) +def test_predict_validate_records_datasets_as_runs_from_configuration(lib_path: str) -> None: + """Tests session statistics when multiple validation datasets are recorded as runs.""" + configuration_file = "configs/language/tests/offline/pubmedqa_combine.yaml" + + with tempfile.TemporaryDirectory() as temp_dir: + predictions_dir = os.path.join(temp_dir, "predictions") + logs_dir = os.path.join(temp_dir, "logs") + + with mock.patch.dict( + os.environ, + { + "N_RUNS": "1", + "BATCH_SIZE": "2", + "PREDICTIONS_OUTPUT_DIR": predictions_dir, + "OUTPUT_ROOT": logs_dir, + "MISSING_LIMIT": "0", + }, + ): + _cli.run_cli_from_main( + cli_args=[ + "predict", + "--config", + os.path.join(lib_path, configuration_file), + ] + ) + _cli.run_cli_from_main( + cli_args=[ + "validate", + "--config", + os.path.join(lib_path, configuration_file), + ] + ) + + with open(_find_results_file(logs_dir), "r") as file: + results = json.load(file) + + assert len(results["metrics"]["val"]) == 1 + metric_statistics = next(iter(results["metrics"]["val"][0].values())) + assert len(metric_statistics["values"]) == 2 + assert metric_statistics["mean"] == pytest.approx( + statistics.mean(metric_statistics["values"]) + ) + + @pytest.fixture(autouse=True) def mock_dependencies(): """Mocks external dependencies to avoid API calls and downloads.""" @@ -125,3 +172,15 @@ def _fake_prepare_data(self): def skip_dataset_validation() -> None: """Mocks the validation step of the datasets.""" datasets.PubMedQA.validate = mock.MagicMock(return_value=None) + + +def _find_results_file(output_dir: str) -> str: + """Returns the path of the generated results file.""" + result_files = [] + for root, _, files in os.walk(output_dir): + for file in files: + if file == "results.json": + result_files.append(os.path.join(root, file)) + + assert len(result_files) == 1 + return result_files[0]