使用 MLflow 评估 LLM RAG 示例笔记本

下载此笔记本

欢迎来到本关于使用 MLflow 评估检索增强生成 (RAG) 系统的全面教程。本教程旨在指导您深入了解评估各种 RAG 系统的复杂性，重点关注如何在现实世界中有效集成和评估它们。无论您是数据科学家、机器学习工程师，还是仅仅是人工智能领域的爱好者，本教程都将提供宝贵的见解和实践知识。

您将学到：

环境设置:
- 学习如何使用所有必需的工具和库（包括 MLflow、OpenAI、ChromaDB、LangChain 等）设置您的开发环境。本节确保您拥有开始使用 RAG 系统所需的一切。
了解 RAG 系统:
- 深入探讨检索增强生成的概念及其在现代 AI 应用中的重要性。了解 RAG 系统如何利用检索和生成能力来提供准确且具有上下文相关性的响应。
使用 Databricks Secrets 安全管理 API 密钥:
- 探索使用 Databricks Secrets 安全管理 API 密钥的最佳实践。这对于确保应用程序的安全性和完整性至关重要。
使用 MLflow 部署和测试 RAG 系统:
- 学习如何使用 MLflow 创建、部署和测试 RAG 系统。这包括设置端点、部署模型以及查询它们以查看它们的响应。
使用 MLflow 评估性能:
- 深入使用 MLflow 的评估工具评估 RAG 系统。了解如何使用相关性和延迟等指标来评估 RAG 系统的性能。
实验不同的分块策略:
- 尝试不同的文本分块策略以优化 RAG 系统的性能。了解文本块的大小如何影响检索准确性和系统响应速度。
创建和使用评估数据集:
- 学习如何创建和利用评估数据集（黄金数据集）来有效地评估 RAG 系统的性能。
结合检索和生成进行问答:
- 深入了解 RAG 系统中的检索和生成组件如何协同工作，以基于给定上下文或文档回答问题。

通过本教程的学习，您将全面了解如何使用 MLflow 评估和优化 RAG 系统。您将掌握部署、测试和改进 RAG 系统的知识，使其适用于各种实际应用。本教程是您迈入高级 AI 模型评估和部署世界的垫脚石。

%pip install mlflow>=2.8.1
%pip install openai
%pip install chromadb==0.4.15
%pip install langchain==0.0.348
%pip install tiktoken
%pip install 'mlflow[genai]'
%pip install databricks-sdk --upgrade

dbutils.library.restartPython()  # noqa: F821

import ast
import os

import chromadb
import pandas as pd
from langchain.chains import RetrievalQA
from langchain.document_loaders import WebBaseLoader
from langchain.embeddings.databricks import DatabricksEmbeddings
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain.llms import Databricks
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma

import mlflow
import mlflow.deployments
from mlflow.deployments import set_deployments_target
from mlflow.metrics.genai.metric_definitions import relevance

# check mlflow version
mlflow.__version__

'2.9.1'

# check chroma version
chromadb.__version__

'0.4.18'

设置 Databricks 工作区密钥

为了使用此笔记本中定义的密钥，请确保按照这里的 Databricks Secrets 指南进行设置。强烈建议使用Databricks CLI在工作区中设置密钥，以获得安全体验。

为了安全地存储和访问您的 Azure OpenAI API 密钥，请确保在注册密钥时设置以下内容

KEY_NAME：您为 Azure OpenAI 密钥设置的名称
SCOPE_NAME：您的密钥在 Databricks Secrets 中所属的引用作用域
OPENAI_API_KEY：您的 Azure OpenAI 密钥

例如，您可以通过终端如下设置这些密钥

    databricks secrets put-secret "<SCOPE_NAME>" "<KEY_NAME>" --string-value "<OPENAI_API_KEY>"

# Set your Scope and Key Names that you used when registering your API KEY from the Databricks CLI
# Do not put your OpenAI API Key in the notebook!
SCOPE_NAME = ...
KEY_NAME = ...

os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(scope=SCOPE_NAME, key=KEY_NAME)  # noqa: F821
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"
# Ensure that you set the name of your OPEN_API_BASE value to the name of your OpenAI instance on Azure
os.environ["OPENAI_API_BASE"] = "https://<NAME_OF_YOUR_INSTANCE>.openai.azure.com/"  # replace this!
os.environ["OPENAI_DEPLOYMENT_NAME"] = "gpt-4o-mini"
os.environ["OPENAI_ENGINE"] = "gpt-4o-mini"

在 MLflow 上为 OpenAI 创建和测试端点

client = mlflow.deployments.get_deploy_client("databricks")

endpoint_name = "<your-endpoint-name>"  # replace this!
client.create_endpoint(
  name=endpoint_name,
  config={
      "served_entities": [
          {
              "name": "test-gpt",  # Provide a unique identifying name for your deployments endpoint
              "external_model": {
                  "name": "gpt-4o-mini",
                  "provider": "openai",
                  "task": "llm/v1/completions",
                  "openai_config": {
                      "openai_api_type": "azure",
                      # replace with your own secrets, for reference see https://docs.databricks.com/en/security/secrets/secrets.html
                      "openai_api_key": "{{secrets/scope/openai_api_key}}",
                      "openai_api_base": "{{secrets/scope/openai_api_base}}",
                      "openai_deployment_name": "gpt-4o-mini",
                      "openai_api_version": "2023-05-15",
                  },
              },
          }
      ],
  },
)

print(
  client.predict(
      endpoint=endpoint_name,
      inputs={
          "prompt": "How is Pi calculated? Be very concise.",
          "max_tokens": 100,
      },
  )
)

使用 LangChain 创建 RAG POC 并使用 MLflow 记录

使用 Langchain 和 Chroma 创建一个 RAG 系统，该系统根据 MLflow 文档回答问题。

loader = WebBaseLoader(
  [
      "https://mlflow.org.cn/docs/latest/index.html",
      "https://mlflow.org.cn/docs/latest/tracking/autolog.html",
      "https://mlflow.org.cn/docs/latest/getting-started/tracking-server-overview/index.html",
      "https://mlflow.org.cn/docs/latest/python_api/mlflow.deployments.html",
  ]
)

documents = loader.load()
CHUNK_SIZE = 1000
text_splitter = CharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=0)
texts = text_splitter.split_documents(documents)

llm = Databricks(
  endpoint_name="<your-endpoint-name>",  # replace this!
  extra_params={
      "temperature": 0.1,
      "top_p": 0.1,
      "max_tokens": 500,
  },  # parameters used in AI Playground
)


# create the embedding function using Databricks Foundation Model APIs
embedding_function = DatabricksEmbeddings(endpoint="databricks-bge-large-en")
docsearch = Chroma.from_documents(texts, embedding_function)

qa = RetrievalQA.from_chain_type(
  llm=llm,
  chain_type="stuff",
  retriever=docsearch.as_retriever(fetch_k=3),
  return_source_documents=True,
)

使用 `mlflow.evaluate()` 评估向量数据库和检索

创建一个评估数据集（黄金数据集）

我们可以利用 LLM 的强大能力生成合成数据进行测试，这提供了一种创新且高效的替代方案。我们向读者和客户强调，创建能反映您的 RAG 应用预期输入和输出的数据集至关重要。这是一段值得投入的旅程，您将从中获得令人难以置信的见解！

EVALUATION_DATASET_PATH = "https://raw.githubusercontent.com/mlflow/mlflow/master/examples/llms/RAG/static_evaluation_dataset.csv"

synthetic_eval_data = pd.read_csv(EVALUATION_DATASET_PATH)

# Load the static evaluation dataset from disk and deserialize the source and retrieved doc ids
synthetic_eval_data["source"] = synthetic_eval_data["source"].apply(ast.literal_eval)
synthetic_eval_data["retrieved_doc_ids"] = synthetic_eval_data["retrieved_doc_ids"].apply(
  ast.literal_eval
)

display(synthetic_eval_data)

使用 MLflow 评估 Embedding 模型

在本教程的这一部分，我们重点评估 embedding 模型在检索系统中的性能。该过程涉及一系列步骤，以评估模型基于给定问题检索相关文档的效率。

创建评估数据

我们首先定义一组问题及其对应的源 URL。此 eval_data DataFrame 作为我们的评估数据集，允许我们测试模型将问题链接到正确源文档的能力。

`evaluate_embedding` 函数

evaluate_embedding 函数旨在评估给定 embedding 函数的性能。
分块策略：该函数首先使用 CharacterTextSplitter 将文档列表分成块。这些块的大小至关重要，因为它会影响检索准确性。
检索器初始化：然后我们使用 Chroma.from_documents 和指定的 embedding 函数创建一个检索器。此检索器负责查找与给定查询相关的文档。
检索过程：该函数定义了一个 retriever_model_function，该函数将检索器应用于评估数据集中的每个问题。它检索模型认为与每个问题最相关的文档 ID。

MLflow 评估

使用 mlflow.start_run()，我们启动一个评估运行。然后调用 mlflow.evaluate 以根据评估数据集评估我们的检索器模型函数。
我们使用默认评估器和指定目标来评估模型的性能。
本次评估的结果存储在 eval_results_of_retriever_df_bge 中并显示出来，提供了关于 embedding 模型在文档检索中的有效性的见解。

使用指标进一步评估

此外，我们使用各种指标（如精度、召回率和不同 'k' 值下的 NDCG）进行更详细的评估。这些指标提供了对模型检索准确性和排名有效性的更深入理解。

此评估步骤是理解 embedding 模型在真实 RAG 系统中的优缺点不可或缺的部分。通过分析这些结果，我们可以就模型调整或优化做出明智的决策，以提高整体系统性能。

eval_data = pd.DataFrame(
  {
      "question": [
          "What is MLflow?",
          "What is Databricks?",
          "How to serve a model on Databricks?",
          "How to enable MLflow Autologging for my workspace by default?",
      ],
      "source": [
          ["https://mlflow.org.cn/docs/latest/index.html"],
          ["https://mlflow.org.cn/docs/latest/getting-started/tracking-server-overview/index.html"],
          ["https://mlflow.org.cn/docs/latest/python_api/mlflow.deployments.html"],
          ["https://mlflow.org.cn/docs/latest/tracking/autolog.html"],
      ],
  }
)

def evaluate_embedding(embedding_function):
  CHUNK_SIZE = 1000
  list_of_documents = loader.load()
  text_splitter = CharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=0)
  docs = text_splitter.split_documents(list_of_documents)
  retriever = Chroma.from_documents(docs, embedding_function).as_retriever()

  def retrieve_doc_ids(question: str) -> list[str]:
      docs = retriever.get_relevant_documents(question)
      return [doc.metadata["source"] for doc in docs]

  def retriever_model_function(question_df: pd.DataFrame) -> pd.Series:
      return question_df["question"].apply(retrieve_doc_ids)

  with mlflow.start_run():
      return mlflow.evaluate(
          model=retriever_model_function,
          data=eval_data,
          model_type="retriever",
          targets="source",
          evaluators="default",
      )


result1 = evaluate_embedding(DatabricksEmbeddings(endpoint="databricks-bge-large-en"))
# To validate the results of a different model, comment out the above line and uncomment the below line:
# result2 = evaluate_embedding(SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2"))

eval_results_of_retriever_df_bge = result1.tables["eval_results_table"]
# To validate the results of a different model, comment out the above line and uncomment the below line:
# eval_results_of_retriever_df_MiniLM = result2.tables["eval_results_table"]
display(eval_results_of_retriever_df_bge)

使用 MLflow 评估不同的 Top K 策略

with mlflow.start_run() as run:
  evaluate_results = mlflow.evaluate(
      data=eval_results_of_retriever_df_bge,
      targets="source",
      predictions="outputs",
      evaluators="default",
      extra_metrics=[
          mlflow.metrics.precision_at_k(1),
          mlflow.metrics.precision_at_k(2),
          mlflow.metrics.precision_at_k(3),
          mlflow.metrics.recall_at_k(1),
          mlflow.metrics.recall_at_k(2),
          mlflow.metrics.recall_at_k(3),
          mlflow.metrics.ndcg_at_k(1),
          mlflow.metrics.ndcg_at_k(2),
          mlflow.metrics.ndcg_at_k(3),
      ],
  )

display(evaluate_results.tables["eval_results_table"])

使用 MLflow 评估分块策略

在 RAG 系统领域，将文本划分为块的策略对检索效率和整体系统性能都起着关键作用。让我们深入探讨为什么要以及如何评估不同的分块策略

分块的重要性：

影响检索准确性：文本的分块方式会显著影响 RAG 系统的检索组件。较小的块可能导致更专注和相关的文档检索，而较大的块可能捕获更广泛的上下文。
影响系统响应速度：文本块的大小也影响文档检索和处理的速度。较小的块可以更快地处理，但可能需要系统整体评估更多的块。

评估不同块大小：

目的：通过评估不同的块大小，我们旨在找到检索准确性和处理效率之间的最佳平衡。这涉及尝试各种块大小，以了解它们如何影响系统的性能。
方法：我们创建不同大小的文本块（例如，1000 个字符，2000 个字符），然后评估每种分块策略对 RAG 系统有什么影响。需要观察的关键方面包括检索到的文档的相关性和系统的延迟。

在下面的示例中，我们使用默认的评估套件对检索到的文档内容的响应质量进行全面裁定，以确定对返回引用的质量有何影响，从而使我们能够探索和调整块大小，以便找到最能处理我们测试问题集的配置。

请注意，下一个代码块中的 embedding 模型已更改。上面我们使用的是 DatabricksEmbeddings(endpoint="databricks-bge-large-en")，而现在我们正在评估 SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") 的性能

def evaluate_chunk_size(chunk_size):
  list_of_documents = loader.load()
  text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=0)
  docs = text_splitter.split_documents(list_of_documents)
  embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
  retriever = Chroma.from_documents(docs, embedding_function).as_retriever()

  def retrieve_doc_ids(question: str) -> list[str]:
      docs = retriever.get_relevant_documents(question)
      return [doc.metadata["source"] for doc in docs]

  def retriever_model_function(question_df: pd.DataFrame) -> pd.Series:
      return question_df["question"].apply(retrieve_doc_ids)

  with mlflow.start_run():
      return mlflow.evaluate(
          model=retriever_model_function,
          data=eval_data,
          model_type="retriever",
          targets="source",
          evaluators="default",
      )


result1 = evaluate_chunk_size(1000)
result2 = evaluate_chunk_size(2000)


display(result1.tables["eval_results_table"])
display(result2.tables["eval_results_table"])

使用 `mlflow.evaluate()` 评估 RAG 系统

在本节中，我们将深入探讨如何使用 mlflow.evaluate() 评估检索增强生成 (RAG) 系统。此评估对于评估 RAG 系统在问答上下文中的有效性和效率至关重要。我们重点关注两个关键指标：relevance_metric 和 latency。

延迟：

衡量内容：latency 指标捕获 RAG 系统的响应时间。它衡量系统在接收查询后生成答案所需的时间。
重要性：响应时间是用户体验的关键因素。在交互式系统中，较低的延迟带来更高效和令人满意的用户体验。相反，高延迟可能对用户满意度产生不利影响。
教程上下文：在本教程中，我们通过 latency 指标评估系统在响应时间方面的效率。此评估对于理解系统在生产环境中的性能至关重要，其中及时的响应与准确性同等重要。

为了开始评估，我们将创建一个简单函数，该函数通过 RAG 链运行每个输入

def model(input_df):
  return input_df["questions"].map(qa).tolist()

创建评估数据集（黄金数据集）

eval_df = pd.DataFrame(
  {
      "questions": [
          "What is MLflow?",
          "What is Databricks?",
          "How to serve a model on Databricks?",
          "How to enable MLflow Autologging for my workspace by default?",
      ],
  }
)
display(eval_df)

使用 LLM 作为裁判和基本指标进行评估

在本教程的最后一部分，我们使用 MLflow 强大的评估工具对 RAG 系统进行最终评估。此评估对于评估问答模型的性能和效率至关重要。

评估的关键步骤：

设置部署目标:
- 部署目标设置为 Databricks，使我们能够检索 Databricks 工作区中的所有端点。这对于访问我们部署的模型至关重要。
相关性指标设置:
- 我们使用托管在 Databricks 上的模型初始化 relevance 指标。此指标评估 RAG 系统生成的答案与输入问题的相关程度。
运行评估:
- 启动 MLflow 运行，并调用 mlflow.evaluate() 以根据准备好的评估数据集评估我们的 RAG 模型。
- 使用默认评估器将模型评估为“问答”系统。
- 指定其他指标，包括 relevance_metric 和 latency。这些指标提供有关答案相关性和模型响应时间的见解。
- evaluator_config 映射输入问题和上下文，确保对 RAG 系统进行正确评估。
结果和指标显示:
- 评估结果（包括关键指标）以表格形式显示，提供了基于相关性和延迟的模型性能的清晰且结构化的视图。

此全面的评估步骤对于理解 RAG 系统的有效性和效率至关重要。通过评估答案的相关性和响应的延迟，我们获得了对模型性能的整体视图，指导任何进一步的优化或部署决策。

set_deployments_target("databricks")  # To retrieve all endpoint in your Databricks Workspace

relevance_metric = relevance(
  model="endpoints:/databricks-llama-2-70b-chat"
)  # You can also use any model you have hosted on Databricks, models from the Marketplace or models in the Foundation model API

with mlflow.start_run():
  results = mlflow.evaluate(
      model,
      eval_df,
      model_type="question-answering",
      evaluators="default",
      predictions="result",
      extra_metrics=[relevance_metric, mlflow.metrics.latency()],
      evaluator_config={
          "col_mapping": {
              "inputs": "questions",
              "context": "source_documents",
          }
      },
  )
  print(results.metrics)

display(results.tables["eval_results_table"])

您将学到：​

设置 Databricks 工作区密钥​

在 MLflow 上为 OpenAI 创建和测试端点​

使用 LangChain 创建 RAG POC 并使用 MLflow 记录​

使用 mlflow.evaluate() 评估向量数据库和检索​

创建一个评估数据集（黄金数据集）​

使用 MLflow 评估 Embedding 模型​

创建评估数据​

evaluate_embedding 函数​

MLflow 评估​

使用指标进一步评估​

使用 MLflow 评估不同的 Top K 策略​

使用 MLflow 评估分块策略​

分块的重要性：​

评估不同块大小：​

使用 mlflow.evaluate() 评估 RAG 系统​

相关性指标：​

延迟：​

创建评估数据集（黄金数据集）​

使用 LLM 作为裁判和基本指标进行评估​

评估的关键步骤：​