diff --git a/autointent/context/optimization_info/_data_models.py b/autointent/context/optimization_info/_data_models.py index d3afcfdb8..942eec213 100644 --- a/autointent/context/optimization_info/_data_models.py +++ b/autointent/context/optimization_info/_data_models.py @@ -130,6 +130,7 @@ class Trial(BaseModel): """Value of the evaluation metric for this trial.""" module_dump_dir: str | None """Directory where the module is dumped.""" + metrics: dict[str, float] class Trials(BaseModel): diff --git a/autointent/context/optimization_info/_optimization_info.py b/autointent/context/optimization_info/_optimization_info.py index 27c1c0207..fad2256f3 100644 --- a/autointent/context/optimization_info/_optimization_info.py +++ b/autointent/context/optimization_info/_optimization_info.py @@ -75,6 +75,7 @@ def log_module_optimization( module_params: dict[str, Any], metric_value: float, metric_name: str, + metrics: dict[str, float], artifact: Artifact, module_dump_dir: str | None, module: "Module | None" = None, @@ -97,6 +98,7 @@ def log_module_optimization( metric_value=metric_value, module_params=module_params, module_dump_dir=module_dump_dir, + metrics=metrics, ) self.trials.add_trial(node_type, trial) self._logger.debug("module %s fitted and saved to optimization info", module_name, extra=trial.model_dump()) diff --git a/autointent/nodes/_optimization/_node_optimizer.py b/autointent/nodes/_optimization/_node_optimizer.py index ca8b9981e..5d3e79012 100644 --- a/autointent/nodes/_optimization/_node_optimizer.py +++ b/autointent/nodes/_optimization/_node_optimizer.py @@ -132,6 +132,7 @@ def objective( config, target_metric, self.target_metric, + all_metrics, module.get_assets(), # retriever name / scores / predictions module_dump_dir, module=module if not context.is_ram_to_clear() else None,