我们如何利用 distilabel 创建 Argilla 2.0 聊天机器人
太长不看 (TL;DR)
了解如何为您选择的工具（本例中为 Argilla 2.0）构建一个聊天机器人，使其能够理解技术文档并与用户就此进行交流。
在本文中，我们将向您展示如何利用 distilabel 并微调一个领域特定的嵌入模型，以创建一个既准确又具吸引力的对话模型。
本文概述了为 Argilla 2.0 创建聊天机器人的过程。我们将：
- 从技术文档中创建一个合成数据集，用于微调特定领域的嵌入模型，
- 创建一个向量数据库来存储和检索文档，以及
- 将最终的聊天机器人部署到 Hugging Face Space，允许用户与其互动，并将互动数据存储在 Argilla 中以进行持续评估和改进。
点击此处跳转至应用程序。
目录
- 为微调领域特定的嵌入模型生成合成数据
- 下载和分块数据
- 使用 distilabel 为我们的嵌入模型生成合成数据
- 在 Argilla 中探索数据集
- 微调嵌入模型
- 向量数据库
- 连接到数据库
- 创建包含文档块的表
- 创建我们的聊天机器人
- Gradio 应用
- 在 Hugging Face Spaces 上部署聊天机器人应用
- 体验我们的聊天机器人
- 后续步骤
为微调自定义嵌入模型生成合成数据
需要快速回顾一下 RAG 吗？可以通过这个方便的入门笔记本来复习基础知识。我们会等你跟上进度！
下载和分块数据
分块数据意味着将您的文本数据分割成每个约 256 个 token 的可管理块（后续 RAG 中使用的块大小）。
让我们进入第一步：处理目标仓库的文档。为了简化此任务，您可以利用像 llama-index 这样的库来读取仓库内容并解析 Markdown 文件。具体来说，langchain 提供了有用的工具，如 MarkdownTextSplitter，而 llama-index 提供了 MarkdownNodeParser 来帮助您提取必要的信息。如果您更喜欢一种更简化的方法，可以考虑使用 davanstrien 的 corpus-creator 应用。
为了使事情更简单高效，我们开发了一个自定义的 Python 脚本来为您完成繁重的工作。您可以在我们的仓库中找到它，地址是 这里。
该脚本自动化了从 GitHub 仓库中检索文档并将其作为数据集存储在 Hugging Face Hub 上的过程。最棒的是，它使用起来非常简单！让我们看看如何运行它。
python docs_dataset.py \
"argilla-io/argilla-python" \
--dataset-name "plaguss/argilla_sdk_docs_raw_unstructured"
虽然该脚本易于使用，但您可以通过使用额外的参数来进一步根据您的需求进行定制。然而，有两个必要的输入您需要提供：
存放文档的 GitHub 仓库路径。
Hugging Face Hub 的数据集 ID，您的数据集将存储在这里。
一旦您提供了这些必需的参数，脚本将负责其余的工作。以下是幕后发生的事情：
该脚本会从指定的 GitHub 仓库下载文档到您的本地目录。默认情况下，它会在
/docs目录中查找文档，但您可以通过指定不同的路径来更改此设置。它从下载的文档中提取所有 Markdown 文件。
将提取的 Markdown 文件分块成可管理的小块。
最后，它将准备好的数据集推送到 Hugging Face Hub，使其准备就绪。
为了让您更好地了解该脚本的内部工作原理，这里有一段代码片段总结了其核心逻辑：
# The function definitions are omitted for brevity, visit the script for more info!
from github import Github
gh = Github()
repo = gh.get_repo("repo_name")
# Download the folder
download_folder(repo, "/folder/with/docs", "dir/to/download/docs")
# Extract the markdown files from the downloaded folder with the documentation from the GitHub repository
md_files = list(docs_path.glob("**/*.md"))
# Loop to iterate over the files and generate chunks from the text pieces
data = create_chunks(md_files)
# Create a dataset to push it to the hub
create_dataset(data, repo_name="name/of/the/dataset")
该脚本包含用于下载文档、从 Markdown 文件创建块以及创建数据集的简短函数。包含更多功能或实现更复杂的分块策略应该会很简单。
您可以查看可用的参数
点击查看 docs_dataset.py 帮助信息
$ python docs_dataset.py -h
usage: docs_dataset.py [-h] [--dataset-name DATASET_NAME] [--docs_folder DOCS_FOLDER] [--output_dir OUTPUT_DIR] [--private | --no-private] repo [repo ...]
Download the docs from a github repository and generate a dataset from the markdown files. The dataset will be pushed to the hub.
positional arguments:
repo Name of the repository in the hub. For example 'argilla-io/argilla-python'.
options:
-h, --help show this help message and exit
--dataset-name DATASET_NAME
Name to give to the new dataset. For example 'my-name/argilla_sdk_docs_raw'.
--docs_folder DOCS_FOLDER
Name of the docs folder in the repo, defaults to 'docs'.
--output_dir OUTPUT_DIR
Path to save the downloaded files from the repo (optional)
--private, --no-private
Whether to keep the repository private or not. Defaults to False.
使用 distilabel 为我们的嵌入模型生成合成数据
我们将从我们的文档中生成合成问题，这些问题可以由每个文档块来回答。我们还将通过生成易于区分的不相关问题来生成困难负例。我们可以使用这些问题、困难负例和文档来构建用于微调数据集的三元组。
完整的管道脚本可以在参考仓库的 pipeline_docs_queries.py 中看到，但让我们回顾一下不同的步骤
load_data:
我们旅程的第一步是获取包含有价值的文档块的数据集。仔细观察后，我们注意到包含这些块的列被恰当地命名为 `chunks`。然而，为了让我们的模型无缝运行，我们需要给这个列赋予一个新的身份。具体来说，我们想将它重命名为 `anchor`，因为这是我们后续步骤所期望的输入。我们将使用 `output_mappings` 来为我们完成这个列的转换
load_data = LoadDataFromHub(
name="load_data",
repo_id="plaguss/argilla_sdk_docs_raw_unstructured",
output_mappings={"chunks": "anchor"},
batch_size=10,
)
generate_sentence_pair
现在，我们来到了我们过程中最有趣的部分，将文档片段转化为合成查询。这就是 GenerateSentencePair 任务大显身手的地方。这个强大的任务为生成高质量的句子对提供了广泛的可能性。我们鼓励您探索其文档以解锁其全部潜力。
在我们的特定用例中，我们将利用 GenerateSentencePair 的能力来制作合成查询，最终提升我们模型的性能。让我们更深入地了解如何配置此任务以实现我们的目标。
llm = InferenceEndpointsLLM(
model_id="meta-llama/Meta-Llama-3-70B-Instruct",
tokenizer_id="meta-llama/Meta-Llama-3-70B-Instruct",
)
generate_sentence_pair = GenerateSentencePair(
name="generate_sentence_pair",
triplet=True, # Generate positive and negative
action="query",
context="The generated sentence has to be related with Argilla, a data annotation tool for AI engineers and domain experts.",
llm=llm,
input_batch_size=10,
output_mappings={"model_name": "model_name_query"},
)
让我们来分解一下上面的代码片段。
通过设置 `triplet=True`，我们指示任务生成一系列三元组，包括一个锚点、一个正向句子和一个负向句子。正如 Sentence Transformers 库的训练概述中所解释的，这种格式非常适合微调。
action="query" 参数是此任务的一个关键方面，因为它指导 LLM 为正向句子生成查询。这就是魔法发生的地方，我们的文档块被转化为有意义的查询。
为了进一步帮助模型，我们加入了 `context` 参数。当锚定句子缺乏足够的上下文时（这在简短的文档块中经常出现），这会向 LLM 提供额外信息。
最后，我们选择利用 `meta-llama/Meta-Llama-3-70B-Instruct` 模型的力量，通过 InferenceEndpointsLLM 组件。这个选择使我们能够利用该模型的能力，生成高质量的合成查询，最终提高我们模型的性能。
multiply_queries
通过使用 `GenerateSentencePair` 步骤，我们获得了与分块数量相同的训练样本，本例中为 251 个。然而，我们认识到这可能不足以微调一个能够准确捕捉我们特定用例细微差别的自定义模型。
为了克服这个限制，我们将使用另一个 LLM 来生成额外的查询。这将允许我们增加训练数据集的大小，为我们的模型提供更丰富的学习基础。
这让我们进入了流程的下一步：`MultipleQueries`，一个我们为了进一步扩充数据集而精心制作的自定义 `Task`。
multiply_queries = MultipleQueries(
name="multiply_queries",
num_queries=3,
system_prompt=(
"You are an AI assistant helping to generate diverse examples. Ensure the "
"generated queries are all in separated lines and preceded by a dash. "
"Do not generate anything else or introduce the task."
),
llm=llm,
input_batch_size=10,
input_mappings={"query": "positive"},
output_mappings={"model_name": "model_name_query_multiplied"},
)
现在，让我们深入了解我们自定义 `Task` 的配置，该任务旨在扩充我们的训练数据集。此任务的关键是 `num_queries` 参数，本例中设置为 3。这意味着我们将为每个示例生成三个额外的“正向”查询，从而有效地将我们的数据集大小增加四倍（假设某些示例可能不会成功）。
为了确保大型语言模型 (LLM) 保持正轨，我们精心设计了一个 `system_prompt`，为我们的指令提供清晰的指导。考虑到所选模型的强大功能和我们示例的简单性，我们不需要使用结构化生成技术。然而，在更复杂的场景中，这可能是一种有价值的方法。
对我们自定义 `Task` 的内部工作原理感到好奇吗？点击下面的下拉菜单探索完整的定义
MultipleQueries 定义
multiply_queries_template = (
"Given the following query:\n{original}\nGenerate {num_queries} similar queries by varying "
"the tone and the phrases slightly. "
"Ensure the generated queries are coherent with the original reference and relevant to the context of data annotation "
"and AI dataset development."
)
class MultipleQueries(Task):
system_prompt: Optional[str] = None
num_queries: int = 1
@property
def inputs(self) -> List[str]:
return ["query"]
def format_input(self, input: Dict[str, Any]) -> ChatType:
prompt = [
{
"role": "user",
"content": multiply_queries_template.format(
original=input["query"],
num_queries=self.num_queries
),
},
]
if self.system_prompt:
prompt.insert(0, {"role": "system", "content": self.system_prompt})
return prompt
@property
def outputs(self) -> List[str]:
return ["queries", "model_name"]
def format_output(
self, output: Union[str, None], input: Dict[str, Any]
) -> Dict[str, Any]:
queries = output.split("- ")
if len(queries) > self.num_queries:
queries = queries[1:]
queries = [q.strip() for q in queries]
return {"queries": queries}
merge_columns
当我们接近流程的最后阶段时，我们的重点转向了数据处理。我们的最终目标是创建一个精炼的数据集，由适合微调的三元组行组成。然而，在生成多个查询后，我们的数据集现在包含两个不同的列：`positive` 和 `queries`。`positive` 列包含原始查询作为单个字符串，而 `queries` 列存储一个字符串列表，表示为同一实体生成的额外查询。
为了将这两列合并成一个单一、连贯的列表，我们将使用 MergeColumns 步骤。这将使我们能够将原始查询与生成的查询结合起来，创建一个统一的列表。
merge_columns = MergeColumns(
name="merge_columns",
columns=["positive", "queries"],
output_column="positive"
)
expand_columns
最后，我们使用 ExpandColumns 将之前的正向列移动到不同的行。因此，每个 `positive` 查询将占据一个单独的行，而 `anchor` 和 `negative` 列将被复制以匹配扩展后的正向查询。这种数据操作将产生一个具有理想微调结构的数据集。
expand_columns = ExpandColumns(columns=["positive"])
点击下拉菜单查看完整的流水线定义
Distilabel 管道
from pathlib import Path
from typing import Any, Dict, List, Union, Optional
from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromHub
from distilabel.llms import InferenceEndpointsLLM
from distilabel.steps.tasks import GenerateSentencePair
from distilabel.steps.tasks.base import Task
from distilabel.steps.tasks.typing import ChatType
from distilabel.steps import ExpandColumns, CombineKeys
multiply_queries_template = (
"Given the following query:\n{original}\nGenerate {num_queries} similar queries by varying "
"the tone and the phrases slightly. "
"Ensure the generated queries are coherent with the original reference and relevant to the context of data annotation "
"and AI dataset development."
)
class MultipleQueries(Task):
system_prompt: Optional[str] = None
num_queries: int = 1
@property
def inputs(self) -> List[str]:
return ["query"]
def format_input(self, input: Dict[str, Any]) -> ChatType:
prompt = [
{
"role": "user",
"content": multiply_queries_template.format(
original=input["query"],
num_queries=self.num_queries
),
},
]
if self.system_prompt:
prompt.insert(0, {"role": "system", "content": self.system_prompt})
return prompt
@property
def outputs(self) -> List[str]:
return ["queries", "model_name"]
def format_output(
self, output: Union[str, None], input: Dict[str, Any]
) -> Dict[str, Any]:
queries = output.split("- ")
if len(queries) > self.num_queries:
queries = queries[1:]
queries = [q.strip() for q in queries]
return {"queries": queries}
with Pipeline(
name="embedding-queries",
description="Generate queries to train a sentence embedding model."
) as pipeline:
load_data = LoadDataFromHub(
name="load_data",
repo_id="plaguss/argilla_sdk_docs_raw_unstructured",
output_mappings={"chunks": "anchor"},
batch_size=10,
)
llm = InferenceEndpointsLLM(
model_id="meta-llama/Meta-Llama-3-70B-Instruct",
tokenizer_id="meta-llama/Meta-Llama-3-70B-Instruct",
)
generate_sentence_pair = GenerateSentencePair(
name="generate_sentence_pair",
triplet=True, # Generate positive and negative
action="query",
context="The generated sentence has to be related with Argilla, a data annotation tool for AI engineers and domain experts.",
llm=llm,
input_batch_size=10,
output_mappings={"model_name": "model_name_query"},
)
multiply_queries = MultipleQueries(
name="multiply_queries",
num_queries=3,
system_prompt=(
"You are an AI assistant helping to generate diverse examples. Ensure the "
"generated queries are all in separated lines and preceded by a dash. "
"Do not generate anything else or introduce the task."
),
llm=llm,
input_batch_size=10,
input_mappings={"query": "positive"},
output_mappings={"model_name": "model_name_query_multiplied"},
)
merge_columns = MergeColumns(
name="merge_columns",
columns=["positive", "queries"],
output_column="positive"
)
expand_columns = ExpandColumns(
columns=["positive"],
)
(
load_data
>> generate_sentence_pair
>> multiply_queries
>> merge_columns
>> expand_columns
)
if __name__ == "__main__":
pipeline_parameters = {
"generate_sentence_pair": {
"llm": {
"generation_kwargs": {
"temperature": 0.7,
"max_new_tokens": 512,
}
}
},
"multiply_queries": {
"llm": {
"generation_kwargs": {
"temperature": 0.7,
"max_new_tokens": 512,
}
}
}
}
distiset = pipeline.run(
parameters=pipeline_parameters
)
distiset.push_to_hub("plaguss/argilla_sdk_docs_queries")
在 Argilla 中探索数据集
现在我们已经生成了数据集，是时候使用 Argilla 深入研究并根据需要进行优化了。要开始，请查看我们的 argilla_datasets.ipynb 笔记本，它提供了如何将数据集上传到 Argilla 的分步指南。
如果您还没有设置 Argilla 实例，别担心！请遵循我们文档中易于理解的指南，在 Hugging Face Space 上创建一个 Argilla 实例。一旦您的 Space 启动并运行，只需通过更新 `api_url` 指向您的 Space 即可连接到它
import argilla as rg
client = rg.Argilla(
api_url="https://plaguss-argilla-sdk-chatbot.hf.space",
api_key="YOUR_API_KEY"
)
一个包含技术文档块的 Argilla 数据集
在您的 Argilla 实例启动并运行后，让我们进入下一步：为您的数据集配置 `Settings`。好消息是，我们将创建的默认 `Settings` 应该可以无缝地适用于您的特定用例，无需进一步调整
settings = rg.Settings(
guidelines="Review the chunks of docs.",
fields=[
rg.TextField(
name="filename",
title="Filename where this chunk was extracted from",
use_markdown=False,
),
rg.TextField(
name="chunk",
title="Chunk from the documentation",
use_markdown=False,
),
],
questions=[
rg.LabelQuestion(
name="good_chunk",
title="Does this chunk contain relevant information?",
labels=["yes", "no"],
)
],
)
让我们仔细看看我们创建的数据集结构。我们将检查 `filename` 和 `chunk` 字段，它们分别包含解析的文件名和生成的块。为了进一步增强我们的数据集，我们可以定义一个简单的标签问题 `good_chunk`，这允许我们手动标记每个块是否有用。这种人在回路的方法使我们能够改进我们的自动化生成过程。有了这些基本要素，我们现在准备好创建我们的数据集了
dataset = rg.Dataset(
name="argilla_sdk_docs_raw_unstructured",
settings=settings,
client=client,
)
dataset.create()
现在，让我们检索我们之前在 Hugging Face Hub 上创建的数据集。还记得我们在数据分块部分生成的数据集吗？我们将下载该数据集并提取我们需要的核心列以便继续
from datasets import load_dataset
data = (
load_dataset("plaguss/argilla_sdk_docs_raw_unstructured", split="train")
.select_columns(["filename", "chunks"])
.to_list()
)
我们已经到达了最后的里程碑！为了将所有东西整合在一起，让我们将记录日志到 Argilla。这将允许我们在 Argilla 界面中可视化我们的数据集，提供一种清晰直观的方式来探索和与我们的数据互动
dataset.records.log(records=data, mapping={"filename": "filename", "chunks": "chunk"})
这些是您可能期望看到的示例类型
一个包含三元组的 Argilla 数据集，用于微调嵌入模型
现在，我们可以使用我们在上一节中生成的准备好进行微调的数据集重复此过程。幸运的是，过程很简单：只需下载相关的数据集并使用其指定的名称上传到 Argilla 即可。有关详细的演练，请参阅 Jupyter 笔记本，其中包含所有必要的说明
settings = rg.Settings(
guidelines="Review the chunks of docs.",
fields=[
rg.TextField(
name="anchor",
title="Anchor (Chunk from the documentation).",
use_markdown=False,
),
rg.TextField(
name="positive",
title="Positive sentence that queries the anchor.",
use_markdown=False,
),
rg.TextField(
name="negative",
title="Negative sentence that may use similar words but has content unrelated to the anchor.",
use_markdown=False,
),
],
questions=[
rg.LabelQuestion(
name="is_positive_relevant",
title="Is the positive query relevant?",
labels=["yes", "no"],
),
rg.LabelQuestion(
name="is_negative_irrelevant",
title="Is the negative query irrelevant?",
labels=["yes", "no"],
)
],
)
让我们仔细看看我们数据集的结构，它由三个基本的 TextFields 组成：`anchor`、`positive` 和 `negative`。`anchor` 字段表示文本块本身，而 `positive` 字段包含一个可以使用锚文本作为参考来回答的查询。相比之下，`negative` 字段包含一个不相关的查询，作为三元组中的负面示例。正面和负面问题在帮助我们的模型区分这些示例并学习有效的嵌入方面起着至关重要的作用。
下图展示了一个例子
到目前为止我们建立的数据集设置主要用于探索我们的数据集，但我们可以更进一步。通过自定义这些设置，我们可以识别和纠正错误的示例，提高生成问题的质量，并迭代改进我们的数据集，使其更适合我们的需求。
一个用于跟踪聊天机器人对话的 Argilla 数据集
现在，让我们创建我们的最终数据集，它将专门用于跟踪用户与我们聊天机器人的互动。*注意*：您可能希望在完成 Gradio 应用后再重新访问此部分，因为它将提供对上下文更全面的理解。现在，让我们看看这个数据集的 `Settings`
settings_chatbot_interactions = rg.Settings(
guidelines="Review the user interactions with the chatbot.",
fields=[
rg.TextField(
name="instruction",
title="User instruction",
use_markdown=True,
),
rg.TextField(
name="response",
title="Bot response",
use_markdown=True,
),
],
questions=[
rg.LabelQuestion(
name="is_response_correct",
title="Is the response correct?",
labels=["yes", "no"],
),
rg.LabelQuestion(
name="out_of_guardrails",
title="Did the model answered something out of the ordinary?",
description="If the model answered something unrelated to Argilla SDK",
labels=["yes", "no"],
),
rg.TextQuestion(
name="feedback",
title="Let any feedback here",
description="This field should be used to report any feedback that can be useful",
required=False
),
],
metadata=[
rg.TermsMetadataProperty(
name="conv_id",
title="Conversation ID",
),
rg.IntegerMetadataProperty(
name="turn",
min=0,
max=100,
title="Conversation Turn",
)
]
)
在这个数据集中，我们将定义两个基本字段：`instruction` 和 `response`。`instruction` 字段将存储初始查询，如果对话被扩展，它将包含到那时为止的整个对话历史。另一方面，`response` 字段将保存聊天机器人的最新响应。为了便于评估和反馈，我们将包括三个问题：一个评估响应的正确性，另一个确定模型是否偏离主题，以及一个可选字段供用户提供对响应的反馈。此外，我们将包括两个元数据属性，以便对对话进行筛选和分析：一个唯一的对话 ID 和对话中的轮次编号。
下图展示了一个例子
一旦我们的聊天机器人获得了大量的用户参与，这个数据集就可以作为宝贵的资源来改进和增强我们的模型，使我们能够根据真实世界的互动迭代和提高其性能。
微调嵌入模型
现在我们的自定义嵌入模型数据集已经准备好了，是时候进入训练过程了。
为了指导我们完成这一步，我们将参考 train_embedding.ipynb 笔记本，该笔记本的灵感来自 Philipp Schmid 关于为 RAG 微调嵌入模型的博客文章。虽然该博客文章对整个过程进行了全面的概述，但我们将重点关注我们用例中特定的关键差异和细微之处。
要更深入地了解底层决策和详细的演练，请务必查看原始博客文章并审阅笔记本以获取分步解释。
准备嵌入数据集
我们将首先下载数据集并选择核心列，这些列恰好已经与 Sentence Transformers 期望的命名约定保持一致。接下来，我们将为每个样本添加一个唯一的 id 列，并将数据集分为训练集和测试集，其中 90% 用于训练，10% 用于测试。最后，我们将格式化后的数据集转换为 JSON 文件，准备好输入训练器进行模型微调。
from datasets import load_dataset
# Load dataset from the hub
dataset = (
load_dataset("plaguss/argilla_sdk_docs_queries", split="train")
.select_columns(["anchor", "positive", "negative"]) # Select the relevant columns
.add_column("id", range(len(dataset))) # Add an id column to the dataset
.train_test_split(test_size=0.1) # split dataset into a 10% test set
)
# Save datasets to disk
dataset["train"].to_json("train_dataset.json", orient="records")
dataset["test"].to_json("test_dataset.json", orient="records")
加载基线模型
数据集准备就绪后，我们现在可以加载基线模型，它将作为我们微调过程的基础。我们将使用参考博客文章中采用的相同模型，以确保我们的自定义嵌入模型开发有一个一致的起点。
from sentence_transformers import SentenceTransformerModelCardData, SentenceTransformer
model = SentenceTransformer(
"BAAI/bge-base-en-v1.5",
model_card_data=SentenceTransformerModelCardData(
language="en",
license="apache-2.0",
model_name="BGE base ArgillaSDK Matryoshka",
),
)
定义损失函数
考虑到我们数据集的结构，我们将利用 `TripletLoss` 函数，它更适合处理我们的 `(anchor-positive-negative)` 三元组。此外，我们还会将它与 `MatryoshkaLoss` 结合使用，这是一个强大的损失函数，已经显示出有希望的结果（要深入了解 `MatryoshkaLoss`，请查看这篇文章）。
from sentence_transformers.losses import MatryoshkaLoss, TripletLoss
inner_train_loss = TripletLoss(model)
train_loss = MatryoshkaLoss(
model, inner_train_loss, matryoshka_dims=[768, 512, 256, 128, 64]
)
定义训练策略
现在我们已经有了基线模型和损失函数，是时候定义将指导微调过程的训练参数了。由于这项工作是在 Apple M2 Pro 上完成的，我们需要进行一些调整以确保训练过程顺利进行。
为了适应我们机器有限的资源，我们将减少 `per_device_train_batch_size` 和 `per_device_eval_batch_size`，与原始博客文章相比。此外，我们需要移除 `tf32` 和 `bf16` 精度选项，因为它们在此设备上不受支持。此外，我们将换掉 `adamw_torch_fused` 优化器，该优化器可以在 Google Colab 笔记本中使用以加快训练速度。通过进行这些修改，我们将能够微调我们的模型。
from sentence_transformers import SentenceTransformerTrainingArguments
# Define training arguments
args = SentenceTransformerTrainingArguments(
output_dir="bge-base-argilla-sdk-matryoshka", # output directory and hugging face model ID
num_train_epochs=3, # number of epochs
per_device_train_batch_size=8, # train batch size
gradient_accumulation_steps=4, # for a global batch size of 512
per_device_eval_batch_size=4, # evaluation batch size
warmup_ratio=0.1, # warmup ratio
learning_rate=2e-5, # learning rate, 2e-5 is a good value
lr_scheduler_type="cosine", # use constant learning rate scheduler
eval_strategy="epoch", # evaluate after each epoch
save_strategy="epoch", # save after each epoch
logging_steps=5, # log every 10 steps
save_total_limit=1, # save only the last 3 models
load_best_model_at_end=True, # load the best model when training ends
metric_for_best_model="eval_dim_512_cosine_ndcg@10", # optimizing for the best ndcg@10 score for the 512 dimension
)
训练并保存最终模型
from sentence_transformers import SentenceTransformerTrainer
trainer = SentenceTransformerTrainer(
model=model, # bg-base-en-v1
args=args, # training arguments
train_dataset=train_dataset.select_columns(
["anchor", "positive", "negative"]
), # training dataset
loss=train_loss,
evaluator=evaluator,
)
# Start training, the model will be automatically saved to the hub and the output directory
trainer.train()
# Save the best model
trainer.save_model()
# Push model to hub
trainer.model.push_to_hub("bge-base-argilla-sdk-matryoshka")
就是这样！我们可以看看这个新模型：plaguss/bge-base-argilla-sdk-matryoshka。请仔细查看数据集卡片，其中包含了关于我们模型的宝贵见解和信息。
但这还不是全部！在下一节中，我们将对我们的模型进行测试，看看它的实际效果。
向量数据库
到目前为止，我们已经取得了显著的进展，为我们的 RAG 聊天机器人创建了一个数据集并微调了一个模型。现在，是时候构建向量数据库了，它将使我们的聊天机器人能够高效地存储和检索相关信息。
在选择向量数据库时，有许多可用的替代方案。为了保持简单明了，我们将使用 lancedb，这是一个轻量级的嵌入式数据库，不需要服务器，类似于 SQLite。正如我们将看到的，lancedb 允许我们创建一个简单的文件来存储我们的嵌入，使其易于移动和快速检索数据，这非常适合我们的用例。
要跟上进度，请参考附带的笔记本：vector_db.ipynb。在这个笔记本中，我们将深入探讨构建和使用我们的向量数据库的细节。
连接到数据库
安装完依赖项后，让我们实例化数据库
import lancedb
# Create a database locally called `lancedb`
db = lancedb.connect("./lancedb")
当我们执行代码时，当前工作目录中应该会出现一个新文件夹，这标志着我们的向量数据库已成功创建。
实例化微调后的模型
现在我们的向量数据库已经建立，是时候加载我们微调过的模型了。我们将利用 `sentence-transformers` 注册表来加载模型，释放其能力并为行动做好准备。
import torch
from lancedb.embeddings import get_registry
model_name = "plaguss/bge-base-argilla-sdk-matryoshka"
device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
model = get_registry().get("sentence-transformers").create(name=model_name, device=device)
创建包含文档块的表
在我们加载了微调模型之后，我们准备创建将存储我们嵌入的表格。为了定义这个表格的模式，我们将使用一个 `LanceModel`，类似于 `pydantic.BaseModel`，来为我们的 `Docs` 实体创建一个健壮的表示。
from lancedb.pydantic import LanceModel, Vector
class Docs(LanceModel):
query: str = model.SourceField()
text: str = model.SourceField()
vector: Vector(model.ndims()) = model.VectorField()
table_name = "docs"
table = db.create_table(table_name, schema=Docs)
之前的代码片段为创建一个包含三个基本列的表格奠定了基础
query: 专用于存储合成查询text: 存放分块的文档文本vector: 与我们微调模型的维度相关联，准备存储嵌入
有了这个表结构，我们现在可以与表进行交互了。
填充表格
在建立了我们的表格结构之后，我们现在准备用数据来填充它。让我们加载最终的数据集，其中包含查询，并将它们连同相应的嵌入一起摄取到我们的数据库中。这个关键步骤将使我们的向量数据库充满活力，使我们的聊天机器人能够高效地存储和检索相关信息。
ds = load_dataset("plaguss/argilla_sdk_docs_queries", split="train")
batch_size = 50
for batch in tqdm.tqdm(ds.iter(batch_size), total=len(ds) // batch_size):
embeddings = model.generate_embeddings(batch["positive"])
df = pd.DataFrame.from_dict({"query": batch["positive"], "text": batch["anchor"], "vector": embeddings})
table.add(df)
在之前的代码片段中，我们以批处理方式遍历数据集，使用我们微调的模型为 `positive` 列中的合成查询生成嵌入。然后，我们创建了一个 Pandas 数据框，以包含 `query`、`text` 和 `vector` 列。这个数据框将 `positive` 和 `anchor` 列与新生成的嵌入分别结合起来。
现在，让我们来测试一下我们的向量数据库！对于一个示例查询，“我如何获取当前用户？”（使用 Argilla SDK），我们将使用我们的自定义嵌入模型生成嵌入。然后，我们将利用 `cosine` 度量来衡量相似度，在我们的表格中搜索前 3 个最相似的出现。最后，我们将提取相关的 `text` 列，它对应于最匹配我们查询的文档块。
query = "How can I get the current user?"
embedded_query = model.generate_embeddings([query])
retrieved = (
table
.search(embedded_query[0])
.metric("cosine")
.limit(3)
.select(["text"]) # Just grab the chunk to use for context
.to_list()
)
点击查看结果
这将是结果
>>> retrieved
[{'text': 'python\nuser = client.users("my_username")\n\nThe current user of the rg.Argilla client can be accessed using the me attribute:\n\npython\nclient.me\n\nClass Reference\n\nrg.User\n\n::: argilla_sdk.users.User\n options:\n heading_level: 3',
'_distance': 0.1881886124610901},
{'text': 'python\nuser = client.users("my_username")\n\nThe current user of the rg.Argilla client can be accessed using the me attribute:\n\npython\nclient.me\n\nClass Reference\n\nrg.User\n\n::: argilla_sdk.users.User\n options:\n heading_level: 3',
'_distance': 0.20238929986953735},
{'text': 'Retrieve a user\n\nYou can retrieve an existing user from Argilla by accessing the users attribute on the Argilla class and passing the username as an argument.\n\n```python\nimport argilla_sdk as rg\n\nclient = rg.Argilla(api_url="", api_key="")\n\nretrieved_user = client.users("my_username")\n```',
'_distance': 0.20401990413665771}]
>>> print(retrieved[0]["text"])
python
user = client.users("my_username")
The current user of the rg.Argilla client can be accessed using the me attribute:
python
client.me
Class Reference
rg.User
::: argilla_sdk.users.User
options:
heading_level: 3
让我们深入研究数据集的第一行，看看能发现什么。乍一看，它似乎包含了与查询相关的信息，这正是我们所期望的。要获取当前用户，我们可以使用 `client.me` 方法。然而，我们也注意到一些无关的内容，这很可能是由于所采用的分块策略造成的。这个策略虽然有效，但可以从一些改进中受益。通过在 Argilla 中审查数据集，我们可以更深入地了解如何优化我们的分块方法，最终得到一个更精简的数据集。不过，就目前而言，这似乎是一个很好的起点。
将数据库存储在 Hugging Face Hub 上
现在我们有了一个数据库，我们将把它作为另一个工件存储在我们的数据集仓库中。您可以访问该仓库以找到可以帮助我们的函数，但它就像运行以下函数一样简单
import Path
import os
local_dir = Path.home() / ".cache/argilla_sdk_docs_db"
upload_database(
local_dir / "lancedb",
repo_id="plaguss/argilla_sdk_docs_queries",
token=os.getenv("HF_API_TOKEN")
)
我们数据库存储之旅的最后一步仅需一个命令！通过运行该函数，我们将创建一个名为 `lancedb.tar.gz` 的全新文件，它将整齐地打包我们的向量数据库。您可以在 Hugging Face Hub 上的 plaguss/argilla_sdk_docs_queries 仓库中预览生成的文件，它与其他重要文件一起存储。
db_path = download_database(repo_id)
关键时刻到了！我们的数据库成功下载后，我们现在可以验证一切是否正常。默认情况下，文件将存储在 `Path.home() / ".cache/argilla_sdk_docs_db"`，但可以轻松自定义。我们可以再次连接到它并检查一切是否按预期工作
db = lancedb.connect(str(db_path))
table = db.open_table(table_name)
query = "how can I delete users?"
retrieved = (
table
.search(query)
.metric("cosine")
.limit(1)
.to_pydantic(Docs)
)
for d in retrieved:
print("======\nQUERY\n======")
print(d.query)
print("======\nDOC\n======")
print(d.text)
# ======
# QUERY
# ======
# Is it possible to remove a user from Argilla by utilizing the delete function on the User class?
# ======
# DOC
# ======
# Delete a user
# You can delete an existing user from Argilla by calling the delete method on the User class.
# ```python
# import argilla_sdk as rg
# client = rg.Argilla(api_url="", api_key="")
# user_to_delete = client.users('my_username')
# deleted_user = user_to_delete.delete()
# ```
用于文档检索的数据库已经完成，接下来让我们来做应用程序吧！
创建我们的聊天机器人
我们聊天机器人的所有部件都已准备就绪；我们需要将它们连接起来，并在一个界面中提供。
Gradio 应用
让我们把 RAG 应用带入现实！使用 gradio，我们可以毫不费力地创建聊天机器人应用。在这种情况下，我们将设计一个简单而有效的界面来展示我们聊天机器人的能力。要查看应用的实际效果，请参阅 GitHub 上 Argilla SDK Chatbot 仓库中的 app.py 脚本。
在我们深入探讨构建聊天机器人应用的细节之前，让我们先退一步，欣赏一下最终的成果。仅用几行代码，我们就成功创建了一个用户友好的界面，让我们的 RAG 聊天机器人栩栩如生。
import gradio as gr
gr.ChatInterface(
chatty,
chatbot=gr.Chatbot(height=600),
textbox=gr.Textbox(placeholder="Ask me about the new argilla SDK", container=False, scale=7),
title="Argilla SDK Chatbot",
description="Ask a question about Argilla SDK",
theme="soft",
examples=[
"How can I connect to an argilla server?",
"How can I access a dataset?",
"How can I get the current user?"
],
cache_examples=True,
retry_btn=None,
).launch()
就是这样！如果您渴望了解更多关于创建自己的聊天机器人的信息，请务必查看 Gradio 的优秀指南使用 Gradio 快速创建聊天机器人。这是一个知识的宝库，能让您在短时间内构建出自己的聊天机器人。
现在，让我们更深入地探讨我们的 `app.py` 脚本的内部工作原理。我们将分解关键组件，重点关注让我们的聊天机器人活起来的基本元素。为了保持简洁，我们将略过一些更精细的细节。
首先，让我们检查 `Database` 类，它是我们聊天机器人知识和功能的支柱。这个组件在存储和检索驱动我们聊天机器人对话的数据方面起着至关重要的作用
点击查看 Database 类
class Database:
def __init__(self, settings: Settings) -> None:
self.settings = settings
self._table: lancedb.table.LanceTable = self.get_table_from_db()
def get_table_from_db(self) -> lancedb.table.LanceTable:
lancedb_db_path = self.settings.LOCAL_DIR / self.settings.LANCEDB
if not lancedb_db_path.exists():
lancedb_db_path = download_database(
self.settings.REPO_ID,
lancedb_file=self.settings.LANCEDB_FILE_TAR,
local_dir=self.settings.LOCAL_DIR,
token=self.settings.TOKEN,
)
db = lancedb.connect(str(lancedb_db_path))
table = db.open_table(self.settings.TABLE_NAME)
return table
def retrieve_doc_chunks(
self, query: str, limit: int = 12, hard_limit: int = 4
) -> str:
# Embed the query to use our custom model instead of the default one.
embedded_query = model.generate_embeddings([query])
field_to_retrieve = "text"
retrieved = (
self._table.search(embedded_query[0])
.metric("cosine")
.limit(limit)
.select([field_to_retrieve]) # Just grab the chunk to use for context
.to_list()
)
return self._prepare_context(retrieved, hard_limit)
@staticmethod
def _prepare_context(retrieved: list[dict[str, str]], hard_limit: int) -> str:
# We have repeated questions (up to 4) for a given chunk, so we may get repeated chunks.
# Request more than necessary and filter them afterwards
responses = []
unique_responses = set()
for item in retrieved:
chunk = item["text"]
if chunk not in unique_responses:
unique_responses.add(chunk)
responses.append(chunk)
context = ""
for i, item in enumerate(responses[:hard_limit]):
if i > 0:
context += "\n\n"
context += f"---\n{item}"
return context
随着我们的 `Database` 类的就位，我们成功地在我们聊天机器人的对话流和存储在我们数据库中的知识之间架起了桥梁。现在，让我们把一切都整合起来！一旦我们下载了我们的嵌入模型（脚本会自动完成），我们就可以实例化 `Database` 类，从而有效地将我们的数据库部署到期望的位置——在这种情况下，是我们的 Hugging Face Space。
这标志着我们聊天机器人开发旅程中的一个重要里程碑。随着我们的数据库集成并准备就绪，我们距离释放聊天机器人的全部潜力仅一步之遥。
database = Database(settings=settings) # The settings can be seen in the following snippet
context = database.retrieve_doc_chunks("How can I delete a user?", limit=2, hard_limit=1)
>>> print(context)
# ---
# Delete a user
# You can delete an existing user from Argilla by calling the delete method on the User class.
# ```python
# import argilla_sdk as rg
# client = rg.Argilla(api_url="", api_key="")
# user_to_delete = client.users('my_username')
# deleted_user = user_to_delete.delete()
# ```
点击查看 Settings 类
@dataclass
class Settings:
LANCEDB: str = "lancedb"
LANCEDB_FILE_TAR: str = "lancedb.tar.gz"
TOKEN: str = os.getenv("HF_API_TOKEN")
LOCAL_DIR: Path = Path.home() / ".cache/argilla_sdk_docs_db"
REPO_ID: str = "plaguss/argilla_sdk_docs_queries"
TABLE_NAME: str = "docs"
MODEL_NAME: str = "plaguss/bge-base-argilla-sdk-matryoshka"
DEVICE: str = (
"mps"
if torch.backends.mps.is_available()
else "cuda"
if torch.cuda.is_available()
else "cpu"
)
MODEL_ID: str = "meta-llama/Meta-Llama-3-70B-Instruct"
拼图的最后一块现在已经就位——我们的数据库已经准备好为我们的聊天机器人对话提供动力。接下来，我们需要准备我们的模型来处理涌入的用户查询。这就是推理端点发挥作用的地方。这些专用端点提供了一种无缝的方式来部署和管理我们的模型，确保它随时准备响应用户输入。
幸运的是，使用推理端点非常简单，这要归功于 `huggingface_hub` 库中的 inference client。
def get_client_and_tokenizer(
model_id: str = settings.MODEL_ID, tokenizer_id: Optional[str] = None
) -> tuple[InferenceClient, AutoTokenizer]:
if tokenizer_id is None:
tokenizer_id = model_id
client = InferenceClient()
base_url = client._resolve_url(model=model_id, task="text-generation")
# Note: We could move to the AsyncClient
client = InferenceClient(model=base_url, token=os.getenv("HF_API_TOKEN"))
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
return client, tokenizer
# Load the client and tokenizer
client, tokenizer = get_client_and_tokenizer()
在我们的组件就位后，我们已经到了准备将被馈送到我们客户端的提示的阶段。这个提示将作为激发我们机器学习模型魔力的输入，引导它生成一个既准确又信息丰富的响应，同时避免回答不相关的问题。在本节中，我们将深入探讨制作一个结构良好的提示的细节，为我们的模型成功做好准备。`prepare_input` 函数将准备对话，应用提示和聊天模板以传递给模型。
def prepare_input(message: str, history: list[tuple[str, str]]) -> str:
# Retrieve the context from the database
context = database.retrieve_doc_chunks(message)
# Prepare the conversation for the model.
conversation = []
for human, bot in history:
conversation.append({"role": "user", "content": human})
conversation.append({"role": "assistant", "content": bot})
conversation.insert(0, {"role": "system", "content": SYSTEM_PROMPT})
conversation.append(
{
"role": "user",
"content": ARGILLA_BOT_TEMPLATE.format(message=message, context=context),
}
)
return tokenizer.apply_chat_template(
[conversation],
tokenize=False,
add_generation_prompt=True,
)[0]
此函数将接受两个参数：`message` 和 `history`，这得益于 gradio 的 ChatInterface。它将从数据库中获取文档片段以帮助 LLM 生成响应，并准备好要传递给我们 `LLM` 模型的提示。
点击查看系统提示和机器人模板
这些是使用的 `system_prompt` 和提示模板。它们深受 Weights and Biases 的 `wandbot` 的启发。
SYSTEM_PROMPT = """\
You are a support expert in Argilla SDK, whose goal is help users with their questions.
As a trustworthy expert, you must provide truthful answers to questions using only the provided documentation snippets, not prior knowledge.
Here are guidelines you must follow when responding to user questions:
##Purpose and Functionality**
- Answer questions related to the Argilla SDK.
- Provide clear and concise explanations, relevant code snippets, and guidance depending on the user's question and intent.
- Ensure users succeed in effectively understanding and using Argilla's features.
- Provide accurate responses to the user's questions.
**Specificity**
- Be specific and provide details only when required.
- Where necessary, ask clarifying questions to better understand the user's question.
- Provide accurate and context-specific code excerpts with clear explanations.
- Ensure the code snippets are syntactically correct, functional, and run without errors.
- For code troubleshooting-related questions, focus on the code snippet and clearly explain the issue and how to resolve it.
- Avoid boilerplate code such as imports, installs, etc.
**Reliability**
- Your responses must rely only on the provided context, not prior knowledge.
- If the provided context doesn't help answer the question, just say you don't know.
- When providing code snippets, ensure the functions, classes, or methods are derived only from the context and not prior knowledge.
- Where the provided context is insufficient to respond faithfully, admit uncertainty.
- Remind the user of your specialization in Argilla SDK support when a question is outside your domain of expertise.
- Redirect the user to the appropriate support channels - Argilla [community](https://join.slack.com/t/rubrixworkspace/shared_invite/zt-whigkyjn-a3IUJLD7gDbTZ0rKlvcJ5g) when the question is outside your capabilities or you do not have enough context to answer the question.
**Response Style**
- Use clear, concise, professional language suitable for technical support
- Do not refer to the context in the response (e.g., "As mentioned in the context...") instead, provide the information directly in the response.
**Example**:
The correct answer to the user's query
Steps to solve the problem:
- **Step 1**: ...
- **Step 2**: ...
...
Here's a code snippet
```python
# Code example
...
```
**Explanation**:
- Point 1
- Point 2
...
"""
ARGILLA_BOT_TEMPLATE = """\
Please provide an answer to the following question related to Argilla's new SDK.
You can make use of the chunks of documents in the context to help you generating the response.
## Query:
{message}
## Context:
{context}
"""
我们已经达到了我们对话式人工智能系统的顶点：`chatty` 函数。这个函数扮演着编排者的角色，将我们到目前为止构建的各种组件整合在一起。它的主要职责是调用 `prepare_input` 函数，该函数会精心制作将被传递给客户端的提示。然后，我们会在文本生成时逐段产生文本流，一旦响应完成，对话历史将被保存，为我们提供宝贵的资源来审查和改进我们的模型，确保它在每次迭代中持续进步。
def chatty(message: str, history: list[tuple[str, str]]) -> Generator[str, None, None]:
prompt = prepare_input(message, history)
partial_response = ""
for token_stream in client.text_generation(prompt=prompt, **client_kwargs):
partial_response += token_stream
yield partial_response
global conv_id
new_conversation = len(history) == 0
if new_conversation:
conv_id = str(uuid.uuid4())
else:
history.append((message, None))
# Register to argilla dataset
argilla_dataset.records.log(
[
{
"instruction": create_chat_html(history) if history else message,
"response": partial_response,
"conv_id": conv_id,
"turn": len(history)
},
]
)
关键时刻到了！我们的应用程序现在已经准备好接受测试了。要查看它的实际效果，只需在您的本地环境中运行 `python app.py`。但在您这样做之前，请确保您可以访问部署在推理端点上的模型。在这个例子中，我们使用的是强大的 Llama 3 70B 模型，但您可以随意尝试其他适合您需求的模型。通过调整模型和微调应用程序，您可以释放其全部潜力，探索人工智能开发的新可能性。
在 Hugging Face Spaces 上部署聊天机器人应用
现在我们的应用程序已经启动并运行，是时候与世界分享它了！为了部署我们的应用程序并使其可供他人访问，我们将遵循 Gradio 的指南中概述的步骤来分享您的应用程序。我们选择的托管平台是 Hugging Face Spaces，这是一个展示人工智能驱动项目的绝佳工具。
首先，我们需要在我们的仓库中添加一个 `requirements.txt` 文件，其中列出了运行我们的应用程序所需的依赖项。这是确保我们的应用程序可以轻松复制和部署的关键一步。您可以在 Hugging Face Spaces spaces dependencies 中了解有关管理依赖项的更多信息。
接下来，我们需要将我们的 Hugging Face API 令牌添加为秘密，请遵循本指南中的说明。这将允许我们的应用程序与 Hugging Face 生态系统进行身份验证。
一旦我们上传了我们的 `app.py` 文件，我们的 Space 将被构建，我们将能够通过以下链接访问我们的应用程序：
https://huggingface.co/spaces/plaguss/argilla-sdk-chatbot-space
请在此处查看我们的示例 Space 文件，了解这一切是如何组合在一起的。通过遵循这些步骤，您将能够与世界分享您自己的人工智能驱动的应用程序，并与 Hugging Face 社区中的其他人合作。
体验我们的聊天机器人
我们现在可以对聊天机器人进行测试。我们提供了一些默认查询供您开始，但您也可以随意尝试自己的问题。例如，您可以问：`新 SDK 中的 Settings 是什么？`
从下面的截图中您可以看到，我们的聊天机器人已经准备好为您的查询提供有用的响应
但这还不是全部！您还可以挑战我们的聊天机器人为特定数据集生成设置，就像我们在本教程前面创建的那个一样。例如，您可以让它为旨在微调嵌入模型的数据集建议设置，类似于我们在一个包含三元组以微调嵌入模型的 Argilla 数据集部分中探讨的那个。
看看下面的截图，看看我们的聊天机器人如何回应这类查询。
来吧，提出您的问题，看看我们的聊天机器人能提供什么样的见解！
下一步
在本教程中，我们成功构建了一个聊天机器人，它可以为有关 Argilla SDK 及其应用的问题提供有用的回答。通过利用 Llama 3 70B 和 Gradio 的强大功能，我们创建了一个用户友好的界面，可以帮助开发人员理解如何处理数据集和微调嵌入模型。
然而，我们的聊天机器人只是一个起点，我们还有很多方法可以改进和扩展其功能。以下是一些可以着手的下一步工作：
改进分块策略：尝试不同的分块策略、参数和大小，以优化聊天机器人的性能和响应质量。
实现去重和过滤：在训练数据集中添加去重和过滤机制，以移除重复和无关信息，确保聊天机器人提供准确简洁的回答。
为响应包含来源：通过在聊天机器人的响应中包含相关文档和来源的链接，增强其功能，让用户可以更深入地探讨主题并进一步探索。
通过解决这些领域的问题，我们可以将我们的聊天机器人提升到一个新的水平，使其成为使用 Argilla SDK 的开发人员更有价值的资源。可能性是无穷的，我们很期待看到这个项目未来的发展。敬请期待未来的更新和改进！
