3939from embeddings .transformation .flair_transformation .split_sample_corpus_transformation import (
4040 SampleSplitsFlairCorpusTransformation ,
4141)
42+ from embeddings .transformation .hf_transformation .class_encode_column_transformation import (
43+ ClassEncodeColumnTransformation ,
44+ )
4245from embeddings .transformation .transformation import DummyTransformation
4346from embeddings .utils .flair_corpus_persister import FlairConllPersister , FlairPicklePersister
4447
@@ -62,6 +65,7 @@ class FlairPreprocessingPipeline(
6265 ignore_test_subset : bool = False
6366 seed : int = 441
6467 load_dataset_kwargs : Optional [Dict [str , Any ]] = None
68+ encode_labels : bool = False
6569
6670 def __post_init__ (self ) -> None :
6771 self .persister = self ._get_persister ()
@@ -71,7 +75,7 @@ def __post_init__(self) -> None:
7175 super (FlairPreprocessingPipeline , self ).__init__ (dataset , data_loader , transformation )
7276
7377 @abc .abstractmethod
74- def _get_base_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
78+ def _get_to_flair_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
7579 pass
7680
7781 @abc .abstractmethod
@@ -89,17 +93,25 @@ def _get_dataloader(self, dataset: Dataset) -> FLAIR_DATALOADERS:
8993
9094 def _get_dataset_transformation (
9195 self , data_loader : FLAIR_DATALOADERS
92- ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
96+ ) -> Optional [ FLAIR_DATASET_TRANSFORMATIONS_TYPE ] :
9397 if isinstance (data_loader , (ConllFlairCorpusDataLoader , PickleFlairCorpusDataLoader )):
94- return DummyTransformation ()
98+ return None
9599
96- return self ._get_base_dataset_transformation ()
100+ return self ._get_to_flair_dataset_transformation ()
97101
98102 def _get_transformations (
99103 self , data_loader : FLAIR_DATALOADERS
100104 ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
101105
102- transformation = self ._get_dataset_transformation (data_loader )
106+ transformation : FLAIR_DATASET_TRANSFORMATIONS_TYPE = DummyTransformation ()
107+ if self .encode_labels :
108+ transformation = transformation .then (
109+ ClassEncodeColumnTransformation (column = self .target_column_name )
110+ )
111+
112+ to_flair_dataset_transformation = self ._get_dataset_transformation (data_loader )
113+ if to_flair_dataset_transformation :
114+ transformation = transformation .then (to_flair_dataset_transformation )
103115
104116 if self .sample_missing_splits :
105117 transformation = transformation .then (
@@ -126,7 +138,7 @@ class FlairTextClassificationPreprocessingPipeline(FlairPreprocessingPipeline):
126138 def _get_persister (self ) -> FLAIR_PERSISTERS_TYPE :
127139 return FlairPicklePersister (self .persist_path )
128140
129- def _get_base_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
141+ def _get_to_flair_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
130142 assert isinstance (self .input_column_name , str )
131143 return ClassificationCorpusTransformation (
132144 input_column_name = self .input_column_name ,
@@ -138,7 +150,7 @@ def _get_base_dataset_transformation(self) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE
138150class FlairTextPairClassificationPreprocessingPipeline (
139151 FlairTextClassificationPreprocessingPipeline
140152):
141- def _get_base_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
153+ def _get_to_flair_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
142154 assert isinstance (self .input_column_name , (tuple , list ))
143155 return PairClassificationCorpusTransformation (
144156 input_columns_names_pair = self .input_column_name ,
@@ -151,7 +163,7 @@ class FlairSequenceLabelingPreprocessingPipeline(FlairPreprocessingPipeline):
151163 def _get_persister (self ) -> FLAIR_PERSISTERS_TYPE :
152164 return FlairConllPersister (self .persist_path )
153165
154- def _get_base_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
166+ def _get_to_flair_dataset_transformation (self ) -> FLAIR_DATASET_TRANSFORMATIONS_TYPE :
155167 assert isinstance (self .input_column_name , str )
156168 return ColumnCorpusTransformation (
157169 input_column_name = self .input_column_name ,
0 commit comments