fix(flair_pipelines): Fix flair pipelines

Author: ktagowski

embeddings/pipeline/flair_classification.py

@@ -23,7 +23,10 @@
 from embeddings.transformation.flair_transformation.split_sample_corpus_transformation import (
     SampleSplitsFlairCorpusTransformation,
 )
-from embeddings.transformation.transformation import Transformation
+from embeddings.transformation.hf_transformation.class_encode_column_transformation import (
+    ClassEncodeColumnTransformation,
+)
+from embeddings.transformation.transformation import DummyTransformation, Transformation
 from embeddings.utils.json_dict_persister import JsonPersister
 
 
@@ -45,14 +48,24 @@ def __init__(
         sample_missing_splits: Optional[Tuple[Optional[float], Optional[float]]] = None,
         seed: int = 441,
         load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        encode_classes: bool = False,
     ):
         output_path = Path(output_path)
         dataset = Dataset(dataset_name, **load_dataset_kwargs if load_dataset_kwargs else {})
         data_loader = HuggingFaceDataLoader()
         transformation: Union[
-            Transformation[datasets.DatasetDict, Corpus], Transformation[Corpus, Corpus]
+            Transformation[datasets.DatasetDict, datasets.DatasetDict],
+            Transformation[datasets.DatasetDict, Corpus],
+            Transformation[Corpus, Corpus],
         ]
-        transformation = ClassificationCorpusTransformation(input_column_name, target_column_name)
+        transformation = DummyTransformation()
+        if encode_classes:
+            transformation = transformation.then(
+                ClassEncodeColumnTransformation(column=target_column_name)
+            )
+        transformation = transformation.then(
+            ClassificationCorpusTransformation(input_column_name, target_column_name)
+        )
         if sample_missing_splits:
             transformation = transformation.then(
                 SampleSplitsFlairCorpusTransformation(*sample_missing_splits, seed=seed)

embeddings/pipeline/flair_pair_classification.py

@@ -23,7 +23,10 @@
 from embeddings.transformation.flair_transformation.split_sample_corpus_transformation import (
     SampleSplitsFlairCorpusTransformation,
 )
-from embeddings.transformation.transformation import Transformation
+from embeddings.transformation.hf_transformation.class_encode_column_transformation import (
+    ClassEncodeColumnTransformation,
+)
+from embeddings.transformation.transformation import DummyTransformation, Transformation
 from embeddings.utils.json_dict_persister import JsonPersister
 
 
@@ -45,15 +48,23 @@ def __init__(
         sample_missing_splits: Optional[Tuple[Optional[float], Optional[float]]] = None,
         seed: int = 441,
         load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        encode_classes: bool = False,
     ):
         output_path = Path(output_path)
         dataset = Dataset(dataset_name, **load_dataset_kwargs if load_dataset_kwargs else {})
         data_loader = HuggingFaceDataLoader()
         transformation: Union[
-            Transformation[datasets.DatasetDict, Corpus], Transformation[Corpus, Corpus]
+            Transformation[datasets.DatasetDict, datasets.DatasetDict],
+            Transformation[datasets.DatasetDict, Corpus],
+            Transformation[Corpus, Corpus],
         ]
-        transformation = PairClassificationCorpusTransformation(
-            input_columns_names_pair, target_column_name
+        transformation = DummyTransformation()
+        if encode_classes:
+            transformation = transformation.then(
+                ClassEncodeColumnTransformation(column=target_column_name)
+            )
+        transformation = transformation.then(
+            PairClassificationCorpusTransformation(input_columns_names_pair, target_column_name)
         )
         if sample_missing_splits:
             transformation = transformation.then(

embeddings/pipeline/flair_sequence_labeling.py

@@ -23,7 +23,10 @@
 from embeddings.transformation.flair_transformation.split_sample_corpus_transformation import (
     SampleSplitsFlairCorpusTransformation,
 )
-from embeddings.transformation.transformation import Transformation
+from embeddings.transformation.hf_transformation.class_encode_column_transformation import (
+    ClassEncodeColumnTransformation,
+)
+from embeddings.transformation.transformation import DummyTransformation, Transformation
 from embeddings.utils.json_dict_persister import JsonPersister
 
 
@@ -47,14 +50,26 @@ def __init__(
         sample_missing_splits: Optional[Tuple[Optional[float], Optional[float]]] = None,
         seed: int = 441,
         load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        encode_classes: bool = True,
     ):
         output_path = Path(output_path)
         dataset = Dataset(dataset_name, **load_dataset_kwargs if load_dataset_kwargs else {})
         data_loader = HuggingFaceDataLoader()
+
         transformation: Union[
-            Transformation[datasets.DatasetDict, Corpus], Transformation[Corpus, Corpus]
+            Transformation[datasets.DatasetDict, datasets.DatasetDict],
+            Transformation[datasets.DatasetDict, Corpus],
+            Transformation[Corpus, Corpus],
         ]
-        transformation = ColumnCorpusTransformation(input_column_name, target_column_name)
+        transformation = DummyTransformation()
+        if encode_classes:
+            transformation = transformation.then(
+                ClassEncodeColumnTransformation(column=target_column_name)
+            )
+        transformation = transformation.then(
+            ColumnCorpusTransformation(input_column_name, target_column_name)
+        )
+
         if sample_missing_splits:
             transformation = transformation.then(
                 SampleSplitsFlairCorpusTransformation(*sample_missing_splits, seed=seed)