Como versionar experimentos de LLM com Scikit-LLM e MLflow

Versionar, comparar e registrar pipelines de aprendizado de máquina que integram LLMs pode virar um pesadelo rápido — cada troca de backend de modelo é uma chance de perder reprodutibilidade. Este tutorial mostra como usar duas ferramentas para domar esse ciclo: o Scikit-LLM (que transforma modelos de linguagem em estimadores no estilo scikit-learn) e o MLflow (framework open-source para gerenciar o ciclo de vida de projetos de ML).

O que você vai fazer

Você vai construir, registrar, comparar e promover pipelines de classificação zero-shot que rodam com LLMs locais, acompanhando cada versão no MLflow e registrando a vencedora no Model Registry.

✅ O que você ganha

• Reprodutibilidade entre trocas de backend de LLM — um processo frequente na prática.
• Um histórico auditável de execuções, incluindo as que falharam.
• Um Model Registry limpo, sem poluição de tentativas ruins.
• Execução local via gpt4all, sem depender de APIs pagas.

⚠️ O que você não ganha

• Não substitui o monitoramento de modelos em produção (deriva de dados, latência).
• O exemplo usa modelos locais pequenos — não reflete a acurácia de LLMs de fronteira.
• Você ainda precisa definir suas próprias métricas de desempenho para escolher “o melhor”.

Passo a passo

1. Instalação e configuração inicial

Instale com o extra que evita problemas de compatibilidade:

pip install "scikit-llm[gpt4all]" mlflow

Depois configure o Scikit-LLM com credenciais dummy (exigidas para execução local via gpt4all) e aponte o MLflow para um backend de banco de dados, que é o que sustenta o Model Registry:

SKLLMConfig.set_openai_key("local-execution-key")
SKLLMConfig.set_openai_org("local-execution-org")
mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_experiment("Scikit-LLM-Versioning")

2. Registrando o pipeline de base

O pipeline de base treina um classificador zero-shot com um LLM leve (Orca Mini). O bloco with mlflow.start_run() registra o backend do LLM e o arquivo do modelo como parâmetros — isso é o que garante a reprodutibilidade. O formato de serialização cloudpickle é usado para contornar a checagem estrita de tipos do skops:

with mlflow.start_run(run_name="Baseline_Orca_Mini") as run_v1:
    mlflow.log_param("llm_backend", "gpt4all")
    mlflow.log_param("llm_model_file", LLM_V1)
    pipeline_v1.fit(X_train, y_train)
    mlflow.sklearn.log_model(pipeline_v1, "model", serialization_format="cloudpickle")

3. Registrando o pipeline “upgrade”

O segundo pipeline troca o Orca Mini por um modelo maior (Falcon). Tudo igual, mas em um run separado nomeado Upgraded_Falcon, gerando um novo ID de execução. É assim que você demonstra a troca de backend do MLflow.

4. Auditando e comparando versões

O mlflow.search_runs() extrai o experimento completo para um DataFrame do Pandas, permitindo ver lado a lado o ID do run, o nome, o arquivo do LLM e o status. O audit mostra não só runs FINISHED, mas também tentativas FAILED — úteis para entender o que deu errado.

5. Promovendo para o Model Registry

Para publicar a vencedora, encontre o run do Upgraded_Falcon e registre o modelo:

best_run_id = runs_df[runs_df['tags.mlflow.runName'] == 'Upgraded_Falcon'].iloc[0]['run_id']
model_uri = f"runs:/{best_run_id}/model"
mlflow.register_model(model_uri=model_uri, name="Production_ZeroShot_Classifier")

Para selecionar automaticamente o melhor por métrica, em vez de hardcode, ordene por acurácia:

best_runs_df = mlflow.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["metrics.accuracy DESC"])

O porquê do fluxo em duas etapas

O padrão “log primeiro, registre depois” evita que o Model Registry oficial vire um depósito de tentativas fracassadas. A tabela de tracking guarda os rascunhos; só o vencedor é publicado no registro para implantação. É uma separação simples, mas é a diferença entre um registro utilizável e um lixo digital.

FAQ

Preciso de GPU? Não — os exemplos usam modelos gpt4all quantizados que rodam em CPU.

Posso usar uma API em vez de modelo local? Sim, o Scikit-LLM também suporta backends como OpenAI; o padrão de versionamento do MLflow é o mesmo.

Por que cloudpickle? Para contornar a checagem estrita de tipos do skops, que pode rejeitar objetos com componentes de LLM.

O que o Model Registry me dá? Um nome de modelo versionado (Production_ZeroShot_Classifier) que você pode promover entre estágios, de staging a produção.

E se um run falhar? Ele aparece no audit como FAILED — você pode investigar sem poluir o registro, já que só registra quem você escolher.

O coração do tutorial é a disciplina: separar o rascunho (tracking) do que entra em produção (registry) e registrar os parâmetros do backend do LLM a cada execução. Com Scikit-LLM e MLflow, a troca de modelo deixa de ser um evento arriscado e vira uma entrada versionada no histórico.