跳转到内容
LlamaIndex Python Documentation
TypeScript

MLflow 追踪与 LlamaIndex 端到端集成

欢迎来到这个关于LlamaIndex与MLflow集成的交互式教程。本教程提供关于LlamaIndex和MLflow核心功能的实践学习体验。

mlflow-tracing 下载此笔记本

为何将LlamaIndex与MLflow结合使用？

Section titled “Why use LlamaIndex with MLflow?”

LlamaIndex与MLflow的集成为开发和管理LlamaIndex应用程序提供了无缝体验：

  • MLflow 追踪是一款强大的可观测性工具，用于监控和调试 LlamaIndex 模型内部运行情况，帮助您快速识别潜在瓶颈或问题。
  • MLflow 实验 允许您在 MLflow 中跟踪您的索引/引擎/工作流，并管理构成您 LlamaIndex 项目的众多动态组件，例如提示词、大语言模型、工具、全局配置等。
  • MLflow 模型 将您的 LlamaIndex 应用程序与其所有依赖版本、输入输出接口及其他关键元数据打包在一起。
  • MLflow 评估 助您高效评估 LlamaIndex 应用程序的性能，确保稳健的性能分析并实现快速迭代。

你将学到什么

Section titled “What you will learn”

在本教程结束时，你将拥有：

  • 在LlamaIndex中创建了一个MVP VectorStoreIndex。
  • 使用索引作为查询引擎进行推理，并通过 MLflow Tracing 进行检查。
  • 已将索引记录到 MLflow 实验中。
  • 探索了 MLflow 用户界面，了解 MLflow 模型如何打包您的 LlamaIndex 应用程序。

这些基础知识将帮助您熟悉在MLflow中使用LlamaIndex的基本流程。如果您想了解更多关于高级用例（例如工具调用智能体）的集成内容，请参阅此高级教程。

设置

Section titled “Setup”
  1. 安装 MLflow 和 LlamaIndex：
%pip install mlflow>=2.18 llama-index>=0.10.44 -q

  1. 打开一个单独的终端并运行 mlflow ui --port 5000 来启动 MLflow UI（如果您尚未这样做）。如果您在云环境中运行此笔记本，请参考如何运行教程指南了解 MLflow 的不同设置方式。

  2. 创建一个MLflow实验并将笔记本连接到该实验

import mlflow


mlflow.set_experiment("llama-index-tutorial")
mlflow.set_tracking_uri(
    "http://localhost:5000"
)  # Or your remote tracking server URI
  1. 将 OpenAI API 密钥设置为环境变量。如果您使用不同的 LLM 提供商，请设置相应的环境变量。
import os
from getpass import getpass


os.environ["OPENAI_API_KEY"] = getpass("Enter your OpenAI API key: ")

启用 MLflow 追踪

Section titled “Enable MLflow Tracing”

只需一行代码即可为LlamaIndex启用MLflow追踪功能。

mlflow.llama_index.autolog()

创建索引

Section titled “Create an Index”

向量存储索引是LlamaIndex的核心组件之一。它们包含已摄取文档块的嵌入向量（有时也包含文档块本身）。在LlamaIndex中，这些向量可以通过不同的引擎类型用于推理任务。

  1. 查询引擎：执行直接查询，根据用户问题检索相关信息。适用于获取简明答案或匹配特定查询的文档，类似于搜索引擎。

  2. 聊天引擎：：用于处理需要跨多次交互保持上下文和历史记录的对话式AI任务。适用于需要维护对话上下文的交互式应用，如客户支持机器人或虚拟助手。

from llama_index.core import Document, VectorStoreIndex
from llama_index.core.llms import ChatMessage


# Create an index with a single dummy document
llama_index_example_document = Document.example()
index = VectorStoreIndex.from_documents([llama_index_example_document])

查询索引

Section titled “Query the Index”

让我们使用这个索引通过查询引擎执行推理。

query_response = index.as_query_engine().query("What is llama_index?")
print(query_response)

除了打印出的响应外，您还应在输出单元格中看到MLflow追踪界面。这提供了查询引擎执行流程的详细而直观的可视化，帮助您理解内部工作原理并调试可能出现的任何问题。

这次让我们使用聊天引擎进行另一个查询，看看执行流程的差异。

chat_response = index.as_chat_engine().chat(
    "What is llama_index?",
    chat_history=[
        ChatMessage(role="system", content="You are an expert on RAG!")
    ],
)
print(chat_response)

如跟踪记录所示，主要区别在于查询引擎执行的是静态工作流（RAG），而聊天引擎使用智能体工作流来动态从索引中获取必要上下文。

您还可以在MLflow用户界面中查看记录的追踪信息，方法是导航到您之前创建的实验并选择Trace标签页。如果您不想在输出单元格中显示追踪信息而仅将其记录在MLflow中，请在笔记本中运行mlflow.tracing.disable_notebook_display()。

使用MLflow保存索引

Section titled “Save the Index with MLflow”

以下代码使用MLflow记录一个LlamaIndex模型，跟踪其参数和一个示例输入，同时使用唯一的model_uri进行注册。这确保了在开发、测试和生产环境中一致且可复现的模型管理，并简化了部署和共享过程。

关键参数：

  • engine_type: 定义pyfunc和spark_udf推理类型
  • input_example: 定义输入签名并通过预测推断输出签名
  • registered_model_name: 定义MLflow模型注册表中模型的名称
with mlflow.start_run() as run:
    model_info = mlflow.llama_index.log_model(
        index,
        artifact_path="llama_index",
        engine_type="query",
        input_example="hi",
        registered_model_name="my_llama_index_vector_store",
    )
    model_uri = model_info.model_uri
    print(f"Model identifier for loading: {model_uri}")

加载索引并执行推理

Section titled “Load the Index Back and Perform Inference”

以下代码展示了使用已加载模型可执行的三种核心推理类型。

  1. 通过LlamaIndex加载并执行推理： 此方法使用mlflow.llama_index.load_model加载模型，并执行直接查询、对话或检索。当您希望充分利用底层llama索引对象的全部功能时，这是理想选择。
  2. 通过 MLflow PyFunc 加载并执行推理： 该方法使用 mlflow.pyfunc.load_model 加载模型，支持以通用 PyFunc 格式进行模型预测，并在记录时指定引擎类型。该方法适用于使用 mlflow.evaluate 评估模型或部署模型进行服务。
  3. 通过 MLflow Spark UDF 加载并执行推理： 该方法使用 mlflow.pyfunc.spark_udf 将模型加载为 Spark UDF，便于在 Spark DataFrame 中对大型数据集进行分布式推理。它非常适合处理大规模数据处理，并且与 PyFunc 推理类似，仅支持在记录时定义的引擎类型。
print("\n------------- Inference via Llama Index   -------------")
index = mlflow.llama_index.load_model(model_uri)
query_response = index.as_query_engine().query("hi")
print(query_response)


print("\n------------- Inference via MLflow PyFunc -------------")
index = mlflow.pyfunc.load_model(model_uri)
query_response = index.predict("hi")
print(query_response)
# Optional: Spark UDF inference
show_spark_udf_inference = False
if show_spark_udf_inference:
    print("\n------------- Inference via MLflow Spark UDF -------------")
    from pyspark.sql import SparkSession


    spark = SparkSession.builder.getOrCreate()


    udf = mlflow.pyfunc.spark_udf(spark, model_uri, result_type="string")
    df = spark.createDataFrame([("hi",), ("hello",)], ["text"])
    df.withColumn("response", udf("text")).toPandas()

探索 MLflow 实验界面

Section titled “Explore the MLflow Experiment UI”

最后，让我们探索 MLflow 的用户界面，查看我们迄今为止记录的内容。您可以通过在浏览器中打开 http://localhost:5000 来访问该界面，或者运行以下单元格在笔记本内显示它。

# Directly renders MLflow UI within the notebook for easy browsing:)
IFrame(src="http://localhost:5000", width=1000, height=600)

让我们导航到屏幕左上角的实验选项卡，并点击我们最近的运行记录，如下图所示。

运行页面显示您实验的总体元数据。您可以进一步导航至 Artifacts 标签页查看已记录的工件（模型）。

MLflow 会在运行过程中记录与您的模型及其环境相关的工件。 大多数记录的文件，例如 conda.yaml、python_env.yml 和 requirements.txt 是所有 MLflow 记录的标准文件，有助于在不同环境之间实现可复现性。 然而，有两组工件是专属于 LlamaIndex 的：

通过存储这些对象，MLflow能够重现您记录模型时所处的环境。

重要提示： MLflow 不会序列化 API 密钥。这些密钥必须作为环境变量存在于您的模型加载环境中。

最后，您可以通过导航到 Tracing 选项卡查看本教程期间记录的所有追踪记录完整列表。点击每一行后，您可以看到详细的追踪视图，类似于之前在输出单元格中显示的视图。

定制化与后续步骤

Section titled “Customization and Next Steps”

在生产系统中工作时，用户通常利用定制化的服务上下文，这可以通过LlamaIndex的设置对象来实现。