在 FiftyOne 中新增了对向量搜索引擎和文本相似度查询的本地支持。这样,用户可以通过简单的自然语言查询,在他们的庞大数据集中(通常包含数百万甚至数千万个样本)找到最相关的图像。
这让我们处于一个有趣的境地:使用开源的 FiftyOne 的人现在可以通过自然语言查询轻松搜索数据集,但是在使用我们的文档时仍然需要传统的关键词搜索。
我们有大量的文档,这既有好处也有坏处。作为一个用户,我有时候发现,由于文档数量庞大,要准确找到我需要的信息需要花费很多时间,其实我不愿意浪费这些时间。
本篇文章的所有代码:https://github.com/voxel51/fiftyone-docs-search
01 将文档转换为统一格式
我公司的文档都以 HTML 文档的形式托管在 https://docs.voxel51.com 上。一个自然的起点是使用 Python 的 requests 库下载这些文档,并使用 BeautifulSoup 解析文档。
但是,作为一名开发人员(也是我们文档的作者之一),我觉得我可以做得更好。我已经在本地计算机上有一个克隆的 GitHub 存储库,其中包含用于生成 HTML 文档的所有原始文件。我们的一些文档使用 Sphinx ReStructured Text(RST)编写,其他文档,如教程,则是从 Jupyter 笔记本转换为 HTML 格式。
我错误地认为,我越接近 RST 和 Jupyter 文件的原始文本,事情就会变得更简单。
RST
在 RST 文档中,sections 通过仅包含 =、- 或_字符串的行进行划分。例如,这是 FiftyOne 用户指南中的一个文档,其中包含了这三种划分符号:
然后,我可以删除所有的 RST 关键词,例如 toctree、code-block 和 button_link(还有很多其他关键词),以及伴随关键词、新块的开头或块描述符的:、:: 和 …。
处理链接也很容易:
no_links_section = re.sub(r”<[^>]+>_?“,””, section)
当我想要从 RST 文件中提取 section 的锚点时,事情开始变得棘手起来。我们的许多节都明确指定了锚点,而其他 section 则在转换为 HTML 时留待推断。
这是一个例子:
… _brain-embeddings-visualization:
Visualizing embeddings
The FiftyOne Brain provides a powerful
:meth:compute_visualization() <fiftyone.brain.compute_visualization>
method
that you can use to generate low-dimensional representations of the samples
and/or individual objects in your datasets.
These representations can be visualized natively in the App’s
:ref:Embeddings panel <app-embeddings-panel>
, where you can interactively
select points of interest and view the corresponding samples/labels of interest
in the :ref:Samples panel <app-samples-panel>
, and vice versa.
… image:: /images/brain/brain-mnist.png
:alt: mnist
:align: center
There are two primary components to an embedding visualization: the method used
to generate the embeddings, and the dimensionality reduction method used to
compute a low-dimensional representation of the embeddings.
Embedding methods
The embeddings
and model
parameters of
:meth:compute_visualization() <fiftyone.brain.compute_visualization>
support a variety of ways to generate embeddings for your data:
在我们的用户指南文档中的 brain.rst 文件(上面部分的一部分)中,可视化嵌入(Visualizing embeddings)部分有一个指定为 #brain-embeddings-visualization 的锚点,使用了 … _brain-embeddings-visualization: 进行定义。然而,紧随其后的 Embedding 方法子部分则使用自动生成的锚点。
另一个很快出现的困难是如何处理 RST 中的表格。列表表格相对来说较为简单。例如,这是我们的查看阶段速查表中的一个列表表格:
… list-table::
- :meth:
match() <fiftyone.core.collections.SampleCollection.match>
- :meth:
match_frames() <fiftyone.core.collections.SampleCollection.match_frames>
- :meth:
match_labels() <fiftyone.core.collections.SampleCollection.match_labels>
- :meth:
match_tags() <fiftyone.core.collections.SampleCollection.match_tags>
另一方面,Grid tables 可能会很快变得混乱。它们为文档编写者提供了很大的灵活性,但这种灵活性也使得解析它们变得麻烦。来看看我们的 Filtering cheat sheet 中的这个表格吧:
±----------------------------------------±----------------------------------------------------------------------+
| Operation | Command |
+=+===============================+
| Filepath starts with “/Users” | … code-block:: |
| | |
| | ds.match(F(“filepath”).starts_with(“/Users”)) |
±----------------------------------------±----------------------------------------------------------------------+
| Filepath ends with “10.jpg” or “10.png” | … code-block:: |
| | |
| | ds.match(F(“filepath”).ends_with((“10.jpg”, “10.png”)) |
±----------------------------------------±----------------------------------------------------------------------+
| Label contains string “be” | … code-block:: |
| | |
| | ds.filter_labels( |
| | “predictions”, |
| | F(“label”).contains_str(“be”), |
| | ) |
±----------------------------------------±----------------------------------------------------------------------+
| Filepath contains “088” and is JPEG | … code-block:: |
| | |
| | ds.match(F(“filepath”).re_match(“088*.jpg”)) |
±----------------------------------------±----------------------------------------------------------------------+
在表格中,行可以占据任意数量的行,而列的宽度可能会不同。Grid table 单元格内的代码块也很难解析,因为它们跨越多行,所以其中的内容与其他列的内容交织在一起。这意味着在解析过程中需要有效地重新构建这些表格中的代码块。
这并不是世界末日,但也不是理想情况。
Jupyter
事实证明,解析 Jupyter 笔记本相对简单。我能够将 Jupyter 笔记本的内容读取为一个字符串列表,每个单元格对应一个字符串:
import json
ifile = “my_notebook.ipynb”
with open(ifile, “r”) as f:
contents = f.read()
contents = json.loads(contents)[“cells”]
contents = [(" ".join(c[“source”]), c[‘cell_type’] for c in contents]
此外,sections 是由以 #开头的 Markdown 单元格来标识的。
尽管如此,考虑到 RST 带来的挑战,我决定转向 HTML,并将我们的所有文档平等对待。
HTML
我使用 bash generate_docs。bash 从本地安装构建了 HTML 文档,并开始使用 Beautiful Soup 解析它们。然而,我很快意识到,当 RST 的代码块和带有内联代码的表格转换为 HTML 时,尽管它们在浏览器中正确显示,但生成的 HTML 非常冗长。以我们的 filtering cheat sheet 为例。
在浏览器中呈现时,filtering cheat sheet 中的日期和时间(Dates and times)节之前的代码块如下所示:
来自开源 FiftyOne 文档中 cheat sheet 的截图。
然而,原始 HTML 的内容如下所示:
RST 的备忘单转换为 HTML。
这并非无法解析,但离理想状态还有差距。
Markdown
幸运的是,我通过使用 markdownify 将所有 HTML 文件转换为 Markdown 格式,成功克服了这些问题。Markdown 具有几个关键优势,使其成为最佳选择。
- 比 HTML 更清晰:代码格式化变得更简洁,不再是一堆 < span > 元素的混乱字符串,而是使用单个反引号 ` 来标记内联代码片段,并使用三个反引号 ```来标记代码块。这使得文本和代码的划分变得容易。
- 仍保留锚点:与原始的 RST 不同,这个 Markdown 包含了节标题的锚点,因为隐式锚点已经生成。这样,我不仅可以链接到包含结果的页面,还可以链接到该页面的特定节或子节。
- 标准化:Markdown 提供了对初始的 RST 和 Jupyter 文档的大部分统一格式,使我们能够在向量搜索应用程序中对它们的内容进行一致处理。这样,它们的内容在格式上具有一致性。
LangChain 注意事项
一些人可能已经了解了用于使用 LLM 构建应用程序的开源库 LangChain,并且可能会想知道为什么我没有直接使用 LangChain 的文档加载器和文本分割器。答案是:我需要更多的控制!
02 处理文件
一旦文档被转换为 Markdown 格式,我就开始清理内容并将其分割成较小的片段。
Cleaning
文档清理的主要任务是去除不必要的元素,包括:
- 页眉和页脚
- 表格中的行和列结构,例如在 |select()| select_by()| 中的 | 符号
- 多余的换行符
- 链接
- 图像
- Unicode 字符
- 加粗格式,即将 文本 转换为 文本
此外,我还删除了转义字符,这些字符用于转义在我们的文档中具有特殊含义的字符:_ 和 。前者在许多方法名中使用,后者则常用于乘法、正则表达式模式和其他许多地方:
document = document.replace( “_” , “_” ).replace( “*” , “” )
将文档拆分为语义块
在清理了文档内容后,我开始将文档分割成适当大小的块。
首先,我将每个文档分割成几个部分。乍一看,这似乎可以通过查找以 # 字符开头的行来完成。在我的应用中,我没有区分 h1、h2、h3 等级的标题(#、##、###),所以检查第一个字符就足够了。然而,当意识到 # 也用于在 Python 代码中添加注释时,这种逻辑就会出现问题。
为了解决这个问题,我将文档分成了文本块和代码块:
text_and_code = page_md.split(‘```’)
text = text_and_code[::2]
code = text_and_code[1::2]
接着,我通过在文本块中以 # 开头的行来确定新部分的起始位置。我从该行中提取了部分的标题和锚点。
def extract_title_and_anchor(header):
header = “ “.join(header.split(” “)[1:])
title = header.split(”[“)[0]
anchor = header.split(”(“)[1].split(” “)[0]
return title, anchor
然后,我将每个文本块或代码分配给适当的部分。
最初,我还尝试将文本块分割为段落,假设一个部分可能包含关于许多不同主题的信息,因此该部分的整体嵌入可能与仅涉及其中一个主题的文本提示的嵌入不相似。然而,这种方法导致大多数搜索查询的前几个匹配结果过度倾向于单行段落,结果发现这并不是非常信息丰富的搜索结果。
请查看附带的 GitHub 存储库:https://github.com/voxel51/fiftyone-docs-search
其中包含这些方法的实现,你可以在自己的文档上尝试使用!
03 使用 OpenAI 嵌入文本和代码块
文档转换、处理和拆分为字符串后,我为每个块生成了一个嵌入向量。由于大语言模型本质上具有灵活性和普适性,我决定将文本块和代码块视为文本片段的同等部分,并使用相同的模型对它们进行嵌入。
我使用了 OpenAI 的 text-embedding-ada-002 模型,因为它易于使用,在 OpenAI 的所有嵌入模型中具有最高的性能(在 BEIR 基准测试中),并且价格也最便宜。
事实上,它非常便宜(每 1,000 个 token 只需 0.0004 美元),以至于为 FiftyOne 文档生成所有嵌入向量只需几美分!正如 OpenAI 自己所说:“我们建议几乎所有的使用场景都使用 text-embedding-ada-002。它更好、更便宜、更简单。”
使用这个嵌入模型,你可以生成一个表示任何输入提示的 1536 维向量,最多可以包含 8,191 个 token(约 30,000 个字符)。
要开始使用,你需要创建一个 OpenAI 账户,在 https://platform.openai.com/account/api-keys 上生成一个 API 密钥,并将此 API 密钥导出为环境变量,可以按以下方式执行:
export OPENAI_API_KEY=“<MY_API_KEY>”
你还需要安装 OpenAI Python 库:
pip install openai
我编写了一个围绕 OpenAI API 的包装器,它接收一个文本提示并返回一个嵌入向量:
MODEL = “text-embedding-ada-002”
def embed_text(text):
response = openai.Embedding.create(
input=text,
model=MODEL
)
embeddings = response[‘data’][0][‘embedding’]
return embeddings
要为所有文档生成嵌入向量,我们只需对所有文档中的各个子部分(文本块和代码块)应用此函数。
04 创建一个 Qdrant 向量索引
有了嵌入向量后,我创建了一个用于搜索的向量索引。选择使用 Qdrant 的原因与我们选择为 FiftyOne 添加原生 Qdrant 支持的原因相同:它是开源的、免费的,而且易于使用。
要开始使用 Qdrant,你可以拉取一个预先构建的 Docker 镜像并运行容器:
docker pull qdrant/qdrant
docker run -d -p 6333:6333 qdrant/qdrant
此外,你还需要安装 Qdrant Python 客户端:
pip install qdrant-client
我创建了 Qdrant 集合:
import qdrant_client as qc
import qdrant_client.http.models as qmodels
client = qc.QdrantClient(url=“localhost”)
METRIC = qmodels.Distance.DOT
DIMENSION = 1536
COLLECTION_NAME = “fiftyone_docs”
def create_index():
client.recreate_collection(
collection_name=COLLECTION_NAME,
vectors_config = qmodels.VectorParams(
size=DIMENSION,
distance=METRIC,
)
)
然后,我为每个子部分(文本块或代码块)创建了一个向量:
import uuid
def create_subsection_vector(
subsection_content,
section_anchor,
page_url,
doc_type
):
vector = embed_text(subsection_content)
id = str(uuid.uuid1().int)[:32]
payload = {
"text": subsection_content,
"url": page_url,
"section_anchor": section_anchor,
"doc_type": doc_type,
"block_type": block_type
}
return id, vector, payload
对于每个向量,你可以在有效载荷中提供额外的上下文。在这种情况下,我包括了可以找到结果的 URL(和锚点),文档类型,以便用户可以指定是搜索所有文档还是只搜索特定类型的文档,以及生成嵌入向量的字符串内容。我还添加了块类型(文本或代码),这样如果用户正在寻找代码片段,他们可以根据需要进行搜索。
然后,我逐个将这些向量添加到索引中,逐页进行操作:
def add_doc_to_index(subsections, page_url, doc_type, block_type):
ids = []
vectors = []
payloads = []
for section_anchor, section_content in subsections.items():
for subsection in section_content:
id, vector, payload = create_subsection_vector(
subsection,
section_anchor,
page_url,
doc_type,
block_type
)
ids.append(id)
vectors.append(vector)
payloads.append(payload)
## Add vectors to collection
client.upsert(
collection_name=COLLECTION_NAME,
points=qmodels.Batch(
ids = ids,
vectors=vectors,
payloads=payloads
),
)
05 查询索引
一旦索引创建完成,可以通过使用相同的 embedding 模型对查询文本进行 embedding,然后在索引中搜索相似的 embedding 向量来运行对已索引文档的搜索。使用 Qdrant 向量索引,可以通过 Qdrant 客户端的 search() 命令执行基本查询。
为了使公司的文档可搜索,我希望用户能够按照文档的部分和编码块的类型进行过滤。在向量搜索的术语中,通过过滤结果并确保返回预定数量的结果(由 top_k 参数指定)被称为预过滤。
为了实现这一目标,我编写了一个程序化的过滤器:
def _generate_query_filter(query, doc_types, block_types):
“”“Generates a filter for the query.
Args:
query: A string containing the query.
doc_types: A list of document types to search.
block_types: A list of block types to search.
Returns:
A filter for the query.
“””
doc_types = _parse_doc_types(doc_types)
block_types = _parse_block_types(block_types)
_filter = models.Filter(
must=[
models.Filter(
should= [
models.FieldCondition(
key="doc_type",
match=models.MatchValue(value=dt),
)
for dt in doc_types
],
),
models.Filter(
should= [
models.FieldCondition(
key="block_type",
match=models.MatchValue(value=bt),
)
for bt in block_types
]
)
]
)
return _filter
内部的 _parse_doc_types() 和 _parse_block_types() 函数处理参数为字符串或列表值,或为 None 的情况。
然后,我编写了一个名为 query_index() 的函数,该函数接受用户的文本查询,进行预过滤,搜索索引,并从有效载荷中提取相关信息。该函数返回一个由元组组成的列表,形式为(url,contents,score),其中分数表示结果与查询文本的匹配程度。
def query_index(query, top_k=10, doc_types=None, block_types=None):
vector = embed_text(query)
_filter = _generate_query_filter(query, doc_types, block_types)
results = CLIENT.search(
collection_name=COLLECTION_NAME,
query_vector=vector,
query_filter=_filter,
limit=top_k,
with_payload=True,
search_params=_search_params,
)
results = [
(
f"{res.payload['url']}#{res.payload['section_anchor']}",
res.payload["text"],
res.score,
)
for res in results
]return results
06 编写搜索包装器
最后一步是为用户提供一个清晰的界面,以便对这些 “向量化” 文档进行语义搜索。
我编写了一个名为 print_results() 的函数,该函数接受查询、来自 query_index() 的结果以及一个分数参数(是否打印相似度分数),并以易于理解的方式打印结果。我使用了 Python 的 rich 包来在终端中格式化超链接,以便在支持超链接的终端中,点击超链接将在默认浏览器中打开页面。如果需要,我还使用 webbrowser 自动打开排名第一的结果的链接。
使用富文本超链接显示搜索结果。
对于基于 Python 的搜索,我创建了一个名为 FiftyOneDocsSearch 的类,用于封装文档搜索行为,这样一旦实例化了一个 FiftyOneDocsSearch 对象(可能使用搜索参数的默认设置):
from fiftyone.docs_search import FiftyOneDocsSearch
fosearch = FiftyOneDocsSearch(open_url=False, top_k=3, score=True)
你可以通过调用此对象在 Python 中进行搜索。例如,要查询文档中的 “如何加载数据集”,只需运行以下代码:
fosearch(“How to load a dataset”)
在 Python 进程中进行语义搜索你公司的文档。
我还使用了 argparse,通过命令行使这个文档搜索功能可用。当包被安装后,可以通过以下方式在命令行中搜索文档:
fiftyone-docs-search query “” <args
只是为了好玩,因为 fiftyone-docs-search query 命令有点繁琐,我在我的。zshrc 文件中添加了一个别名:
alias fosearch=‘fiftyone-docs-search query’
使用这个别名,可以通过命令行搜索文档:
fosearch “” args
07 总结
进入这个项目时,我已经认为自己是我们公司开源 Python 库 FiftyOne 的高级用户。我撰写了许多文档,并且每天都在使用(并继续使用)该库。但是将我们的文档转化为可搜索的数据库的过程迫使我更深入地了解我们的文档。这总是令人欣喜的,当你为他人构建东西,并且最终也能帮助到自己!
以下是我学到的一些内容:
- Sphinx RST 很繁琐:它可以生成漂亮的文档,但解析起来有点麻烦。
- 不要过度预处理:OpenAI 的 text-embeddings-ada-002 模型非常擅长理解文本字符串的含义,即使它们具有稍微不典型的格式。不再需要进行词干提取、费力地去除停用词和其他杂项字符。
- 最好使用小而有意义的片段:将文档拆分为最小可能的有意义的片段,并保留上下文。对于较长的文本,很可能在索引中搜索查询与文本片段的重叠部分被不相关的文本所遮盖。如果将文档拆分得太小,可能导致索引中的许多条目包含的语义信息很少。
- 向量搜索非常强大:几乎没有额外的工作量,也没有进行任何微调,我就能显著提升我们文档的可搜索性。根据初步估计,这种改进的文档搜索相对于旧的关键字搜索方法来说,返回相关结果的几率超过两倍。此外,向量搜索方法的语义特性意味着用户现在可以使用任意短语、任意复杂的查询进行搜索,并确保得到指定数量的结果。
如果你发现自己(或他人)经常在大量文档中搜索特定的信息,我鼓励你根据自己的需求来适应这个过程。你可以修改它以适用于个人文档或公司的档案,我保证你会以全新的视角看待你的文档!
以下是你可以为自己的文档扩展这个过程的几种方式: - 混合搜索:将向量搜索与传统的关键字搜索结合使用。
- 全球化:使用 Qdrant Cloud 在云端存储和查询文集。
- 整合网络数据:使用 requests 直接从网络下载 HTML 内容。
- 自动更新:使用 Github Actions 在底层文档更改时自动触发嵌入向量的重新计算。
- Embed:将其包装为一个 Javascript 元素,并将其作为传统搜索栏的替代品。