GenAI高级索引管道

RAG Nov 2, 2024

管理复杂的 GenAI 管道（尤其是具有多个模型和数据源的管道）可能是一项艰巨的任务。本文演示了如何将 LlamaIndex 与 Qdrant 和 MLflow 集成以简化 GenAI 应用程序的管理和部署。你将探索 MLflow 的功能（例如跟踪、模型打包和评估）如何实现 LlamaIndex 引擎的无缝处理。通过逐步实施，了解如何简化检索增强生成 (RAG) 工作流程、确保性能一致性并优化索引系统以实现更好的可扩展性和效率。

1、系统架构

该架构集成了多个关键组件，以构建高效且可扩展的检索增强生成 (RAG) 系统。其核心是利用 LlamaIndex 进行索引和检索、利用 Qdrant 进行矢量存储和搜索以及利用 MLflow 在整个生命周期内注册、服务和跟踪所有组件的模型。这种设计旨在管理来自各种来源的大量数据，同时确保最终用户能够快速检索和准确推理。

管道从提取来自各种来源（例如 Web 文档、内部文档和数据库记录）的文档开始。这些文档由变更数据捕获 (CDC) 机制处理，该机制跟踪数据的变化，确保新文档和更新的文档始终反映在索引过程中。数据工程师负责系统的这一部分，管理数据管道以将传入的文档输入到由 LlamaIndex 支持的索引器中。LlamaIndex 和 Qdrant 索引器创建这些文档的语义索引，并将向量存储在 Qdrant 数据库中，该数据库充当向量存储和检索引擎。LlamaIndex 和 Qdrant 之间的集成实现了高效、高性能的索引和检索。

索引完成后，系统使用 MLflow 保存模型，这在 LlamaIndex 模型的生命周期管理中起着关键作用。 MLflow 允许跟踪所有必要的元数据，包括模型的版本控制、模型跟踪和模型服务，确保模型开发、测试和部署的不同阶段的一致性和透明度。

在检索方面，RAG-Bot 使用相同的 LlamaIndex 和 Qdrant 设置与用户的查询进行交互。RAG-Bot 处理传入的请求并将其发送到 MLflow 的 LlamaIndex 检索器，后者查询 Qdrant 存储以根据查询嵌入获取最相关的索引文档。此检索器旨在高速工作，从而能够快速响应用户。然后，将检索到的数据作为推理过程的一部分发送回用户。

2、MLflow 设置

首先安装 MLflow

pip install mlflow

要开始使用 MLflow，你需要先安装它。MLflow 是一个用于管理机器学习实验、跟踪指标、打包代码和部署模型的开源工具。 pip install mlflow 命令将安装所有必要的依赖项，使 MLflow 可以在你的环境中使用。

用如下命令启动 MLflow 服务器：

mlflow server # this will run the server on port 5000

mlflow server --host 127.0.0.1 --port 3001  # runs the server on custom port

编写代码以使用 MLflow 跟踪实验：

import mlflow
import mlflow.sklearn
from sklearn.datasets import load_diabetes
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

# Load the dataset
data = load_diabetes()
X = data.data
y = data.target

# Split the data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Set model parameters
alpha = 0.5
l1_ratio = 0.5

# set mlflow tracking uri
mlflow.set_tracking_uri('http://127.0.0.1:3001/')

# Start an MLflow experiment
with mlflow.start_run():
    # Train the model
    model = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
    model.fit(X_train, y_train)

    # Predict on the test set
    predictions = model.predict(X_test)

    # Calculate metrics
    mse = mean_squared_error(y_test, predictions)
    r2 = r2_score(y_test, predictions)

    # Log parameters and metrics with MLflow
    mlflow.log_param("alpha", alpha)
    mlflow.log_param("l1_ratio", l1_ratio)
    mlflow.log_metric("mse", mse)
    mlflow.log_metric("r2", r2)

    # Log the model itself
    mlflow.sklearn.log_model(model, "model")

    # Advanced Use: Generate a plot for residuals
    plt.scatter(y_test, predictions)
    plt.xlabel("True Values")
    plt.ylabel("Predictions")
    plt.savefig("residuals_plot.png")

    # Log the plot as an artifact
    mlflow.log_artifact("residuals_plot.png")

print("MLflow run complete!")

我们创建一个 Python 脚本 (train.py)，使用 MLflow 来跟踪模型的训练和评估机器学习模型。我们首先导入必要的库，如 MLflow、scikit-learn 等。我们加载糖尿病数据集，将其分成训练集和测试集，并定义模型参数。使用 mlflow.start_run()，我们开始实验运行，训练 ElasticNet 回归模型并进行预测。我们计算评估指标（MSE 和 R2 分数）并使用 MLflow 的日志记录功能记录模型参数、指标和训练模型本身。这使我们能够轻松跟踪实验、比较不同的运行并存储模型以供日后使用。下面成功运行的代码之一是可以在 MLflow UI 中看到的跟踪。

如果你想要永久保留实验并想要追溯历史记录，请按如下方式运行 mlflow 服务器：

# this will plugin the mlflow with postgres and log all details to postgres.
mlflow server --backend-store-uri postgresql://root:root@localhost:5432/mlflow --default-artifact-root ./mlruns --host 127.0.0.1 --port 3001

3、设置 Qdrant

转到 qdrant.tech 并单击主屏幕右上角的登录：

在登录页面上，选择适当的方法进入云托管 qdrant 平台。登录后，你应该进入以下环境，如下所示。

如果你有一个集群，可以重新启动它，或者可以免费创建一个集群：

4、设置 Ollama

如果你使用的是 Mac，请使用以下命令，否则请转到这里选择合适的平台：

https://ollama.com/download/Ollama-darwin.zip # if on mac

安装后，通过运行以下命令检查 Ollama 是否已成功安装：

ollama --version

如果一切设置正确，此命令将打印已安装的版本。

接下来设置模型：

# LLM model
ollama pull llama3.2:latest

# Embedding Model
ollama pull nomic-embed-text:latest

检查系统中是否包含所有模型：

ollama list

5、 Playground 代码

现在我们已经运行了 MLflow 服务器、运行了 Qdrant 和运行了 Ollama，我们可以使用 LlamaIndex 在所有这些之间创建一个简单的集成，以测试所有工具是否都朝着同一个方向对齐。项目脚手架如下：

.
├── .env
├── data
│   └── vllm.pdf
├── model_indexing.py
├── model_logging.py
└── requirements.txt

让我们来编写 .env 文件：

OLLAMA_URL='http://localhost:11434'
EMBED_MODEL='nomic-embed-text:latest'
LANGUAGE_MODEL='llama3.2:latest'
QDRANT_URL='https://5496bdf1-fe1b-4e36-8715-aa5319aa1bf7.us-east4-0.gcp.cloud.qdrant.io:6333'
QDRANT_API_KEY='YbKL_l28******************************************'
COLLECTION_NAME='mlflow_collection'

让我们看看下面的 requirements.txt：

mlflow==2.17.1
psycopg2-binary==2.9.10
psycopg2==2.9.10
llama-index==0.11.20
llama-index-llms-ollama==0.3.4
llama-index-embeddings-ollama==0.3.1
llama-index-vector-stores-qdrant==0.3.2
qdrant-client==1.12.0
python-dotenv==1.0.1

让我们编写 model_indexing.py：

from llama_index.core import Settings, SimpleDirectoryReader, StorageContext, VectorStoreIndex
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.llms.ollama import Ollama
from llama_index.vector_stores.qdrant import QdrantVectorStore
from dotenv import load_dotenv, find_dotenv
import qdrant_client
import mlflow
import os

load_dotenv(find_dotenv())

Settings.llm = Ollama(model=os.environ.get('LANGUAGE_MODEL'), base_url=os.environ.get('OLLAMA_URL'))
Settings.embed_model = OllamaEmbedding(model_name=os.environ.get('EMBED_MODEL'), base_url=os.environ.get('OLLAMA_URL'))

documents = SimpleDirectoryReader(input_dir='data', required_exts=['.pdf']).load_data(show_progress=True)

client = qdrant_client.QdrantClient(url=os.environ.get('QDRANT_URL'), api_key=os.environ.get('QDRANT_API_KEY'))

vector_store = QdrantVectorStore(client=client, collection_name=os.environ.get('COLLECTION_NAME'))

storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents=documents, storage_context=storage_context, show_progress=True)

mlflow.models.set_model(index)

让我们编写 model_logging.py：

from llama_index.core import VectorStoreIndex
import mlflow

mlflow.set_tracking_uri('http://127.0.0.1:3000')
mlflow.set_experiment(experiment_name='llamaindex-qdrant-rag-1')

mlflow.llama_index.autolog()

with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        "model_indexing.py",
        artifact_path="index",
        engine_type="chat",  # chat, query & retrieve
    )

    print(model_info.model_uri)

    index: VectorStoreIndex = mlflow.llama_index.load_model(model_info.model_uri)
    
    print(index)

在运行 model_logging.py 之前，我们的 MLflow 服务器应该正在运行