跳到主要内容

MLflow LlamaIndex 插件

简介​

LlamaIndex 🦙 是一个以数据为中心的强大框架，旨在将自定义数据源无缝连接到大型语言模型 (LLM)。它提供了一套全面的数据结构和工具，简化了将私有或特定领域的数据摄取、结构化和访问的过程，以便与 LLM 一起使用。LlamaIndex 通过提供高效的索引和检索机制，在支持上下文感知的 AI 应用程序方面表现出色，从而更轻松地构建需要集成外部知识的高级问答系统、聊天机器人和其他 AI 驱动的应用程序。

Overview of LlamaIndex and MLflow integration

为什么要在 MLflow 中使用 LlamaIndex？​

将 LlamaIndex 库与 MLflow 集成，可以为管理和部署 LlamaIndex 引擎提供无缝的体验。以下是在 MLflow 中使用 LlamaIndex 的一些主要优势：

  • MLflow Tracking 允许您在 MLflow 中跟踪您的索引，并管理构成 LlamaIndex 项目的许多移动部件，例如提示、LLM、工作流、工具、全局配置等。
  • MLflow Model 将您的 LlamaIndex 索引/引擎/工作流与其所有依赖版本、输入输出接口和其他重要元数据一起打包。这使您可以轻松地部署 LlamaIndex 模型以进行推理，从而确保在 ML 生命周期不同阶段的环境一致性。
  • MLflow Evaluate 在 MLflow 中提供了评估 GenAI 应用程序的本机功能。此功能有助于对 LlamaIndex 模型的推理结果进行高效评估，确保稳健的性能分析并促进快速迭代。
  • MLflow Tracing 是一个强大的可观察性工具，用于监视和调试 LlamaIndex 模型内部发生的情况，帮助您快速识别潜在的瓶颈或问题。凭借其强大的自动记录功能，您无需添加任何代码，只需运行一个命令即可为 LlamaIndex 应用程序添加插装。

开始使用​

在这些入门教程中，您将了解 LlamaIndex 最基本组件，以及如何利用与 MLflow 的集成来为您的 LlamaIndex 应用程序带来更好的可维护性和可观察性。

MLflow LlamaIndex 工作流

通过构建一个简单的代理工作流，开始使用 MLflow 和 LLamaIndex。了解如何记录和加载工作流以进行推理，以及启用跟踪以获得可观察性。

使用 MLflow 构建索引

通过探索最简单的 VectorStoreIndex 配置，开始使用 MLflow 和 LlamaIndex。

概念​

注意

工作流集成仅在 LlamaIndex >= 0.11.0 和 MLflow >= 2.17.0 中可用。

Workflow 🆕​

Workflow 是 LlamaIndex 的事件驱动编排框架。它被设计为一个灵活且可解释的框架，用于构建任意 LLM 应用程序，例如代理、RAG 流、数据提取管道等。MLflow 支持跟踪、评估和跟踪 Workflow 对象，这使其更具可观察性和可维护性。

Index​

Index 对象是为快速信息检索而索引的文档集合，为检索增强生成 (RAG) 和代理等应用程序提供功能。Index 对象可以直接记录到 MLflow 运行中，并重新加载以用作推理引擎。

Engine​

Engine 是建立在 Index 对象之上的通用接口，它提供了一组 API 来与索引进行交互。LlamaIndex 提供两种类型的引擎：QueryEngine 和 ChatEngine。QueryEngine 仅接受单个查询并根据索引返回响应。ChatEngine 是为对话代理设计的，它会跟踪对话历史记录。

Settings​

Settings 对象是一个全局服务上下文，它捆绑了 LlamaIndex 应用程序中常用的资源。它包括 LLM 模型、嵌入模型、回调等设置。记录 LlamaIndex 索引/引擎/工作流时，MLflow 会跟踪 Settings 对象的状态，以便您在重新加载模型以进行推理时轻松重现相同的结果（请注意，某些对象，如 API 密钥、不可序列化的对象等，不会被跟踪）。

用法​

在 MLflow 实验中保存和加载索引​

创建索引​

index 对象是 LlamaIndex 和 MLflow 集成的核心。使用 LlamaIndex，您可以从文档集合或外部向量存储创建索引。以下代码从 LlamaIndex 存储库中可用的 Paul Graham 的论文数据创建了一个示例索引。

mkdir -p data
curl -L https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt -o ./data/paul_graham_essay.txt
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

documents = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(documents)

将索引记录到 MLflow​

您可以使用 mlflow.llama_index.log_model() 函数将 index 对象记录到 MLflow 实验中。

这里的一个关键步骤是指定 engine_type 参数。引擎类型的选择不影响索引本身，但决定了重新加载它以进行推理时查询索引的接口。

  1. QueryEngine (engine_type="query") 专为简单的查询-响应系统而设计，该系统接受单个查询字符串并返回响应。
  2. ChatEngine (engine_type="chat") 专为对话代理而设计，该代理会跟踪对话历史记录并响应用户消息。
  3. Retriever (engine_type="retriever") 是一个低级组件，它返回与查询匹配的前 k 个相关文档。

以下代码是一个使用 chat 引擎类型将索引记录到 MLflow 的示例。

import mlflow

mlflow.set_experiment("llama-index-demo")

with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        index,
        name="index",
        engine_type="chat",
        input_example="What did the author do growing up?",
    )
注意

上面的代码片段将索引对象直接传递给 log_model 函数。此方法仅适用于默认的 SimpleVectorStore 向量存储，该存储仅将嵌入的文档保存在内存中。如果您的索引使用了诸如 QdrantVectorStore 或 DatabricksVectorSearch 等外部向量存储，则可以使用 Model-from-Code 记录方法。有关更多详细信息，请参阅 如何使用外部向量存储记录和加载索引。

MLflow artifacts for the LlamaIndex index

提示

在底层，MLflow 调用索引对象的 as_query_engine() / as_chat_engine() / as_retriever() 方法将其转换为相应的引擎实例。

重新加载索引以进行推理​

可以使用 mlflow.pyfunc.load_model() 函数重新加载已保存的索引以进行推理。此函数提供了一个由 LlamaIndex 引擎支持的 MLflow Python 模型，引擎类型在记录时指定。

import mlflow

model = mlflow.pyfunc.load_model(model_info.model_uri)

response = model.predict("What was the first program the author wrote?")
print(response)
# >> The first program the author wrote was on the IBM 1401 ...

# The chat engine keeps track of the conversation history
response = model.predict("How did the author feel about it?")
print(response)
# >> The author felt puzzled by the first program ...
提示

要重新加载索引本身而不是引擎，请使用 mlflow.llama_index.load_model() 函数。

index = mlflow.llama_index.load_model("runs:/<run_id>/index")

启用跟踪​

您可以通过调用 mlflow.llama_index.autolog() 函数来为 LlamaIndex 代码启用跟踪。MLflow 会自动将 LlamaIndex 执行的输入和输出记录到活动的 MLflow 实验中，为您提供模型行为的详细视图。

import mlflow

mlflow.llama_index.autolog()

chat_engine = index.as_chat_engine()
response = chat_engine.chat("What was the first program the author wrote?")

然后，您可以导航到 MLflow UI，选择实验，然后打开“Traces”选项卡，以查找引擎预测所记录的跟踪。看到聊天引擎如何协调和执行一系列任务来回答您的问题，真是令人印象深刻！

Trace view in MLflow UI

您可以通过运行具有 disable 参数设置为 True 的同一函数来禁用跟踪。

mlflow.llama_index.autolog(disable=True)
注意

跟踪支持异步预测和流式响应，但不支持异步和流式的组合，例如 astream_chat 方法。

常见问题​

如何使用外部向量存储记录和加载索引？​

如果您的索引使用默认的 SimpleVectorStore，则可以使用 mlflow.llama_index.log_model() 函数将索引直接记录到 MLflow。MLflow 将内存中的索引数据（嵌入文档）持久化到 MLflow 构件存储中，这允许在不重新索引文档的情况下重新加载索引。

但是，当索引使用 DatabricksVectorSearch 和 QdrantVectorStore 等外部向量存储时，索引数据存储在远程，并且不支持本地序列化。因此，您无法直接记录这些存储的索引。对于这种情况，您可以使用 Model-from-Code 记录，它提供了对索引保存过程的更多控制，并允许您使用外部向量存储。

要使用 Model-from-Code 记录，您首先需要创建一个单独的 Python 文件来定义索引。如果您在 Jupyter Notebook 中，可以使用 %%writefile 魔术命令将单元格代码保存到 Python 文件。

# %%writefile index.py

# Create Qdrant client with your own settings.
client = qdrant_client.QdrantClient(
    host="localhost",
    port=6333,
)

# Here we simply load vector store from the existing collection to avoid
# re-indexing documents, because this Python file is executed every time
# when the model is loaded. If you don't have an existing collection, create
# a new one by following the official tutorial:
# https://docs.llamaindex.org.cn/en/stable/examples/vector_stores/QdrantIndexDemo/
vector_store = QdrantVectorStore(client=client, collection_name="my_collection")
index = VectorStoreIndex.from_vector_store(vector_store=vector_store)

# IMPORTANT: call set_model() method to tell MLflow to log this index
mlflow.models.set_model(index)

然后，您可以将 Python 文件路径传递给 mlflow.llama_index.log_model() 函数来记录索引。全局 Settings 对象正常保存为模型的一部分。

import mlflow

with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        "index.py",
        name="index",
        engine_type="query",
    )

可以使用 mlflow.llama_index.load_model() 或 mlflow.pyfunc.load_model() 函数以与本地索引相同的方式重新加载记录的索引。

index = mlflow.llama_index.load_model(model_info.model_uri)
index.as_query_engine().query("What is MLflow?")
注意

传递给 set_model() 方法的对象必须是与记录期间指定的引擎类型兼容的 LlamaIndex 索引。未来的版本将增加对更多对象的支持。

如何记录和加载 LlamaIndex 工作流？​

MLflow 支持通过 Model-from-Code 功能记录和加载 LlamaIndex 工作流。有关记录和加载 LlamaIndex 工作流的详细示例，请参阅 MLflow LlamaIndex 工作流 Notebook。

import mlflow

with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        "/path/to/workflow.py",
        name="model",
        input_example={"input": "What is MLflow?"},
    )

可以使用 mlflow.llama_index.load_model() 或 mlflow.pyfunc.load_model() 函数重新加载记录的工作流。

# Use mlflow.llama_index.load_model to load the workflow object as is
workflow = mlflow.llama_index.load_model(model_info.model_uri)
await workflow.run(input="What is MLflow?")

# Use mlflow.pyfunc.load_model to load the workflow as a MLflow Pyfunc Model
# with standard inference APIs for deployment and evaluation.
pyfunc = mlflow.pyfunc.load_model(model_info.model_uri)
pyfunc.predict({"input": "What is MLflow?"})
警告

MLflow PyFunc Model 不支持异步推理。当您使用 mlflow.pyfunc.load_model() 加载工作流时，predict 方法将变为同步，并会一直阻塞直到工作流执行完成。当使用 MLflow Deployment 或 Databricks Model Serving 将记录的 LlamaIndex 工作流部署到生产端点时，这也适用。

我有一个使用 query 引擎类型记录的索引。我能将其加载为 chat 引擎吗？​

虽然无法就地更新记录模型的引擎类型，但您可以随时加载索引并以所需的引擎类型重新记录它。此过程无需重新创建索引，因此是切换不同引擎类型的有效方法。

import mlflow

# Log the index with the query engine type first
with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        index,
        name="index-query",
        engine_type="query",
    )

# Load the index back and re-log it with the chat engine type
index = mlflow.llama_index.load_model(model_info.model_uri)
with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        index,
        name="index-chat",
        # Specify the chat engine type this time
        engine_type="chat",
    )

或者，您可以利用加载的 LlamaIndex本机索引对象的标准推理 API，特别是

  • index.as_chat_engine().chat("hi")
  • index.as_query_engine().query("hi")
  • index.as_retriever().retrieve("hi")

如何使用不同的 LLM 进行加载引擎的推理？​

当将索引保存到 MLflow 时，它会将全局 Settings 对象作为模型的一部分进行持久化。此对象包含引擎使用的 LLM 和嵌入模型等设置。

import mlflow
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI

Settings.llm = OpenAI("gpt-4o-mini")

# MLflow saves GPT-4o-Mini as the LLM to use for inference
with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(index, name="index", engine_type="chat")

然后，当您稍后重新加载索引时，持久化的设置也会全局应用。这意味着加载的引擎将使用与记录时相同的 LLM。

但是，有时您可能希望使用不同的 LLM 进行推理。在这种情况下，您可以在加载索引后直接更新全局 Settings 对象。

import mlflow

# Load the index back
loaded_index = mlflow.llama_index.load_model(model_info.model_uri)

assert Settings.llm.model == "gpt-4o-mini"


# Update the settings to use GPT-4 instead
Settings.llm = OpenAI("gpt-4")
query_engine = loaded_index.as_query_engine()
response = query_engine.query("What is the capital of France?")