Feature/wandb logging

embeddings/data/datamodule.py

@@ -28,7 +28,7 @@
 class BaseDataModule(abc.ABC, pl.LightningDataModule, Generic[Data]):
     dataset: Data
 
-    def __init__(self) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         # ignoring the type to avoid calling to untyped function "__init__" in typed context error
         # caused by pl.LightningDataModule __init__ method not being typed
         super().__init__()  # type: ignore
@@ -64,9 +64,8 @@ def __init__(
         seed: int = 441,
         **kwargs: Any,
     ) -> None:
-        super().__init__()
-        self.tokenizer_name_or_path = tokenizer_name_or_path
         self.dataset_name_or_path = dataset_name_or_path
+        self.tokenizer_name_or_path = tokenizer_name_or_path
         self.target_field = target_field
         self.max_seq_length = max_seq_length
         self.train_batch_size = train_batch_size
@@ -82,6 +81,10 @@ def __init__(
         )
         self.load_dataset_kwargs = load_dataset_kwargs if load_dataset_kwargs else {}
         self.seed = seed
+        dataset_info = self.load_dataset()["train"].info
+        super().__init__(
+            dataset_info=dataset_info, dataset_version=dataset_info.version.version_str
+        )
 
     @abc.abstractmethod
     def prepare_labels(self) -> None:

embeddings/model/base_model.py

@@ -1,4 +1,4 @@
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 from embeddings.embedding.embedding import Embedding
 from embeddings.model.model import Model
@@ -16,6 +16,6 @@ def __init__(
         self.embedding = embedding
         self.task = task
 
-    def execute(self, data: Input) -> Output:
+    def execute(self, data: Input, **kwargs: Any) -> Output:
         embedded = self.embedding.embed(data)
         return self.task.fit_predict(embedded)

embeddings/model/flair_model.py

@@ -21,7 +21,7 @@ def __init__(
         self.task = task
         self.predict_subset = predict_subset
 
-    def execute(self, data: Corpus) -> Dict[str, nptyping.NDArray[Any]]:
+    def execute(self, data: Corpus, **kwargs: Any) -> Dict[str, nptyping.NDArray[Any]]:
         self.task.build_task_model(
             embedding=self.embedding, y_dictionary=self.task.make_y_dictionary(data)
         )

embeddings/model/lightning_model.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from numpy import typing as nptyping
 
@@ -18,6 +18,8 @@ def __init__(
         self.task = task
         self.predict_subset = predict_subset
 
-    def execute(self, data: HuggingFaceDataModule) -> Dict[str, nptyping.NDArray[Any]]:
+    def execute(
+        self, data: HuggingFaceDataModule, run_name: Optional[str] = None, **kwargs: Any
+    ) -> Dict[str, nptyping.NDArray[Any]]:
         self.task.build_task_model()
-        return self.task.fit_predict(data, self.predict_subset)
+        return self.task.fit_predict(data, self.predict_subset, run_name=run_name)

embeddings/model/lightning_module/lightning_module.py

@@ -6,6 +6,7 @@
 import torch
 from numpy import typing as nptyping
 from pytorch_lightning.utilities.types import STEP_OUTPUT
+from torch.nn.functional import softmax
 from torch.optim import AdamW, Optimizer
 from torch.utils.data import DataLoader
 from torchmetrics import MetricCollection
@@ -57,21 +58,25 @@ def validation_step(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:
     def test_step(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:
         pass
 
-    def predict_step(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:
+    def predict_step(self, *args: Any, **kwargs: Any) -> Optional[Tuple[STEP_OUTPUT, STEP_OUTPUT]]:
         batch, batch_idx = args
         loss, logits, preds = self.shared_step(**batch)
-        return preds
+        return logits, preds
 
     def predict(
         self, dataloader: DataLoader[HuggingFaceDataset]
     ) -> Dict[str, nptyping.NDArray[Any]]:
         assert self.trainer is not None
-        predictions = self.trainer.predict(dataloaders=dataloader, return_predictions=True)
+        logits_predictions = self.trainer.predict(
+            dataloaders=dataloader, return_predictions=True, ckpt_path="best"
+        )
+        logits, predictions = zip(*logits_predictions)
+        probabilities = softmax(torch.cat(logits), dim=1).numpy()
         predictions = torch.cat(predictions).numpy()
-        assert isinstance(predictions, np.ndarray)
         ground_truth = torch.cat([x["labels"] for x in dataloader]).numpy()
-        assert isinstance(ground_truth, np.ndarray)
-        return {"y_pred": predictions, "y_true": ground_truth}
+        result = {"y_pred": predictions, "y_true": ground_truth, "y_probabilities": probabilities}
+        assert all(isinstance(x, np.ndarray) for x in result.values())
+        return result
 
     def configure_metrics(self) -> None:
         if self.metrics is None:

embeddings/model/lightning_module/sequence_labeling.py

@@ -67,7 +67,7 @@ def test_step(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:
         loss, logits, preds = self.shared_step(**batch)
         if -1 not in labels:
             self.test_metrics(
-                preds[labels != self.IGNORE_INDEX], labels[labels != self.IGNORE_INDEX]
+                preds[labels != self.ignore_index], labels[labels != self.ignore_index]
             )
             self.log("test/Loss", loss, on_epoch=True)
         else:

embeddings/model/model.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 Input = TypeVar("Input")
 Output = TypeVar("Output")
@@ -10,5 +10,5 @@ def __init__(self) -> None:
         pass
 
     @abc.abstractmethod
-    def execute(self, data: Input) -> Output:
+    def execute(self, data: Input, **kwargs: Any) -> Output:
         pass

embeddings/model/sklearn_model.py

@@ -19,7 +19,7 @@ def __init__(
         self.task = task
         self.predict_subset = predict_subset
 
-    def execute(self, data: Dict[str, Any]) -> Dict[str, Any]:
+    def execute(self, data: Dict[str, Any], **kwargs: Any) -> Dict[str, Any]:
         self.embedding.fit(data["train"]["x"])
         self.task.build_task_model(self.embedding)
         return self.task.fit_predict(data, self.predict_subset)

embeddings/pipeline/evaluation_pipeline.py

@@ -46,7 +46,7 @@ def __init__(
         self.model = model
         self.evaluator = evaluator
 
-    def run(self) -> EvaluationResult:
+    def run(self, **kwargs: Any) -> EvaluationResult:
         loaded_data = self.data_loader.load(self.dataset)
         model_result = self.model.execute(loaded_data)
         return self.evaluator.evaluate(model_result)

embeddings/pipeline/flair_hps_pipeline.py

@@ -2,10 +2,11 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Dict, Generic, Optional, Tuple
+from typing import Any, Dict, Generic, Optional, Tuple
 
 import datasets
 from flair.data import Corpus
+from numpy import typing as nptyping
 
 from embeddings.data.io import T_path
 from embeddings.evaluator.sequence_labeling_evaluator import (
@@ -109,6 +110,8 @@ class OptimizedFlairClassificationPipeline(
         str,
         datasets.DatasetDict,
         Corpus,
+        Dict[str, nptyping.NDArray[Any]],
+        Dict[str, Any],
     ],
     AbstractOptimizedFlairClassificationPipeline,
     _OptimizedFlairPipelineBase[FlairTextClassificationConfigSpace],
@@ -164,7 +167,7 @@ def _get_metadata(self, parameters: SampledParameters) -> FlairClassificationPip
         return metadata
 
     def _get_evaluation_metadata(
-        self, parameters: SampledParameters
+        self, parameters: SampledParameters, **kwargs: Any
     ) -> FlairClassificationEvaluationPipelineMetadata:
         (
             embedding_name,
@@ -200,6 +203,8 @@ class OptimizedFlairPairClassificationPipeline(
         str,
         datasets.DatasetDict,
         Corpus,
+        Dict[str, nptyping.NDArray[Any]],
+        Dict[str, Any],
     ],
     AbstractOptimizedFlairClassificationPipeline,
     _OptimizedFlairPairClassificationPipelineBase[FlairTextClassificationConfigSpace],
@@ -257,7 +262,7 @@ def _get_metadata(
         return metadata
 
     def _get_evaluation_metadata(
-        self, parameters: SampledParameters
+        self, parameters: SampledParameters, **kwargs: Any
     ) -> FlairClassificationEvaluationPipelineMetadata:
         (
             embedding_name,
@@ -293,6 +298,8 @@ class OptimizedFlairSequenceLabelingPipeline(
         str,
         datasets.DatasetDict,
         Corpus,
+        Dict[str, nptyping.NDArray[Any]],
+        Dict[str, Any],
     ],
     AbstractHuggingFaceOptimizedPipeline[FlairSequenceLabelingConfigSpace],
     _OptimizedFlairPipelineDefaultsBase,
@@ -379,7 +386,7 @@ def _get_metadata(self, parameters: SampledParameters) -> FlairSequenceLabelingP
         return metadata
 
     def _get_evaluation_metadata(
-        self, parameters: SampledParameters
+        self, parameters: SampledParameters, **kwargs: Any
     ) -> FlairSequenceLabelingEvaluationPipelineMetadata:
         (
             embedding_name,

embeddings/pipeline/hps_pipeline.py

@@ -1,5 +1,6 @@
 import abc
 import logging
+import os
 from abc import ABC
 from dataclasses import dataclass, field
 from tempfile import TemporaryDirectory
@@ -18,7 +19,7 @@
 from embeddings.pipeline.preprocessing_pipeline import PreprocessingPipeline
 from embeddings.pipeline.standard_pipeline import LoaderResult, ModelResult, TransformationResult
 from embeddings.utils.hps_persister import HPSResultsPersister
-from embeddings.utils.utils import PrimitiveTypes
+from embeddings.utils.utils import PrimitiveTypes, standardize_name
 
 EvaluationResult = TypeVar("EvaluationResult", bound=Dict[str, Dict[str, PrimitiveTypes]])
 
@@ -28,7 +29,7 @@ def __init__(self) -> None:
         pass
 
     @abc.abstractmethod
-    def run(self) -> Tuple[pd.DataFrame, Metadata]:
+    def run(self, **kwargs: Any) -> Tuple[pd.DataFrame, Metadata]:
         pass
 
     def persisting(
@@ -46,15 +47,24 @@ def __init__(
             best_params_path=best_params_path, log_path=log_path
         )
 
-    def run(self) -> Tuple[pd.DataFrame, Metadata]:
-        result = self.base_pipeline.run()
+    def run(self, **kwargs: Any) -> Tuple[pd.DataFrame, Metadata]:
+        result = self.base_pipeline.run(**kwargs)
         self.persister.persist(result)
         return result
 
 
 class OptunaPipeline(
     OptimizedPipeline[Metadata],
-    Generic[ConfigSpace, Metadata, EvaluationMetadata, Data, LoaderResult, TransformationResult],
+    Generic[
+        ConfigSpace,
+        Metadata,
+        EvaluationMetadata,
+        Data,
+        LoaderResult,
+        TransformationResult,
+        ModelResult,
+        EvaluationResult,
+    ],
 ):
     def __init__(
         self,
@@ -64,7 +74,7 @@ def __init__(
         ],
         evaluation_pipeline: Union[
             Type[ModelEvaluationPipeline[Data, LoaderResult, ModelResult, EvaluationResult]],
-            Type[LightningPipeline[Data, ModelResult, EvaluationResult]],
+            Type[LightningPipeline[TransformationResult, ModelResult, EvaluationResult]],
         ],
         pruner: optuna.pruners.BasePruner,
         sampler: optuna.samplers.BaseSampler,
@@ -88,7 +98,9 @@ def _get_metadata(self, parameters: SampledParameters) -> Metadata:
         pass
 
     @abc.abstractmethod
-    def _get_evaluation_metadata(self, parameters: SampledParameters) -> EvaluationMetadata:
+    def _get_evaluation_metadata(
+        self, parameters: SampledParameters, **kwargs: Any
+    ) -> EvaluationMetadata:
         pass
 
     def get_best_paramaters(self, study: Study) -> Metadata:
@@ -99,18 +111,17 @@ def get_best_paramaters(self, study: Study) -> Metadata:
 
     def run(
         self,
+        run_name: Optional[str] = None,
+        catch: Tuple[Type[Exception], ...] = (Exception,),
+        **kwargs: Any,
     ) -> Tuple[pd.DataFrame, Metadata]:
         self._pre_run_hook()
         if self.preprocessing_pipeline is not None:
             self.preprocessing_pipeline.run()
         study: Study = optuna.create_study(
-            direction="maximize",
-            sampler=self.sampler,
-            pruner=self.pruner,
-        )
-        study.optimize(
-            self.objective, n_trials=self.n_trials, show_progress_bar=True, catch=(Exception,)
+            direction="maximize", sampler=self.sampler, pruner=self.pruner, study_name=run_name
         )
+        study.optimize(self.objective, n_trials=self.n_trials, show_progress_bar=True, catch=catch)
 
         if isinstance(self.dataset_path, TemporaryDirectory):
             self.dataset_path.cleanup()
@@ -120,15 +131,25 @@ def run(
         return study.trials_dataframe(), metadata
 
     def objective(self, trial: optuna.trial.Trial) -> float:
+        trial_name = standardize_name(f"study_{trial.study.study_name}_trial_{trial.number}")
         parameters = self.config_space.sample_parameters(trial=trial)
         parsed_params = self.config_space.parse_parameters(parameters)
-        args = self._get_evaluation_metadata(parsed_params)
-        pipeline = self.evaluation_pipeline(**args)
-        results = pipeline.run()
+        kwargs = self._get_evaluation_metadata(parsed_params, trial_name=trial_name)
+        os.makedirs(kwargs["output_path"], exist_ok=True)
+        pipeline = self._get_evaluation_pipeline(**kwargs)
+        results = pipeline.run(run_name=trial_name)
         metric = results[self.metric_name][self.metric_key]
         assert isinstance(metric, float)
         return metric
 
+    def _get_evaluation_pipeline(
+        self, **kwargs: Any
+    ) -> Union[
+        ModelEvaluationPipeline[Data, LoaderResult, ModelResult, EvaluationResult],
+        LightningPipeline[TransformationResult, ModelResult, EvaluationResult],
+    ]:
+        return self.evaluation_pipeline(**kwargs)
+
     def _pre_run_hook(self) -> None:
         logging.getLogger("optuna").setLevel(logging.WARNING)