update new features (#17)

* enhance pool loader instantiation * enhance pool loader instantiation * add logger save_to_parquet method * add logger save_to_parquet method * enhance datastores -- enforce default column names * make BADGE an hybrid strategy * add warning when index is bigger than dataset * update lock file * make estimator call hooks defined on Self * rename target to labels * format clustering utils * rename target to labels * rename target to labels * format * lint * format and lint * condense pool-based steps * fix typo and lint * lint * fix seals * fix train loader when anchoral
pietrolesci · Oct 17, 2023 · 279ae46 · 279ae46
1 parent e063ede
commit 279ae46
Show file tree

Hide file tree

Showing 25 changed files with 1,529 additions and 1,525 deletions.
diff --git a/energizer/active_learning/clustering_utilities.py b/energizer/active_learning/clustering_utilities.py
@@ -2,7 +2,8 @@
 
 import numpy as np
 from numpy.random import RandomState
-from scipy.spatial.distance import cdist
+
+# from scipy.spatial.distance import cdist
 from sklearn.cluster import KMeans, kmeans_plusplus
 from sklearn.metrics import silhouette_score
 from sklearn.preprocessing import StandardScaler
@@ -81,6 +82,18 @@ def kmeans_silhouette_sampling(X: np.ndarray, num_clusters: int, rng: RandomStat
     return _kmeans(X, num_clusters, rng, True, normalize)
 
 
+def kmeans_pp_sampling(X: np.ndarray, num_clusters: int, rng: RandomState, *args, **kwargs) -> List[int]:
+    _, indices = kmeans_plusplus(X, num_clusters, random_state=rng)
+
+    unique_ids = list(set(indices.tolist()))
+
+    # can this generate duplicates?
+    if len(unique_ids) != len(indices):
+        print(f"Kmeans++ returned duplicates. {X.shape=} {num_clusters=}")
+
+    return unique_ids
+
+
 # def kmeans_pp_sampling(X: np.ndarray, num_clusters: int, rng: RandomState, normalize: bool = True) -> List[int]:
 #     """kmeans++ seeding algorithm.
 
@@ -111,15 +124,3 @@ def kmeans_silhouette_sampling(X: np.ndarray, num_clusters: int, rng: RandomStat
 #         centers_ids.append(new_center_id)
 
 #     return centers_ids
-
-
-def kmeans_pp_sampling(X: np.ndarray, num_clusters: int, rng: RandomState, *args, **kwargs) -> List[int]:
-    _, indices = kmeans_plusplus(X, num_clusters, random_state=rng)
-
-    unique_ids = list(set(indices.tolist()))
-
-    # can this generate duplicates?
-    if len(unique_ids) != len(indices):
-        print(f"Kmeans++ returned duplicates. {X.shape=} {num_clusters=}")
-
-    return unique_ids
diff --git a/energizer/active_learning/datastores/base.py b/energizer/active_learning/datastores/base.py
@@ -132,7 +132,7 @@ def pool_dataset(self, round: Optional[int] = None, with_indices: Optional[List[
         if with_indices is not None:
             mask = mask & self._train_data[SpecialKeys.ID].isin(with_indices)
         return Dataset.from_pandas(
-            self._train_data.loc[mask, [i for i in self._train_data.columns if i != InputKeys.TARGET]]
+            self._train_data.loc[mask, [i for i in self._train_data.columns if i != InputKeys.LABELS]]
         )
 
     def label(
@@ -155,11 +155,11 @@ def label(
         # train-validation split
         if validation_perc is not None:
             n_val = floor(validation_perc * len(indices)) or 1  # at least add one
-            currentdata = self._train_data.loc[mask, [SpecialKeys.ID, InputKeys.TARGET]]
+            currentdata = self._train_data.loc[mask, [SpecialKeys.ID, InputKeys.LABELS]]
             val_indices = sample(
                 indices=currentdata[SpecialKeys.ID].tolist(),
                 size=n_val,
-                labels=currentdata[InputKeys.TARGET].tolist(),
+                labels=currentdata[InputKeys.LABELS].tolist(),
                 mode=validation_sampling,
                 random_state=self._rng,
             )
@@ -180,13 +180,13 @@ def sample_from_pool(
         mask = self._pool_mask(round)
         if with_indices:
             mask = mask & self._train_data[SpecialKeys.ID].isin(with_indices)
-        data = self._train_data.loc[mask, [SpecialKeys.ID, InputKeys.TARGET]]
+        data = self._train_data.loc[mask, [SpecialKeys.ID, InputKeys.LABELS]]
 
         return sample(
             indices=data[SpecialKeys.ID].tolist(),
             size=size,
             random_state=random_state or self._rng,
-            labels=data[InputKeys.TARGET].tolist(),
+            labels=data[InputKeys.LABELS].tolist(),
             **kwargs,
         )
 
@@ -200,19 +200,19 @@ def save_labelled_dataset(self, save_dir: Union[str, Path]) -> None:
     """
 
     def _labelled_mask(self, round: Optional[int] = None) -> pd.Series:
-        mask = self._train_data[SpecialKeys.IS_LABELLED] == True
+        mask = self._train_data[SpecialKeys.IS_LABELLED] == True  # noqa: E712
         if round is not None:
             mask = mask & (self._train_data[SpecialKeys.LABELLING_ROUND] <= round)
         return mask
 
     def _train_mask(self, round: Optional[int] = None) -> pd.Series:
-        return self._labelled_mask(round) & (self._train_data[SpecialKeys.IS_VALIDATION] == False)
+        return self._labelled_mask(round) & (self._train_data[SpecialKeys.IS_VALIDATION] == False)  # noqa: E712
 
     def _validation_mask(self, round: Optional[int] = None) -> pd.Series:
-        return self._labelled_mask(round) & (self._train_data[SpecialKeys.IS_VALIDATION] == True)
+        return self._labelled_mask(round) & (self._train_data[SpecialKeys.IS_VALIDATION] == True)  # noqa: E712
 
     def _pool_mask(self, round: Optional[int] = None) -> pd.Series:
-        mask = self._train_data[SpecialKeys.IS_LABELLED] == False
+        mask = self._train_data[SpecialKeys.IS_LABELLED] == False  # noqa: E712
         if round is not None:
             mask = mask | (self._train_data[SpecialKeys.LABELLING_ROUND] > round)
         return mask
@@ -233,7 +233,8 @@ def get_pool_embeddings(self, ids: List[int]) -> np.ndarray:
 
     def get_train_embeddings(self, ids: List[int]) -> np.ndarray:
         # check all the ids are training ids
-        assert len(set(self.get_train_ids()).intersection(set(ids))) == len(ids)  # type: ignore
+        train_ids = self.get_train_ids()  # type: ignore
+        assert all(i in train_ids for i in ids), set(train_ids).difference(set(ids))
 
         # now that we are sure, let's unmask them and get the items
         self.unmask_ids_from_index(ids)

diff --git a/energizer/active_learning/datastores/classification.py b/energizer/active_learning/datastores/classification.py
@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional
 
 from datasets import Dataset
 from transformers import PreTrainedTokenizerBase
@@ -13,21 +13,20 @@ class ActivePandasDataStoreForSequenceClassification(SequenceClassificationMixin
     @classmethod
     def from_datasets(
         cls,
-        input_names: Union[str, List[str]],
-        target_name: str,
         tokenizer: PreTrainedTokenizerBase,
-        train_dataset: Dataset,
-        validation_dataset: Optional[Dataset] = None,
-        test_dataset: Optional[Dataset] = None,
         uid_name: Optional[str] = None,
         on_cpu: Optional[List[str]] = None,
         seed: Optional[int] = 42,
+        train_dataset: Optional[Dataset] = None,
+        validation_dataset: Optional[Dataset] = None,
+        test_dataset: Optional[Dataset] = None,
     ) -> Self:
-        obj = cls(seed)
+        obj = cls(seed)  # type: ignore
         obj = _from_datasets(
             obj=obj,
-            input_names=input_names,
-            target_name=target_name,
+            mandatory_input_names=cls.MANDATORY_INPUT_NAMES,
+            optional_input_names=cls.OPTIONAL_INPUT_NAMES,
+            mandatory_target_name=cls.MANDATORY_TARGET_NAME,
             tokenizer=tokenizer,
             uid_name=uid_name,
             on_cpu=on_cpu,