stage progress

Author: voorhs

autointent/nodes/_optimization/_node_optimizer.py

@@ -107,11 +107,15 @@ def fit(self, context: Context) -> None:
         self._logger.info("%s node optimization is finished! saving best assets", self.node_info.node_type)
         # TODO refactor the following code (via implementing `autointent.load_module(path)` utility)
         trial_idx = context.optimization_info.get_best_trial_idx(self.node_type)
-        trial = context.optimization_info.trials.get_trial(self.node_type, trial_idx)  # type: ignore[arg-type]
-        module_type, module_kwargs = scored_modules[trial_idx]  # type: ignore[index]
-        best_module: Module = module_type(**module_kwargs)
-        best_module.load(trial.module_dump_dir)  # type: ignore[arg-type]
-        context.optimization_info.artifacts.add_artifact(self.node_type, best_module.get_artifact(context))
+        if context.is_ram_to_clear():
+            trial = context.optimization_info.trials.get_trial(self.node_type, trial_idx)  # type: ignore[arg-type]
+            module_type, module_kwargs = scored_modules[trial_idx]  # type: ignore[index]
+            best_module: Module = module_type(**module_kwargs)
+            best_module.load(trial.module_dump_dir)  # type: ignore[arg-type]
+        else:
+            best_module = context.optimization_info.modules.get(self.node_type)[trial_idx]
+        artifact = best_module.get_artifact(context)
+        context.optimization_info.artifacts.add_artifact(self.node_type, artifact)
 
     def get_module_dump_dir(self, dump_dir: Path, module_name: str, j_combination: int) -> str:
         """

tests/modules/embedding/test_logreg.py

@@ -1,12 +1,6 @@
 from autointent.modules.embedding import LogregAimedEmbedding
 
 
-def test_get_artifact_returns_correct_artifact_for_logreg():
-    module = LogregAimedEmbedding(embedder_name="sergeyzh/rubert-tiny-turbo")
-    artifact = module.get_artifact()
-    assert artifact.embedder_name == "sergeyzh/rubert-tiny-turbo"
-
-
 def test_fit_trains_model():
     module = LogregAimedEmbedding(embedder_name="sergeyzh/rubert-tiny-turbo")
 

tests/modules/embedding/test_retrieval.py

@@ -4,12 +4,6 @@
 from tests.conftest import setup_environment
 
 
-def test_get_artifact_returns_correct_artifact():
-    module = RetrievalAimedEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo")
-    artifact = module.get_artifact()
-    assert artifact.embedder_name == "sergeyzh/rubert-tiny-turbo"
-
-
 def test_dump_and_load_preserves_model_state():
     project_dir = setup_environment()
     module = RetrievalAimedEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo")

tests/nodes/conftest.py

@@ -82,7 +82,7 @@ def get_context(multilabel):
     if multilabel:
         dataset = dataset.to_multilabel()
     res.set_dataset(dataset)
-    res.configure_logging(LoggingConfig(project_dir=project_dir, keep_in_ram=False))
+    res.configure_logging(LoggingConfig(project_dir=project_dir, keep_in_ram=True))
     res.configure_vector_index(VectorIndexConfig(), EmbedderConfig(device="cpu"))
     res.configure_cross_encoder(CrossEncoderConfig())
     return res