Pyfunc Chains: load_context, predict, and Pre/Post-Processing Hooks

1.Why a chain becomes a pyfunc model

MLflow packages models in *flavors* that serving and inference platforms understand. Examples are python-function, pytorch and sklearn. A built-in flavor such as sklearn can only serve what the library's own predict returns. A generative AI application usually needs more than that. It cleans up or rewrites the user's request, calls one or more models, and then shapes the raw result into a form the caller can use. MLflow's Python function flavor, pyfunc, is built for this, because it "provides flexibility to deploy any piece of Python code or any Python model."

Databricks lists the situations where you would write custom pyfunc code:

- The model needs preprocessing before inputs can reach its predict function. - MLflow doesn't natively support the model framework. - The model's raw outputs have to be post-processed before anyone can use them. - The model has per-request branching logic. - You want to deploy fully custom code as a model.

The first and third items are what this objective is about. When a chain has a preprocessing stage, a model call and a postprocessing stage, all three travel together as one pyfunc model. A Model Serving endpoint then serves that model, which Databricks calls a *custom model*.

Sources1

2.The two methods: load_context and predict

To write a pyfunc model, you subclass mlflow.pyfunc.PythonModel. The Databricks guide names two functions to implement when you package arbitrary Python code this way.

**load_context(self, context)** runs once, when the model loads. Put anything here that only has to be loaded once for the model to work, such as model weights, a tokenizer, or helper objects built from packaged files. The reason is speed. Loading these things up front keeps the number of artifacts loaded during predict to a minimum, and that makes inference faster.

**predict(self, context, model_input)** "houses all the logic that is run every time an input request is made." Preprocessing, the call to the underlying model and postprocessing all happen here, on every request.

The rule is: one-time setup goes in load_context, and per-request work goes in predict. Suppose a chain reloads its weights or rebuilds its tokenizer inside predict. It still returns correct answers, but it repeats the slowest work on every call, which is exactly what load_context exists to avoid.

Sources1

3.Writing the pre- and post-processing hooks

The Databricks example below shows the standard shape. Preprocessing and postprocessing each live in their own helper method, and predict calls them in order: format the input, call the model, format the output. Look at what load_context does too. It loads the weights and builds a tokenizer from paths in context.artifacts. It also imports a tokenizer class from a shared preprocessing_utils module that was packaged with the model.

class CustomModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        self.model = torch.load(context.artifacts["model-weights"])
        from preprocessing_utils.my_custom_tokenizer import CustomTokenizer
        self.tokenizer = CustomTokenizer(context.artifacts["tokenizer_cache"])

    def format_inputs(self, model_input):
        # insert some code that formats your inputs
        pass

    def format_outputs(self, outputs):
        predictions = (torch.sigmoid(outputs)).data.numpy()
        return predictions

    def predict(self, context, model_input):
        model_input = self.format_inputs(model_input)
        outputs = self.model.predict(model_input)
        return self.format_outputs(outputs)

Sources1

4.Where pre- and post-processing sit in a RAG chain

Databricks uses the term *RAG chain* for the series of steps that runs at inference time when a user submits a request. Two of those steps line up directly with the hooks you just saw.

1. (Optional) User query preprocessing. The query is reshaped so it works better for searching the vector database. That can mean putting it into a template, using another model to rewrite the request, or extracting keywords. The output is a *retrieval query*. 2. Retrieval. The retrieval query is turned into an embedding with the same embedding model that embedded the document chunks. The most similar chunks come back. 3. Prompt augmentation. The retrieved context is combined with the user's query in a prompt template. 4. LLM generation. The LLM produces a response grounded in that context. 5. (Optional) Post-processing. The response "may be processed further to apply additional business logic, add citations, or otherwise refine the generated text."

In a pyfunc chain, step 1 is your input-formatting code and step 5 is your output-formatting code. Steps 2 to 4 are what predict does between the two. Guardrails can also be added anywhere along the chain, for example filtering requests or checking user permissions before data sources are accessed. Those also fit naturally inside predict.

Sources2

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.Loading weights or tokenizers inside predict is fine, because the answers come out the same.Why is that wrong?

   predict runs on every request. One-time loading belongs in load_context so that fewer artifacts are loaded during predict and inference stays fast.

   Covered in The two methods: load_context and predict