0

I have a custom model class that calls mlflow.llm.log_predictions in the predict method like so:

class Model:
    ...

    def predict(input) -> Output:
        ...
        mlflow.llm.log_predictions(...)
        ...

I'm using that predict method in two different contexts:

  1. After training, before serializing the model - in order to do cross-validation.
  2. After serializing - to serve "production" predictions.

In the first context (cross-val), I wrap all the code in with mlflow.start_run(): in order to log data into MLFlow.

In the second context (serving), I'd like to suppress all the MLFlow calls and not log anything.

Moving mlflow.llm.log_predictions out of predict won't work for me, as the LLM I'm using is a submodel in the overall model, so I don't exactly have access to the LLM inputs/outputs outside of predict.

0 Answers0