help with categorical encoding

[g00phy]

Trying to port a kaggle housing regressor from Tensorflow to Jax but have hit a snag.

The dataset has a few categorical features(like "MSZoning", "Street","Alley") which need to hot encoded. Tensorflow has a few preprocessors like TextVectorization, StringLookup, or IntegerLookup to handle categorical features of various types.

I could build a text category encoder as shown below which would adapt on a particular categorical feature.

And finally plug the encoder in the model as shown below.

Is there anything similar to StringLookup/textvectorizer in jax/flax which could adapt on a category and can be plugged in a model?

[jakevdp]

JAX doesn't include any tools for this kind of text vectorization. I would probably use something like sklearn.preprocessing.OneHotEncoder.

For doing things more at scale, the maxtext repo has some good examples (e.g. https://github.com/google/maxtext/blob/main/getting_started/Data_Input_Pipeline.md)

[g00phy]

Trying to follow the tensorflow guides when preparing dataset .

Two options are suggested.

1. Preprocessing data before the model: Using preprocessor before the model creation and process the dataset. The preprocessor could be sklearn OneHotEncoder or tensorflow Stringlookup.
   The whole pipeline would be proprocessor creation --> preprocessing dataset ---> model creation ----> train on the processed dataset.
   Only disadvantage is preprocessing might happen on CPU as preprocessor is not part of the model.
2. Preprocessing data in the model: Using the already trained preprocessor and plug it in model creation and process the batches in the model.
   The pipeline would be proprocessor creation --> train the preprocesor on the dataset ---> model creation and plug the processor as the first layer ----> training on the raw dataset
   There might be some speedup as the pre-processing happens as the first step of model traversal.

With JAX, Scikit OneHotEncoder could be used. But, we might be taxed at the preprocessing.

Let me if there's a way to plug tensorflow stringlookup/textvectorizer with JAX model.

Again, thanks for the input.