multabench 0.0.1__py3-none-any.whl

multabench/baselines/preprocessing/categorical.py

@@ -0,0 +1,24 @@
+from typing import Set, Dict
+
+from pandas import Series, DataFrame
+from sklearn.preprocessing import LabelEncoder
+
+from multabench.baselines.preprocessing.nulls import MISSING_VALUE
+
+
+def fit_categorical_encoders(x: DataFrame, categorical_features: Set[str]) -> Dict[str, LabelEncoder]:
+    return {col: fit_encode_categorical(s=x[col]) for col in categorical_features}
+
+def fit_encode_categorical(s: Series) -> LabelEncoder:
+    encoder = LabelEncoder()
+    train_values = set(s).union({MISSING_VALUE})
+    encoder.fit(list(train_values))
+    return encoder
+
+def transform_categorical_features(x: DataFrame, categorical_encoders: Dict[str, LabelEncoder]) -> DataFrame:
+    x = x.copy()
+    for col, encoder in categorical_encoders.items():
+        s = x[col].apply(lambda v: v if v in encoder.classes_ else MISSING_VALUE)
+        s = encoder.transform(s).astype(int)
+        x[col] = s
+    return x

multabench/baselines/preprocessing/dates.py

@@ -0,0 +1,29 @@
+"""Datetime columns expanded into numeric parts (copied from tabstar 1.1.15, tabstar/preprocessing/dates.py)."""
+from typing import Dict, Set
+
+import pandas as pd
+from pandas import DataFrame
+from skrub import DatetimeEncoder
+
+
+def transform_date_features(x: DataFrame, date_transformers: Dict[str, DatetimeEncoder]) -> DataFrame:
+    rows = x.shape[0]
+    for col, dt_encoder in date_transformers.items():
+        dt_df = dt_encoder.transform(x[col])
+        dt_df.index = x.index
+        x = x.drop(columns=[col])
+        x = pd.concat([x, dt_df], axis=1)
+        if x.shape[0] != rows:
+            raise ValueError(f"Row mismatch after transforming date column {col}: expected {rows}, got {x.shape}")
+    return x
+
+
+def fit_date_encoders(x: DataFrame, date_features: Set[str]) -> Dict[str, DatetimeEncoder]:
+    date_encoders = {}
+    for col in [c for c in x.columns if c in date_features]:
+        # Adds: "year", "month", "day", "hour", "total_seconds", "weekday"
+        encoder = DatetimeEncoder(add_weekday=True, add_total_seconds=True)
+        encoder.fit(x[col])
+        date_encoders[col] = encoder
+    return date_encoders
+

multabench/baselines/preprocessing/feature_types.py

@@ -0,0 +1,48 @@
+"""Feature types, read from the stored dtypes and the dataset's image column, and the dtype conversion per type."""
+from dataclasses import dataclass, field
+from typing import Optional, Set
+
+import pandas as pd
+from pandas import DataFrame
+from pandas.api.types import is_bool_dtype, is_datetime64_any_dtype, is_numeric_dtype
+
+from multabench.baselines.preprocessing.nulls import MISSING_VALUE
+
+
+@dataclass
+class FeatureTypes:
+    numerical_features: Set[str] = field(default_factory=set)
+    categorical_features: Set[str] = field(default_factory=set)
+    text_features: Set[str] = field(default_factory=set)
+    date_features: Set[str] = field(default_factory=set)
+    image_features: Set[str] = field(default_factory=set)
+
+
+def detect_feature_types(x: DataFrame, image_column: Optional[str]) -> FeatureTypes:
+    result = FeatureTypes()
+    for col in x.columns:
+        if col == image_column:
+            result.image_features.add(col)
+        elif is_datetime64_any_dtype(x[col]):
+            result.date_features.add(col)
+        elif isinstance(x[col].dtype, pd.CategoricalDtype):
+            result.categorical_features.add(col)
+        elif is_numeric_dtype(x[col]) or is_bool_dtype(x[col]):
+            result.numerical_features.add(col)
+        else:
+            result.text_features.add(col)
+    return result
+
+
+def transform_feature_types(x: DataFrame, numerical_features: set[str], image_features: set[str]) -> DataFrame:
+    new_x = {}
+    for col in x.columns:
+        if col in numerical_features:
+            new_x[col] = x[col].astype(float)
+        elif col in image_features:
+            new_x[col] = x[col]
+        else:
+            new_x[col] = x[col].astype(object).fillna(MISSING_VALUE).astype(str)
+    new_x = DataFrame(new_x, index=x.index)
+    ordered_x = new_x[x.columns]
+    return ordered_x

multabench/baselines/preprocessing/image_embeddings.py

@@ -0,0 +1,208 @@
+import numpy as np
+from PIL import Image
+from typing import Dict, Tuple, Optional, List, Any, Set
+
+import pandas as pd
+import torch
+from pandas import DataFrame, Series
+from sklearn.decomposition import PCA
+from sklearn.preprocessing import StandardScaler
+from multabench.constants import SEED
+from multabench.datasets.all_datasets import MulTaBenchDatasetID
+from transformers import AutoImageProcessor, AutoModel, DINOv3ViTImageProcessorFast, DINOv3ViTModel
+
+from multabench.dino.constants import DINOV3_SMALL
+from multabench.dino.dino_finetune import finetune_dino_with_lora
+from multabench.dino.image_loading import load_images
+from multabench.embeddings.hub import CachedEmbeddings, cached_embeddings, cached_encoder
+from multabench.preprocessing.discretize import discretize_numerical
+from multabench.utils.warnings import suppress_channel_dimension_warning
+from multabench.preprocessing.splits import split_to_val
+from multabench.utils.timing import embedding_step
+
+# Alternatives: facebook/dinov3-vits16plus-pretrain-lvd1689m, facebook/dinov3-convnext-tiny-pretrain-lvd1689m
+
+PCA_COMPONENTS = 30
+IMAGE_BATCH_SIZE = 64
+
+IMAGE_ENCODER = {"encoder": None, "processor": None}
+
+
+class _ScaledPCA:
+    """StandardScaler then PCA; exposes .transform(X) and .n_components for drop-in use. Holds the column's cached
+    embeddings when it has them."""
+    def __init__(self, scaler: StandardScaler, pca: PCA, cache: Optional[CachedEmbeddings] = None):
+        self.scaler = scaler
+        self.pca = pca
+        self.n_components = pca.n_components
+        self.cache = cache
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        return self.pca.transform(self.scaler.transform(X))
+
+
+def fit_image_encoders(
+    x: DataFrame,
+    image_features: Set[str],
+    device: torch.device,
+    image_folder: Optional[str],
+    is_cls: bool,
+    d_output: int,
+    y: Optional[Series] = None,
+    tune_dino: bool = False,
+    dino_train_kwargs: Optional[Dict[str, Any]] = None,
+    dino_model_name: str = DINOV3_SMALL,
+    dataset: Optional[MulTaBenchDatasetID] = None,
+) -> Tuple[Dict[str, Any], Optional[Any], Optional[DINOv3ViTImageProcessorFast], Dict[str, np.ndarray]]:
+    """
+    Fit PCA per image column. If tune_dino and y is provided,
+    finetune DINO with LoRA on train/val split then use the best model for encoding.
+    Otherwise columns with cached embeddings on the Hub are read from there instead of being encoded.
+    Returns (image_encoders, tuned_dino_model_or_none, tuned_processor_or_none, train_embeddings per column).
+    """
+    image_encoders: Dict[str, Any] = {}
+    train_embeddings: Dict[str, np.ndarray] = {}
+    image_features = [c for c in x.columns if c in image_features]
+    tuned_model = None
+    tuned_processor: Optional[DINOv3ViTImageProcessorFast] = None
+
+    if not image_features:
+        return image_encoders, tuned_model, tuned_processor, train_embeddings
+    if image_folder is None:
+        raise ValueError(f"Image columns {image_features} need an image_folder")
+
+    if tune_dino and y is not None:
+        if not is_cls:
+            y = discretize_numerical(y, n_bins=20)
+        x_tr, x_val, y_tr, y_val = split_to_val(x=x, y=y, is_cls=True)
+        first_col = image_features[0]
+        train_imgs = load_images(s=x_tr[first_col], image_folder=image_folder)
+        val_imgs = load_images(s=x_val[first_col], image_folder=image_folder)
+        kwargs = dino_train_kwargs or {}
+        tuned_model, tuned_processor = finetune_dino_with_lora(
+            train_images=train_imgs,
+            train_y=y_tr.values,
+            val_images=val_imgs,
+            val_y=y_val.values,
+            device=device,
+            processor=AutoImageProcessor.from_pretrained(dino_model_name, use_fast=True),
+            model_name=dino_model_name,
+            **kwargs,
+        )
+        tuned_model.to(device)
+        dino_model, img_processor = tuned_model, tuned_processor
+    caches = {} if tune_dino else _caches_per_column(dino_model_name, dataset)
+    if not tune_dino and not set(caches).issuperset(image_features):
+        dino_model, img_processor = get_image_encoder(model_name=dino_model_name)
+        dino_model.to(device)
+
+    for col in image_features:
+        cache = caches.get(col)
+        if cache is None:
+            images = load_images(s=x[col], image_folder=image_folder)
+            embeddings = encode_images_in_batches(images=images, processor=img_processor, model=dino_model)
+        else:
+            embeddings = cache.lookup(x[col].astype(str).tolist(), feature=col)
+        scaler = StandardScaler()
+        scaler.fit(embeddings)
+        scaled = scaler.transform(embeddings)
+        pca_col = PCA(n_components=PCA_COMPONENTS, random_state=SEED)
+        pca_col.fit(scaled)
+        image_encoders[col] = _ScaledPCA(scaler, pca_col, cache=cache)
+        train_embeddings[col] = embeddings
+
+    return image_encoders, tuned_model, tuned_processor, train_embeddings
+
+
+def _caches_per_column(dino_model_name: str, dataset: Optional[MulTaBenchDatasetID]) -> Dict[str, CachedEmbeddings]:
+    """The cached embeddings of each image column that has them."""
+    encoder = cached_encoder(dino_model_name, tuned=False)
+    if encoder is None or dataset is None:
+        return {}
+    cache = cached_embeddings(encoder, dataset)
+    print(f"Reading the image embeddings of {cache.meta['features']} from the cached {encoder} embeddings")
+    return {col: cache for col in cache.meta["features"]}
+
+
+def transform_image_features(
+    x: DataFrame,
+    image_transformers: Dict[str, Any],
+    device: torch.device,
+    image_folder: str,
+    dino_model: Optional[DINOv3ViTModel] = None,
+    dino_processor: Optional[DINOv3ViTImageProcessorFast] = None,
+    dino_model_name: str = DINOV3_SMALL,
+    train_embeddings: Optional[Dict[str, np.ndarray]] = None,
+) -> DataFrame:
+    # TODO: In realistic scenarios where we tune the model repeatedly, the efficient way would be to fit the encoder once,
+    # cache the embeddings, and then apply PCA on the cached embeddings over and over for every train split in every run.
+    if not image_transformers:
+        return x
+    train_embeddings = train_embeddings or {}
+    for image_col, image_pca in image_transformers.items():
+        assert image_col in x.columns, f"Image column {image_col} not found in DataFrame"
+        s = x[image_col]
+        embeddings = train_embeddings.get(image_col)
+        if embeddings is None and image_pca.cache is not None:
+            embeddings = image_pca.cache.lookup(s.astype(str).tolist(), feature=image_col)
+        if embeddings is None:
+            if dino_model is None or dino_processor is None:
+                dino_model, dino_processor = get_image_encoder(model_name=dino_model_name)
+            dino_model.to(device)
+            embeddings = image_urls_to_embeddings(s=s, image_folder=image_folder, processor=dino_processor, model=dino_model)
+        assert len(embeddings) == len(x)
+        pca_vec = image_pca.transform(embeddings).astype(np.float32)
+        pca_cols = [f"{image_col}_img_pca_{i}" for i in range(pca_vec.shape[1])]
+        pca_df = pd.DataFrame(pca_vec, index=s.index, columns=pca_cols)
+        x = x.drop(columns=[image_col])
+        x = pd.concat([x, pca_df], axis=1)
+    return x
+
+
+@embedding_step
+def get_image_encoder(model_name: str) -> Tuple[DINOv3ViTModel, DINOv3ViTImageProcessorFast]:
+    if IMAGE_ENCODER["encoder"] is None:
+        model = AutoModel.from_pretrained(model_name)
+        model.eval()
+        IMAGE_ENCODER["encoder"] = model
+        IMAGE_ENCODER["processor"] = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
+    dino_model = IMAGE_ENCODER["encoder"]
+    img_processor = IMAGE_ENCODER["processor"]
+    if not isinstance(dino_model, DINOv3ViTModel):
+        raise TypeError(f"Expected DINOv3ViTModel, got {type(dino_model)}")
+    if not isinstance(img_processor, (DINOv3ViTImageProcessorFast)):
+        raise TypeError(f"Expected BitImageProcessor, got {type(img_processor)}")
+    return dino_model, img_processor
+
+
+def image_urls_to_embeddings(
+    s: Series,
+    image_folder: str,
+    processor: DINOv3ViTImageProcessorFast,
+    model: DINOv3ViTModel,
+) -> np.ndarray:
+    images = load_images(s=s, image_folder=image_folder)
+    embeddings = encode_images_in_batches(images=images, processor=processor, model=model)
+    return embeddings
+
+
+@embedding_step
+def encode_images_in_batches(
+    images: List[Image.Image],
+    processor: DINOv3ViTImageProcessorFast,
+    model: DINOv3ViTModel,
+) -> np.ndarray:
+    embeddings = []
+    for i in range(0, len(images), IMAGE_BATCH_SIZE):
+        batch_images = images[i:i + IMAGE_BATCH_SIZE]
+        with suppress_channel_dimension_warning():
+            inputs = processor(images=batch_images, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            outputs = model(**inputs)
+        last_hidden = outputs["last_hidden_state"] if isinstance(outputs, dict) else outputs.last_hidden_state
+        batch_embeddings = last_hidden[:, 0, :].cpu().numpy()
+        embeddings.append(batch_embeddings)
+    embeddings = np.concatenate(embeddings, axis=0)
+    d_model = model.config.hidden_size
+    assert embeddings.shape == (len(images), d_model), f"Expected ({len(images)}, {d_model}), got {embeddings.shape}"
+    return embeddings

multabench/baselines/preprocessing/nulls.py

@@ -0,0 +1,10 @@
+"""Missing values (copied from tabstar 1.1.15, tabstar/preprocessing/nulls.py)."""
+from pandas import Series
+
+MISSING_VALUE = "Unknown Value"
+
+
+def raise_if_null_target(y: Series):
+    y_missing = y.isnull().sum()
+    if y_missing > 0:
+        raise ValueError(f"Target variable {y.name} has {y_missing} null values, please handle them before training.")

multabench/baselines/preprocessing/numerical.py

@@ -0,0 +1,13 @@
+from typing import Set, Dict
+
+from pandas import DataFrame
+
+
+def fit_numerical_median(x: DataFrame, numerical_features: Set[str]) -> Dict[str, float]:
+    return {col: x[col].median() for col in numerical_features}
+
+def transform_numerical_features(x: DataFrame, numerical_medians: Dict[str, float]) -> DataFrame:
+    x = x.copy()
+    for col, train_median in numerical_medians.items():
+        x[col] = x[col].fillna(train_median)
+    return x

multabench/baselines/preprocessing/target.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+import numpy as np
+from pandas import Series
+from sklearn.preprocessing import LabelEncoder
+
+
+def fit_preprocess_y(y: Series, is_cls: bool) -> Optional[LabelEncoder]:
+    if not is_cls:
+        return None
+    label_encoder = LabelEncoder()
+    label_encoder.fit(y)
+    return label_encoder
+
+
+def transform_preprocess_y(y: Series | np.ndarray, scaler: Optional[LabelEncoder]) -> Series:
+    if scaler is None:
+        return y
+    assert isinstance(scaler, LabelEncoder), f"Scaler must be LabelEncoder if not None, but got {type(scaler)}"
+    y = y.copy()
+    return Series(scaler.transform(y), name=y.name, index=y.index)

multabench/baselines/preprocessing/text_embeddings.py

@@ -0,0 +1,271 @@
+import os
+
+from sklearn.decomposition import PCA
+
+from multabench.constants import SEED
+
+from multabench.preprocessing.discretize import discretize_numerical
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Suppresses warning, avoids deadlock
+from typing import Dict, Set, Optional, Any, Tuple
+
+import numpy as np
+import pandas as pd
+from pandas import DataFrame, Series
+import torch
+
+from multabench.datasets.all_datasets import MulTaBenchDatasetID
+from multabench.e5.constants import E5_SMALL_V2, TF_IDF
+from multabench.e5.e5_finetune import encode_texts_with_e5, get_vanilla_e5
+from multabench.embeddings.hub import CachedEmbeddings, cached_embeddings, cached_encoder, e5_passages
+from multabench.utils.timing import embedding_step
+
+PCA_COMPONENTS = 30
+
+
+class SkrubColumnEncoder:
+    """Per-column text encoder using skrub.StringEncoder (TF-IDF + TruncatedSVD). CPU-only, no tokenizer needed."""
+
+    def __init__(self, string_encoder: Any, col_name: str, n_components: int):
+        self.string_encoder = string_encoder
+        self.col_name = col_name
+        self.n_components = n_components
+        self.encoder = self  # so wrapper.encoder.transform(X) delegates to self.transform(X)
+
+    @embedding_step
+    def encode_texts(self, texts: list[str], device) -> np.ndarray:
+        """Encode texts with skrub StringEncoder (device ignored — CPU-only)."""
+        result = self.string_encoder.transform(pd.Series(texts, dtype=str))
+        return np.asarray(result)
+
+    def embed(self, x: DataFrame, device) -> np.ndarray:
+        return self.encode_texts(x[self.col_name].astype(str).fillna("").tolist(), device)
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        return X  # encode_texts already returns final (N, n_components) array
+
+
+class E5ColumnEncoder:
+    """Per-column text encoder: holds E5 model, tokenizer (processor), and PCA (or identity). Uses passage: col_name: col_val format."""
+
+    def __init__(self, model: Any, tokenizer: Any, encoder: Any, col_name: str,
+                 cache: Optional[CachedEmbeddings] = None):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.encoder = encoder
+        self.col_name = col_name
+        self.n_components = encoder.n_components
+        self.cache = cache
+
+    def encode_texts(self, texts: list[str], device: torch.device) -> np.ndarray:
+        """Encode texts with this column's E5 model and tokenizer (passage: col_name: col_val)."""
+        return encode_texts_with_e5(
+            texts=texts,
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device=device,
+            col_name=self.col_name,
+        )
+
+    def embed(self, x: DataFrame, device: torch.device) -> np.ndarray:
+        """The embeddings of x's texts, looked up in the cache when there is one."""
+        texts = x[self.col_name].astype(str).fillna("").tolist()
+        if self.cache is None:
+            return self.encode_texts(texts, device)
+        return self.cache.lookup(e5_passages(self.col_name, texts), feature=self.col_name)
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        return self.encoder.transform(X)
+
+
+@embedding_step
+def fit_text_encoders_skrub(
+    x: DataFrame,
+    text_features_list: list[str],
+) -> Tuple[Dict[str, SkrubColumnEncoder], Dict[str, np.ndarray]]:
+    """Fit one SkrubColumnEncoder per column using skrub.StringEncoder (TF-IDF + TruncatedSVD)."""
+    from skrub import StringEncoder
+    text_encoders: Dict[str, SkrubColumnEncoder] = {}
+    for col in text_features_list:
+        col_series = x[col].astype(str).fillna("")
+        print(f"Fitting SkrubStringEncoder for column {col} with n_components={PCA_COMPONENTS} for {len(col_series)} texts")
+        string_enc = StringEncoder(n_components=PCA_COMPONENTS, random_state=SEED)
+        string_enc.fit(col_series)
+        text_encoders[str(col)] = SkrubColumnEncoder(
+            string_encoder=string_enc,
+            col_name=str(col),
+            n_components=PCA_COMPONENTS,
+        )
+    return text_encoders, {}
+
+
+def fit_text_encoders_vanilla(
+    x: DataFrame,
+    text_features_list: list[str],
+    device: torch.device,
+    e5_model_name: str = E5_SMALL_V2,
+    dataset: Optional[MulTaBenchDatasetID] = None,
+) -> Tuple[Dict[str, E5ColumnEncoder], Dict[str, np.ndarray]]:
+    """Fit one E5ColumnEncoder per column using shared vanilla E5 + PCA per column. Uses passage: col_name: col_val format.
+    Columns with cached embeddings on the Hub are read from there instead of being encoded."""
+    caches = _caches_per_column(e5_model_name, dataset)
+    model, tokenizer = None, None
+    if not set(caches).issuperset(text_features_list):
+        model, tokenizer = get_vanilla_e5(device, model_name=e5_model_name)
+    text_encoders: Dict[str, E5ColumnEncoder] = {}
+    train_embeddings: Dict[str, np.ndarray] = {}
+    for col in text_features_list:
+        pca = PCA(n_components=PCA_COMPONENTS, random_state=SEED)
+        wrapper = E5ColumnEncoder(model=model, tokenizer=tokenizer, encoder=pca, col_name=str(col), cache=caches.get(col))
+        col_embeddings = wrapper.embed(x, device)
+        pca.fit(col_embeddings)
+        text_encoders[str(col)] = wrapper
+        train_embeddings[str(col)] = col_embeddings
+    return text_encoders, train_embeddings
+
+
+def _caches_per_column(e5_model_name: str, dataset: Optional[MulTaBenchDatasetID]) -> Dict[str, CachedEmbeddings]:
+    """The cached embeddings of each text column that has them."""
+    encoder = cached_encoder(e5_model_name, tuned=False)
+    if encoder is None or dataset is None:
+        return {}
+    cache = cached_embeddings(encoder, dataset)
+    print(f"Reading the text embeddings of {cache.meta['features']} from the cached {encoder} embeddings")
+    return {col: cache for col in cache.meta["features"]}
+
+
+def fit_text_encoders_tuned(
+    x: DataFrame,
+    text_features_list: list[str],
+    device: torch.device,
+    y: Series,
+    e5_train_kwargs: Dict[str, Any],
+    is_cls: bool,
+    d_output: int,
+    e5_model_name: str = E5_SMALL_V2,
+) -> Tuple[Dict[str, E5ColumnEncoder], Dict[str, np.ndarray]]:
+    """Fit a single E5 model for all text columns with passage: col_name: col_val format. Each column gets an E5ColumnEncoder sharing the same tuned model."""
+    from transformers import AutoTokenizer
+
+    from multabench.e5.e5_finetune import finetune_e5_with_lora
+    from multabench.preprocessing.splits import split_to_val
+
+    if not is_cls:
+        y = discretize_numerical(y, n_bins=20)
+    x_tr, x_val, y_tr, y_val = split_to_val(x=x, y=y, is_cls=True)
+    kwargs = e5_train_kwargs or {}
+
+    # Build combined train/val: for each row, for each text col, create (col_name: col_val, label)
+    train_texts: list[str] = []
+    train_y: list = []
+    for idx in x_tr.index:
+        row_y = y_tr.loc[idx]
+        for col in text_features_list:
+            val = str(x_tr.loc[idx, col]).strip() if pd.notna(x_tr.loc[idx, col]) else ""
+            train_texts.append(f"{col}: {val}")
+            train_y.append(row_y)
+
+    val_texts: list[str] = []
+    val_y: list = []
+    for idx in x_val.index:
+        row_y = y_val.loc[idx]
+        for col in text_features_list:
+            val = str(x_val.loc[idx, col]).strip() if pd.notna(x_val.loc[idx, col]) else ""
+            val_texts.append(f"{col}: {val}")
+            val_y.append(row_y)
+
+    train_y_arr = np.array(train_y)
+    val_y_arr = np.array(val_y)
+
+    tokenizer = AutoTokenizer.from_pretrained(e5_model_name)
+    print(f"Finetuning single E5 ({e5_model_name}) for {len(text_features_list)} columns with {len(train_texts)} train examples (passage: col_name: col_val)")
+    tuned_model, tuned_tokenizer = finetune_e5_with_lora(
+        train_texts=train_texts,
+        train_y=train_y_arr,
+        val_texts=val_texts,
+        val_y=val_y_arr,
+        device=device,
+        tokenizer=tokenizer,
+        model_name=e5_model_name,
+        **kwargs,
+    )
+    tuned_model.to(device)
+
+    text_encoders: Dict[str, E5ColumnEncoder] = {}
+    train_embeddings: Dict[str, np.ndarray] = {}
+    for col in text_features_list:
+        texts = x[col].astype(str).fillna("").tolist()
+        col_embeddings = encode_texts_with_e5(texts=texts, model=tuned_model, tokenizer=tuned_tokenizer, device=device, col_name=str(col))
+        encoder = PCA(n_components=PCA_COMPONENTS, random_state=SEED)
+        encoder.fit(col_embeddings)
+        text_encoders[col] = E5ColumnEncoder(model=tuned_model, tokenizer=tuned_tokenizer, encoder=encoder, col_name=str(col))
+        train_embeddings[col] = col_embeddings
+    return text_encoders, train_embeddings
+
+
+def fit_text_encoders(
+    x: DataFrame,
+    text_features: Set[str],
+    device: torch.device,
+    y: Optional[Series] = None,
+    tune_e5: bool = False,
+    e5_train_kwargs: Optional[Dict[str, Any]] = None,
+    is_cls: bool = True,
+    d_output: int = 2,
+    e5_model_name: str = E5_SMALL_V2,
+    dataset: Optional[MulTaBenchDatasetID] = None,
+) -> Tuple[Dict[str, E5ColumnEncoder], Dict[str, np.ndarray]]:
+    """
+    Fit one E5 model per text column (or vanilla E5 shared across columns when not tuning).
+    Each column gets an E5ColumnEncoder wrapper holding model, tokenizer, and PCA.
+    Returns text_encoders mapping column -> E5ColumnEncoder, and the training rows' embeddings per column.
+    """
+    text_features_list = sorted(text_features)
+    if not text_features_list:
+        return {}, {}
+    if e5_model_name == TF_IDF:
+        return fit_text_encoders_skrub(
+            x=x,
+            text_features_list=text_features_list,
+        )
+    if tune_e5 and y is not None:
+        return fit_text_encoders_tuned(
+            x=x,
+            text_features_list=text_features_list,
+            device=device,
+            y=y,
+            e5_train_kwargs=e5_train_kwargs or {},
+            is_cls=is_cls,
+            d_output=d_output,
+            e5_model_name=e5_model_name,
+        )
+    return fit_text_encoders_vanilla(
+        x=x,
+        text_features_list=text_features_list,
+        device=device,
+        e5_model_name=e5_model_name,
+        dataset=dataset,
+    )
+
+
+def transform_text_features(
+    x: DataFrame,
+    text_encoders: Dict[str, E5ColumnEncoder],
+    device: torch.device,
+    train_embeddings: Optional[Dict[str, np.ndarray]] = None,
+) -> DataFrame:
+    train_embeddings = train_embeddings or {}
+    for text_col, wrapper in text_encoders.items():
+        embeddings = train_embeddings.get(text_col)
+        if embeddings is None:
+            embeddings = wrapper.embed(x, device)
+        assert len(embeddings) == len(x)
+        n_components = wrapper.n_components
+        pca_vec = wrapper.encoder.transform(embeddings)
+        pca_cols = [f"{text_col}_txt_pca_{i}" for i in range(n_components)]
+        pca_df = pd.DataFrame(pca_vec, index=x.index, columns=pca_cols)
+        cols_before = len(x.columns)
+        x = x.drop(columns=[text_col])
+        x = pd.concat([x, pca_df], axis=1)
+        assert len(x.columns) == cols_before + n_components - 1
+    return x