fix label encoder

Author: kaylode

theseus/ml/preprocessors/encoder.py

@@ -1,6 +1,6 @@
 import json
 import os
-import os.path as osp
+import pickle
 
 from sklearn.preprocessing import LabelEncoder, OneHotEncoder, OrdinalEncoder
 
@@ -12,7 +12,9 @@
 
 
 class LabelEncode(Preprocessor):
-    def __init__(self, encoder_type="le", save_folder=None, **kwargs):
+    def __init__(
+        self, encoder_type="le", pickle_path=None, engine: str = "pandas", **kwargs
+    ):
         super().__init__(**kwargs)
 
         assert encoder_type in [
@@ -22,48 +24,110 @@ def __init__(self, encoder_type="le", save_folder=None, **kwargs):
         ], "Encoder type not supported"
 
         self.encoder_type = encoder_type
-        self.save_folder = save_folder
-        self.mapping_dict = {}
-
-        if self.encoder_type == "le":
-            self.encoder = LabelEncoder()
-        elif self.encoder_type == "onehot":
-            self.encoder = OneHotEncoder()
+        self.pickle_path = pickle_path
+        self.engine = engine
+        if self.engine == "polars":
+            import polars as pl
+
+        if self.pickle_path is not None:
+            with open(self.pickle_path, "rb") as fb:
+                config = pickle.load(fb)
+            self.column_names = config["column_names"]
+            self.encoder_type = config["encoder_type"]
+            self.engine = config["engine"]
+            self.encoders = config["encoders"]
+            self.log(f"Loaded mapping dict from {self.pickle_path}")
         else:
-            self.encoder = OrdinalEncoder()
+            self.encoders = {}
+            if self.encoder_type == "le":
+                encoder = LabelEncoder()
+            elif self.encoder_type == "onehot":
+                encoder = OneHotEncoder()
+            else:
+                encoder = OrdinalEncoder()
+
+            for column in self.column_names:
+                self.encoders[column] = encoder
 
     @classmethod
-    def from_json(cls, json_path: str):
-        return cls(json_path=json_path, encoder_type="json_mapping")
-
-    def create_mapping_dict(self, column_name):
-        le_name_mapping = dict(
-            zip(
-                self.encoder.classes_,
-                [int(i) for i in self.encoder.transform(self.encoder.classes_)],
+    def from_pickle(cls, pickle_path: str):
+        return cls(pickle_path=pickle_path)
+
+    def save_pickle(self, pickle_path: str):
+        with open(pickle_path, "wb") as fb:
+            pickle.dump(
+                {
+                    "column_names": self.column_names,
+                    "encoder_type": self.encoder_type,
+                    "engine": self.engine,
+                    "encoders": self.encoders,
+                },
+                fb,
             )
-        )
-        if self.save_folder is not None:
-            os.makedirs(self.save_folder, exist_ok=True)
-            json.dump(
-                le_name_mapping,
-                open(osp.join(self.save_folder, column_name + ".json"), "w"),
-                indent=4,
+        self.log(f"Saved encoder to {pickle_path}")
+
+    def save_json(self, json_path: str):
+        os.makedirs(os.path.dirname(json_path), exist_ok=True)
+        mapping_dict = {}
+        for column_name in self.column_names:
+            class_mapping = dict(
+                zip(
+                    self.encoders[column_name].classes_,
+                    [
+                        int(i)
+                        for i in self.encoders[column_name].transform(
+                            self.encoders[column_name].classes_
+                        )
+                    ],
+                )
             )
-        return le_name_mapping
+            mapping_dict[column_name] = class_mapping
+        with open(json_path, "w") as fb:
+            json.dump(mapping_dict, fb, indent=4)
+        self.log(f"Saved mapping dict to {json_path}")
 
     def encode_corpus(self, df):
         for column_name in self.column_names:
-            df[column_name] = self.encoder.fit_transform(df[column_name].values).copy()
-            mapping_dict = self.create_mapping_dict(column_name)
-            self.mapping_dict[column_name] = mapping_dict
+            encoder = self.encoders[column_name]
+            if self.engine == "pandas":
+                df[column_name] = encoder.fit_transform(df[column_name].values).copy()
+            elif self.engine == "polars":
+                import polars as pl
+
+                encoder.fit_transform(df[column_name].to_numpy())
+                le_name_mapping = dict(
+                    zip(
+                        encoder.classes_,
+                        [int(i) for i in encoder.transform(encoder.classes_)],
+                    )
+                )
+                df = df.with_columns(
+                    pl.col(column_name).replace_strict(
+                        le_name_mapping, return_dtype=pl.Int32, default=None
+                    )
+                )
+
         return df
 
     def encode_query(self, df):
         for column_name in self.column_names:
-            df[column_name] = self.apply(
-                df[column_name], lambda x: self.mapping_dict[column_name].get(x, -1)
-            ).copy()
+            encoder = self.encoders[column_name]
+            if self.engine == "pandas":
+                df[column_name] = encoder.transform(df[column_name].values).copy()
+            elif self.engine == "polars":
+                import polars as pl
+
+                le_name_mapping = dict(
+                    zip(
+                        encoder.classes_,
+                        [int(i) for i in encoder.transform(encoder.classes_)],
+                    )
+                )
+                df = df.with_columns(
+                    pl.col(column_name).replace_strict(
+                        le_name_mapping, return_dtype=pl.Int32, default=None
+                    )
+                )
         return df
 
     def run(self, df):
@@ -75,7 +139,7 @@ def run(self, df):
                 level=LoggerObserver.WARN,
             )
             self.column_names = [col for col, dt in df.dtypes.items() if dt == object]
-        self.encode_corpus(df)
+        df = self.encode_corpus(df)
 
         self.log(f"Label-encoded columns: {self.column_names}")
         return df

theseus/nlp/retrieval/models/tf_idf.py

@@ -14,7 +14,7 @@ def identity_tokenizer(text):
 class TFIDFEncoder(BaseRetrieval):
     def __init__(
         self,
-        min_df: int = 0,
+        min_df: int = 0.0,
         max_df: int = 1.0,
         model_path: str = None,
         ngram_range: Tuple[int] = (1, 1),