Scikit‑LLM et MLflow : versionner et auditer des pipelines LLM

Un exemple de travail montre comment consigner et comparer plusieurs variantes d’un pipeline scikit-learn intégrant des LLMs en local avec gpt4all, puis auditer les exécutions sous MLflow. L’expérience nommée « Scikit-LLM-Versioning » enregistre les paramètres clés, les statuts et les modèles, avec des IDs de runs distincts pour un modèle Orca Mini et une version Falcon.

Un audit des runs expose réussites et échecs sous MLflow

Les exécutions de l’expérience sont récupérées via l’API de recherche de MLflow et présentées dans un DataFrame pandas, en ne gardant que les colonnes run_id, tags.mlflow.runName, params.llm_model_file et status. L’exemple affiché comprend cinq lignes, avec des runs FINISHED et un run FAILED. Deux lignes concernent Upgraded_Falcon avec le modèle Falcon et un statut FINISHED, deux autres Baseline_Orca_Mini terminées avec Orca Mini, et une ligne Baseline_Orca_Mini en échec. L’expérience est retrouvée par son nom « Scikit-LLM-Versioning », puis les runs sont listés et sous-sélectionnés pour l’audit, le nom de run figurant dans tags.mlflow.runName, le fichier de modèle dans params.llm_model_file et le statut d’exécution dans status.

Objectif : versionnage reproductible malgré les mises à jour de backends

Le dispositif vise d’abord la reproductibilité et le versionnage des modèles lorsque les backends LLM évoluent, un cas fréquent et potentiellement complexe en production. L’association de Scikit-LLM et de MLflow permet de construire, suivre, comparer et consigner les variantes de pipelines, tout en facilitant la comparaison de plusieurs backends au sein d’une même expérience.

Environnement local et expérience MLflow configurés

L’installation repose sur « scikit-llm[gpt4all] » et mlflow, l’option gpt4all étant indiquée pour éviter des incompatibilités. Scikit-LLM est initialisé avec une clé et une organisation fictives afin d’autoriser l’exécution locale via gpt4all. Côté traçabilité, le registre de modèles MLflow s’appuie sur une base SQLite déclarée par sqlite:///mlflow.db, et l’expérience de suivi porte le nom « Scikit-LLM-Versioning ».

Jeu d’entraînement minimal pour une classification zéro-shot

Un petit corpus d’entraînement est défini pour une tâche de classification zéro-shot, composé de trois textes étiquetés. Les trois classes utilisées sont « bug », « praise » et « feedback », ce qui suffit ici à illustrer l’ajustement des pipelines et leur consignation.

Première variante : Orca Mini consignée avec cloudpickle

Un pipeline de base recourt à un classifieur zéro-shot s’appuyant sur le fichier « gpt4all::orca-mini-3k-71m-q4_0.gguf ». Intégré dans une sklearn.pipeline.Pipeline, il est exécuté dans un run MLflow nommé « Baseline_Orca_Mini », avec llm_backend fixé à « gpt4all » et llm_model_file au chemin du modèle. Le pipeline est ajusté sur le corpus et enregistré via mlflow.sklearn.log_model en « cloudpickle », présenté comme une variante du classique pickle (.pkl). L’exécution produit un identifiant de run, illustré par l’exemple 0852aaec23364725b433f09973a3d911.

Seconde variante : Falcon exécutée dans un run dédié

Une variante plus lourde cible « gpt4all::ggml-model-gpt4all-falcon-q4_0.bin ». Elle est lancée dans un run distinct « Upgraded_Falcon », avec llm_backend renseigné à « gpt4all » et llm_model_file pointant vers le modèle Falcon. Le pipeline est entraîné et archivé avec le même format de sérialisation, et génère un identifiant de run illustré par ee892572d0a641f89201c33479b98746.

Comparaison des backends et perspective de promotion en registre

Les variantes consignées peuvent être comparées grâce aux paramètres enregistrés par MLflow, permettant d’identifier la version la plus adaptée. Le processus décrit inclut la possibilité de promouvoir le pipeline le plus performant vers le registre de modèles MLflow en vue d’un déploiement. L’ensemble s’appuie sur MLflow, cadre open-source de gestion du cycle de vie des projets d’apprentissage automatique, et sur Scikit-LLM pour intégrer les LLMs aux pipelines scikit-learn.