Convert Preprocess subdirectory to polars; drop keras_preprocessing; add pytest

fuxictr/preprocess/build_dataset.py

@@ -31,9 +31,9 @@ def split_train_test(train_ddf=None, valid_ddf=None, test_ddf=None, valid_size=0
     Supports sequential (by index) or random splitting.
 
     Args:
-        train_ddf (pd.DataFrame): Full training data.
-        valid_ddf (pd.DataFrame, optional): Pre-existing validation data.
-        test_ddf (pd.DataFrame, optional): Pre-existing test data.
+        train_ddf (pl.DataFrame or pl.LazyFrame): Full training data.
+        valid_ddf (pl.DataFrame or pl.LazyFrame, optional): Pre-existing validation data.
+        test_ddf (pl.DataFrame or pl.LazyFrame, optional): Pre-existing test data.
         valid_size (int or float): Validation set size. If ``< 1``, treated as
             a fraction. Default: ``0``.
         test_size (int or float): Test set size. If ``< 1``, treated as a
@@ -43,7 +43,10 @@ def split_train_test(train_ddf=None, valid_ddf=None, test_ddf=None, valid_size=0
     Returns:
         tuple: ``(train_ddf, valid_ddf, test_ddf)``.
     """
-    num_samples = len(train_ddf)
+    if isinstance(train_ddf, pl.LazyFrame):
+        train_ddf = train_ddf.collect()
+
+    num_samples = train_ddf.height
     train_size = num_samples
     instance_IDs = np.arange(num_samples)
     if split_type == "random":
@@ -52,16 +55,16 @@ def split_train_test(train_ddf=None, valid_ddf=None, test_ddf=None, valid_size=0
         if test_size < 1:
             test_size = int(num_samples * test_size)
         train_size = train_size - test_size
-        test_ddf = train_ddf.loc[instance_IDs[train_size:], :].reset_index()
+        test_ddf = train_ddf.select(pl.all().gather(instance_IDs[train_size:]))
         instance_IDs = instance_IDs[0:train_size]
     if valid_size > 0:
         if valid_size < 1:
             valid_size = int(num_samples * valid_size)
         train_size = train_size - valid_size
-        valid_ddf = train_ddf.loc[instance_IDs[train_size:], :].reset_index()
+        valid_ddf = train_ddf.select(pl.all().gather(instance_IDs[train_size:]))
         instance_IDs = instance_IDs[0:train_size]
     if valid_size > 0 or test_size > 0:
-        train_ddf = train_ddf.loc[instance_IDs, :].reset_index()
+        train_ddf = train_ddf.select(pl.all().gather(instance_IDs))
     return train_ddf, valid_ddf, test_ddf
 
 
@@ -70,14 +73,14 @@ def transform_block(feature_encoder, df_block, filename):
 
     Args:
         feature_encoder (FeatureProcessor): Fitted feature processor.
-        df_block (pd.DataFrame): Data block to transform.
+        df_block (pl.DataFrame): Data block to transform.
         filename (str): Output filename relative to ``data_dir``.
     """
     df_block = feature_encoder.transform(df_block)
     data_path = os.path.join(feature_encoder.data_dir, filename)
     logging.info("Saving data to parquet: " + data_path)
     os.makedirs(os.path.dirname(data_path), exist_ok=True)
-    df_block.to_parquet(data_path, index=False, engine="pyarrow")
+    df_block.write_parquet(data_path)
 
 
 def transform(feature_encoder, ddf, filename, block_size=0):
@@ -90,18 +93,15 @@ def transform(feature_encoder, ddf, filename, block_size=0):
         block_size (int): Rows per block for parallel writing. ``0`` disables
             blocking. Default: ``0``.
     """
-    ddf = ddf.collect().to_pandas()
+    ddf = ddf.collect()
     if block_size > 0:
         pool = mp.Pool(mp.cpu_count() // 2)
-        block_id = 0
-        for idx in range(0, len(ddf), block_size):
-            df_block = ddf.iloc[idx:(idx + block_size)]
+        for block_id,df_block in enumerate(ddf.iter_slices(block_size)):
             pool.apply_async(
                 transform_block,
                 args=(feature_encoder, df_block,
                       '{}/part_{:05d}.parquet'.format(filename, block_id))
             )
-            block_id += 1
         pool.close()
         pool.join()
     else:

fuxictr/preprocess/feature_processor.py

@@ -183,7 +183,7 @@ def fit(self, train_ddf, min_categr_count=1, num_buckets=10, rebuild_dataset=Tru
             if col["active"]:
                 logging.info("Processing column: {}".format(col))
                 col_series = (
-                    train_ddf.select(name).collect().to_series().to_pandas() if self.rebuild_dataset
+                    train_ddf.select(name).collect().to_series() if self.rebuild_dataset
                     else None
                 )
                 if col["type"] == "meta": # e.g. set group_id in gAUC
@@ -279,7 +279,7 @@ def fit_numeric_col(self, col, col_series):
         if "normalizer" in col:
             normalizer = Normalizer(col["normalizer"])
             if self.rebuild_dataset:
-                normalizer.fit(col_series.dropna().values)
+                normalizer.fit(col_series.drop_na())
             self.processor_dict[name + "::normalizer"] = normalizer
 
     def fit_embedding_col(self, col):
@@ -351,7 +351,7 @@ def fit_categorical_col(self, col, col_series, min_categr_count=1, num_buckets=1
                 num_buckets = col.get("num_buckets", num_buckets)
                 qtf = sklearn_preprocess.QuantileTransformer(n_quantiles=num_buckets + 1)
                 if self.rebuild_dataset:
-                    qtf.fit(col_series.values)
+                    qtf.fit(col_series)
                     boundaries = qtf.quantiles_[1:-1]
                     self.processor_dict[name + "::boundaries"] = boundaries
                 self.feature_map.features[name]["vocab_size"] = num_buckets
@@ -424,35 +424,34 @@ def transform(self, ddf):
         for feature, feature_spec in self.feature_map.features.items():
             if feature in ddf.columns:
                 feature_type = feature_spec["type"]
-                col_series = ddf[feature]
+                col_series = ddf.get_column(feature)
                 if feature_type == "meta":
                     if feature + "::tokenizer" in self.processor_dict:
                         tokenizer = self.processor_dict[feature + "::tokenizer"]
-                        ddf[feature] = tokenizer.encode_meta(col_series)
+                        feature_data = tokenizer.encode_meta(col_series)
                         # Update vocab in tokenizer
                         self.processor_dict[feature + "::tokenizer"] = tokenizer
                 elif feature_type == "numeric":
                     normalizer = self.processor_dict.get(feature + "::normalizer")
                     if normalizer:
-                        ddf[feature] = normalizer.transform(col_series.values)
+                        feature_data = normalizer.transform(col_series)
                 elif feature_type == "categorical":
                     category_processor = feature_spec.get("category_processor")
                     if category_processor is None:
-                        ddf[feature] = (
-                            self.processor_dict.get(feature + "::tokenizer")
-                            .encode_category(col_series)
-                        )
+                        feature_data = (self.processor_dict.get(feature + "::tokenizer")
+                                        .encode_category(col_series))
                     elif category_processor == "numeric_bucket":
                         raise NotImplementedError
                     elif category_processor == "hash_bucket":
                         raise NotImplementedError
                 elif feature_type == "sequence":
-                    ddf[feature] = (self.processor_dict.get(feature + "::tokenizer")
+                    feature_data = (self.processor_dict.get(feature + "::tokenizer")
                                     .encode_sequence(col_series))
                 elif feature_type == "embedding":
                     continue
                 else:
                     raise NotImplementedError
+                ddf = ddf.with_columns(feature_data.alias(feature))
         return ddf
 
     def load_pickle(self, pickle_file=None):

fuxictr/preprocess/tokenizer.py

@@ -16,11 +16,11 @@
 # =========================================================================
 
 from collections import Counter
+from typing import Iterable
 import numpy as np
 import h5py
 from tqdm import tqdm
 import polars as pl
-from keras_preprocessing.sequence import pad_sequences
 from concurrent.futures import ProcessPoolExecutor, as_completed
 import multiprocessing as mp
 
@@ -69,7 +69,7 @@ def fit_on_texts(self, series):
             chunk_size = 1000000
             tasks = []
             for idx in range(0, len(series), chunk_size):
-                data_chunk = series.iloc[idx: (idx + chunk_size)]
+                data_chunk = series.slice(idx, chunk_size)
                 tasks.append(executor.submit(count_tokens, data_chunk, self._splitter))
             for future in tqdm(as_completed(tasks), total=len(tasks)):
                 chunk_word_counts, chunk_max_len = future.result()
@@ -153,48 +153,55 @@ def encode_meta(self, series):
         """Encode a meta column series to integer indices.
 
         Args:
-            series (pandas.Series): Raw meta values.
+            series (polars.Series): Raw meta values.
 
         Returns:
-            numpy.ndarray: Encoded integer values.
+            polars.Series: Encoded integer values.
         """
         word_counts = dict(series.value_counts())
         if len(self.vocab) == 0:
             self.build_vocab(word_counts)
         else: # for considering meta data in test data
             self.update_vocab(word_counts.keys())
-        series = series.map(lambda x: self.vocab.get(x, self.vocab["__OOV__"]))
-        return series.values
+
+        series = self.encode_category(series)
+        return series
 
     def encode_category(self, series):
         """Encode a categorical series to integer indices.
 
         Args:
-            series (pandas.Series): Raw categorical values.
+            series (polars.Series): Raw categorical values.
 
         Returns:
-            numpy.ndarray: Encoded integer values.
+            polars.Series: Encoded integer values.
         """
-        series = series.map(lambda x: self.vocab.get(x, self.vocab["__OOV__"]))
-        return series.values
+        vocab = {key: self.vocab[key] for key in set(series.unique()) & set(self.vocab.keys())}
+        # polars complains if vocab keys are of different type than series (ie "__PAD__" and numeric series)
+        series = series.replace_strict(vocab, default=self.vocab["__OOV__"])
+        return series
 
     def encode_sequence(self, series):
         """Encode a sequence series to padded integer arrays.
 
         Args:
-            series (pandas.Series): Raw sequence strings.
+            series (polars.Series): Raw sequence strings.
 
         Returns:
-            list: List of padded integer sequences.
+            series (polars.Series): padded integer sequences.
         """
-        series = series.map(
-            lambda text: [self.vocab.get(x, self.vocab["__OOV__"]) if x != self._na_value \
-            else self.vocab["__PAD__"] for x in text.split(self._splitter)]
+
+        series = (
+            series.str.split(self._splitter)
+            .list.eval(
+                pl.when(pl.element()!=self._na_value)
+                .then(pl.element().replace_strict(self.vocab,default=self.vocab["__OOV__"]))
+                .otherwise(self.vocab["__PAD__"]))
         )
         seqs = pad_sequences(series.to_list(), maxlen=self.max_len,
-                             value=self.vocab["__PAD__"],
-                             padding=self.padding, truncating=self.padding)
-        return seqs.tolist()
+                              value=self.vocab["__PAD__"],
+                              padding=self.padding, truncating=self.padding)
+        return seqs
 
     def load_pretrained_vocab(self, feature_dtype, pretrain_path, expand_vocab=True):
         """Load pretrained embedding keys and optionally expand vocabulary.
@@ -217,11 +224,39 @@ def load_pretrained_vocab(self, feature_dtype, pretrain_path, expand_vocab=True)
                     vocab_size += 1
 
 
+def pad_sequences(sequences: Iterable[Iterable[int]], maxlen=None,
+                  padding='pre', truncating='pre', value=0.):
+    if not isinstance(sequences,pl.Series):
+        sequences = pl.Series(sequences)
+    sequence_lengths = sequences.list.len()
+    if maxlen is None:
+        maxlen = sequence_lengths.max()
+    sequence_lengths = sequence_lengths.clip(upper_bound=maxlen)
+    if truncating == 'pre':
+        sequences = sequences.list.slice(0, maxlen)
+    elif truncating == 'post': 
+        sequences = sequences.list.slice(-maxlen)
+    else:
+        raise ValueError(f'Truncating type "{truncating}" not understood')
+    padder = pl.select(pl.repeat(value,len(sequences)).repeat_by(maxlen - sequence_lengths).alias("sequence")).to_series()
+    # convert to type
+    # test for 0 repeat
+    # sample_shape?
+    if padding == 'pre':
+        sequences = padder.list.concat(sequences)
+    elif padding == 'post':
+        sequences = sequences.list.concat(padder)
+    else: 
+        raise ValueError(f'Padding type "{padding}" not understood')
+    sequences = sequences.list.to_array(maxlen) # can be converted to 2d numpy array by .to_numpy()
+    return sequences
+
+
 def count_tokens(series, splitter=None):
     """Count token frequencies and max sequence length in a series.
 
     Args:
-        series (pandas.Series): Text data series.
+        series (polars.Series): Text data series.
         splitter (str, optional): Delimiter for splitting sequences.
 
     Returns:
@@ -230,12 +265,12 @@ def count_tokens(series, splitter=None):
     """
     max_len = 0
     if splitter is not None: # for sequence
-        series = series.map(lambda text: text.split(splitter))
-        max_len = series.str.len().max()
-        word_counts = series.explode().value_counts()
+        series = series.str.split(splitter)
+        max_len = series.list.len().max()
+        word_counts = series.list.explode().value_counts()
     else:
         word_counts = series.value_counts()
-    return dict(word_counts), max_len
+    return dict(word_counts.iter_rows()), max_len
 
 
 def load_pretrain_emb(pretrain_path, keys=["key", "value"]):
@@ -248,7 +283,7 @@ def load_pretrain_emb(pretrain_path, keys=["key", "value"]):
         keys (list): Keys to read from the file. Default: ``["key", "value"]``.
 
     Returns:
-        numpy.ndarray or tuple: Loaded embedding data.
+        numpy.ndarray if single embedding else list[numpy.ndarray]: Loaded embedding data.
 
     Raises:
         ValueError: If the file format is not supported.
@@ -262,8 +297,8 @@ def load_pretrain_emb(pretrain_path, keys=["key", "value"]):
         npz = np.load(pretrain_path)
         values = [npz[k] for k in keys]
     elif pretrain_path.endswith("parquet"):
-        df = pd.read_parquet(pretrain_path)
-        values = [df[k].values for k in keys]
+        df = pl.read_parquet(pretrain_path)
+        values = [df.get_column(k).to_numpy() for k in keys]
     else:
         raise ValueError(f"Embedding format not supported: {pretrain_path}")
     return values[0] if len(values) == 1 else values

model_zoo/LongCTR/longctr_dataloader.py

@@ -18,7 +18,7 @@
 import numpy as np
 from torch.utils.data import Dataset, DataLoader
 from torch.utils.data.dataloader import default_collate
-from keras_preprocessing.sequence import pad_sequences
+from  fuxictr.preprocess.tokenizer import pad_sequences
 import pandas as pd
 import torch
 
@@ -183,5 +183,5 @@ def padding_seqs(self, user_seqs, seq_lens):
             batch_seqs.append(seq[:l])
         max_len = min(max(seq_lens), self.max_len)
         batch_seqs = pad_sequences(batch_seqs, maxlen=max_len,
-                                   value=0, padding=self.padding, truncating=self.padding)
+                                   value=0, padding=self.padding, truncating=self.padding).to_numpy()
         return batch_seqs

requirements.txt

@@ -7,4 +7,4 @@ h5py
 tqdm
 pyarrow
 polars>=1.40.1
-keras_preprocessing
+pytest

setup.py

@@ -17,8 +17,8 @@
         exclude=["model_zoo", "tests", "data", "docs", "demo"]),
     include_package_data=True,
     python_requires=">=3.10",
-    install_requires=["keras_preprocessing", "pandas", "PyYAML>=6.0.1", "scikit-learn",
-                      "numpy", "h5py", "tqdm", "pyarrow", "polars"],
+    install_requires=["pandas", "PyYAML>=6.0.1", "scikit-learn",
+                      "numpy", "h5py", "tqdm", "pyarrow", "polars", "pytest"],
     classifiers=(
         "License :: OSI Approved :: Apache Software License",
         "Operating System :: OS Independent",