Re: Custom sentence transformer for indexing

mark_ott · ‎11-07-2025

To use a custom MLflow pyfunc model for sentence-transformers with preprocessing, you need to comply with the expected interface of mlflow.pyfunc.PythonModel, especially the predict method. The method signature, data handling, and serialization are key points. Below is a direct answer with practical explanation and guidelines.

Required Methods for mlflow.pyfunc.PythonModel

The only method you must implement is predict(self, context, model_input).

context: MLflow-provided info (artifacts, configs, etc.).
model_input: The input passed during inference (usually Pandas DataFrame, NumPy array, or Python native types).

Guidelines and Typical Pattern

Load everything needed in load_context, which runs once when the model is loaded by MLflow.
Accept both batch (DataFrame/array) and single-input cases in predict.
The output of predict should be directly serializable (ideally array-like or DataFrame).

Example Template

python

import mlflow.pyfunc
from sentence_transformers import SentenceTransformer
import pandas as pd

class CustomSentenceTransformerModel(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        self.model = SentenceTransformer('all-MiniLM-L6-v2')

    def preprocess(self, row):
        # Custom preprocessing - join columns, etc.
        return f"{row['field1']} {row['field2']} {row['field3']}"

    def predict(self, context, model_input):
        # Accept DataFrame, Series, or list
        # If DataFrame, apply preprocessing
        if isinstance(model_input, pd.DataFrame):
            texts = model_input.apply(self.preprocess, axis=1).tolist()
        elif isinstance(model_input, list):
            texts = [self.preprocess(x) if isinstance(x, dict) else x for x in model_input]
        else:
            texts = [str(model_input)]

        return self.model.encode(texts)

Key Points for Indexing Tables

When serving/inferencing, the input must be a DataFrame, array, or compatible structure; MLflow Model Serving expects this.
If you want to process tables, accept a DataFrame in predict, preprocess each row, and then encode.
All logic for optional preprocessing must be inside predict.

Troubleshooting the "Index creation failed" Error

The error likely means predict does not consume the input structure as expected, or the output is not serializable.
Ensure you return standard Python objects (lists, arrays, DataFrames); avoid returning custom objects or types that cannot be serialized easily.
Check that your model serving environment has all dependencies (sentence-transformers, etc.).

Final Recommendations

Implement only load_context and predict, where predict handles any preprocessing.
Return vector outputs in formats compatible with downstream tooling (usually NumPy arrays or lists).
Test your model locally first:

python

import pandas as pd data = pd.DataFrame([{"field1": "hello", "field2": 3, "field3": 4.2}]) model.predict(None, data)