Scikit-LLM and MLflow: Versioning and Auditing LLM Pipelines

⚡

Key Takeaways

1Two scikit-learn pipelines integrating LLMs are logged under MLflow, one with Orca Mini and the other with Falcon

2The "Scikit-LLM-Versioning" experiment records parameters, statuses, and models, and exposes successful runs and one failure

3Variants are compared via MLflow, with the option to promote to the model registry

💡Why it matters — This setup ensures traceability and reproducibility of LLM pipelines during backend changes, facilitating their management and deployment.

An example of work demonstrates how to log and compare multiple variants of a scikit-learn pipeline integrating LLMs locally with gpt4all, and then audit the executions under MLflow. The experiment named "Scikit-LLM-Versioning" records key parameters, statuses, and models, with distinct run IDs for an Orca Mini model and a Falcon version.

An audit of runs reveals successes and failures under MLflow

The executions of the experiment are retrieved via the MLflow search API and presented in a pandas DataFrame, keeping only the columns run_id, tags.mlflow.runName, params.llm_model_file, and status. The displayed example includes five rows, with runs marked as FINISHED and one run marked as FAILED. Two rows pertain to Upgraded_Falcon with the Falcon model and a FINISHED status, two others are Baseline_Orca_Mini completed with Orca Mini, and one row is Baseline_Orca_Mini that failed. The experiment is found by its name "Scikit-LLM-Versioning," after which the runs are listed and selected for auditing, with the run name appearing in tags.mlflow.runName, the model file in params.llm_model_file, and the execution status in status.

Objective: reproducible versioning despite backend updates

The setup primarily aims for reproducibility and versioning of models when LLM backends evolve, a common and potentially complex case in production. The combination of Scikit-LLM and MLflow allows for building, tracking, comparing, and logging pipeline variants, while facilitating the comparison of multiple backends within the same experiment.

Local environment and MLflow experiment configured

The installation relies on "scikit-llm[gpt4all]" and mlflow, with the gpt4all option specified to avoid incompatibilities. Scikit-LLM is initialized with a fictitious key and organization to allow local execution via gpt4all. For traceability, the MLflow model registry relies on an SQLite database declared as sqlite:///mlflow.db, and the tracking experiment is named "Scikit-LLM-Versioning."

Minimal training set for zero-shot classification

A small training corpus is defined for a zero-shot classification task, consisting of three labeled texts. The three classes used are "bug," "praise," and "feedback," which is sufficient here to illustrate the adjustment of pipelines and their logging.

First variant: Orca Mini logged with cloudpickle

A basic pipeline uses a zero-shot classifier relying on the file "gpt4all::orca-mini-3k-71m-q4_0.gguf." Integrated into a sklearn.pipeline.Pipeline, it is executed in an MLflow run named "Baseline_Orca_Mini," with llm_backend set to "gpt4all" and llm_model_file pointing to the model path. The pipeline is trained on the corpus and logged via mlflow.sklearn.log_model in "cloudpickle," presented as a variant of the classic pickle (.pkl). The execution produces a run ID, illustrated by the example 0852aaec23364725b433f09973a3d911.

Second variant: Falcon executed in a dedicated run

A heavier variant targets "gpt4all::ggml-model-gpt4all-falcon-q4_0.bin." It is launched in a separate run "Upgraded_Falcon," with llm_backend specified as "gpt4all" and llm_model_file pointing to the Falcon model. The pipeline is trained and archived with the same serialization format, generating a run ID illustrated by ee892572d0a641f89201c33479b98746.

Comparison of backends and perspective of promotion to registry

The logged variants can be compared thanks to the parameters recorded by MLflow, allowing for the identification of the most suitable version. The described process includes the possibility of promoting the best-performing pipeline to the MLflow model registry for deployment. The entire setup relies on MLflow, an open-source framework for managing the lifecycle of machine learning projects, and on Scikit-LLM to integrate LLMs into scikit-learn pipelines.