使用MLflow跟踪LangChain实验

MLflow是一个功能强大的开源平台，用于管理机器学习生命周期中的工作流和工件。它提供了与许多流行的机器学习库的内置集成，也可以与任何库、算法或部署工具配合使用。在LangChain集成的背景下，MLflow提供以下功能：

• 实验跟踪：MLflow跟踪并存储LangChain实验的工件，包括模型、代码、提示、指标等。
• 依赖管理：MLflow自动记录模型依赖，确保开发和生产环境的一致性。
• 模型评估：MLflow提供原生功能来评估LangChain应用程序。
• 跟踪：MLflow允许您通过LangChain链、代理、检索器或其他组件可视化数据流。

注意：跟踪功能仅在MLflow版本2.14.0及以上可用。

本文将演示如何使用MLflow跟踪LangChain实验。

设置

安装MLflow Python包:

%pip install google-search-results num
%pip install mlflow -qU
%pip install langchain-openai -qU  # 可选：使用OpenAI LLM

环境设置

import os

# 如果运行MLflow Tracking Server，请设置MLflow跟踪URI
os.environ["MLFLOW_TRACKING_URI"] = "http://localhost:5000"
os.environ["OPENAI_API_KEY"] = "your-api-key"

创建MLflow实验

为了组织我们的运行和工件，我们建议创建一个专属的MLflow实验，而不是使用默认实验。

import mlflow

mlflow.set_experiment("LangChain MLflow Integration")

集成方法

方法1: 自动记录 (Autologging)

自动记录是利用MLflow提供的无缝追踪功能，推荐作为使用LangChain MLflow集成的首选。只需调用mlflow.langchain.autolog()即可开始使用。

import mlflow

mlflow.langchain.autolog(
    log_models=True,
    log_input_examples=True,
    log_model_signatures=True,
)

定义链 (Define a Chain)

from langchain.schema.output_parser import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful assistant that translates {input_language} to {output_language}."),
        ("human", "{input}"),
    ]
)
parser = StrOutputParser()

chain = prompt | llm | parser

调用链 (Invoke the Chain)

test_input = {
    "input_language": "English",
    "output_language": "German",
    "input": "I love programming.",
}

result = chain.invoke(test_input)
print(result)  # 'Ich liebe das Programmieren.'

探索MLflow Tracking UI，可以更深入地理解被记录的信息。

应用场景分析

可以通过MLflow UI分析运行的调用堆栈，了解LangChain组件的执行情况。此外，MLflow还能记录模型定义、依赖和输入输出样例，确保实验的可复现性和透明性。

实践建议

• 使用mlflow.langchain.autolog()进行自动记录以简化追踪过程。
• 为每个实验创建专属的实验空间以保持工作流的清晰结构。
• 在复杂场景中使用MLflow Callbacks进行半自动化追踪。

结束语：如果遇到问题欢迎在评论区交流。

—END—