检索与查询数据
搜索与查询数据
相关源文件
以下文件为本维基页面的生成提供了上下文:
argilla-server/README.mdargilla-server/src/argilla_server/api/schemas/v1/suggestions.pyargilla/README.mdargilla/docs/assets/images/getting_started/dataset_configurator.pngargilla/docs/assets/images/how_to_guides/custom_field/3d_object_viewer.pngargilla/docs/assets/images/how_to_guides/custom_field/images_in_two_columns.pngargilla/docs/assets/images/how_to_guides/custom_field/metadata_table.pngargilla/docs/getting_started/faq.mdargilla/docs/getting_started/how-to-configure-argilla-on-huggingface.mdargilla/docs/getting_started/how-to-deploy-argilla-with-docker.mdargilla/docs/getting_started/quickstart.mdargilla/docs/how_to_guides/annotate.mdargilla/docs/how_to_guides/custom_fields.mdargilla/docs/how_to_guides/distribution.mdargilla/docs/how_to_guides/query.mdargilla/docs/index.mdargilla/docs/reference/argilla/client.mdargilla/docs/reference/argilla/search.mdargilla/src/argilla/_models/_dataset_progress.pyargilla/src/argilla/_models/_record/_record.pyargilla/src/argilla/records/_dataset_records.pyargilla/src/argilla/records/_io/_datasets.pyargilla/src/argilla/records/_io/_generic.pyargilla/src/argilla/records/_io/_json.pyargilla/src/argilla/records/_resource.pyargilla/src/argilla/records/_search.pyargilla/src/argilla/settings/_question.pyargilla/tests/integration/test_export_dataset.pyargilla/tests/integration/test_export_records.pyargilla/tests/integration/test_list_records.pyargilla/tests/integration/test_query_records.pyargilla/tests/integration/test_search_records.pyargilla/tests/unit/export/test_record_export_import_compatibillity.pyargilla/tests/unit/test_io/test_generic.pyargilla/tests/unit/test_io/test_hf_datasets.pyargilla/tests/unit/test_record_fields.pyargilla/tests/unit/test_resources/test_records.py
Argilla 提供了强大的搜索与查询数据集功能。本指南将介绍如何使用这些功能,基于文本查询、过滤条件和向量相似度搜索等不同条件,高效地过滤、搜索和检索记录。
关于创建和管理数据集的信息,请参阅创建与管理数据集。关于标注工作流的信息,请参阅标注工作流。
搜索功能概览
Argilla 的搜索功能基于以下几个核心组件构建:
来源:argilla/src/argilla/records/_search.py:14-201, argilla/src/argilla/records/_dataset_records.py:29-154
Argilla 中搜索与查询的主要组件包括:
- 查询(Query):用于基于文本的搜索,并可结合过滤条件或相似度搜索
- 过滤(Filter):用于根据特定条件过滤记录
- 相似(Similar):用于基于向量的相似度搜索
这些组件可以独立使用,也可以组合使用以创建复杂的搜索查询。
使用 查询 类进行基本查询
Query 类是搜索数据集中记录的主要接口。它可以接受文本查询、过滤条件或相似度搜索参数。
基于文本的搜索
要搜索包含特定文本的记录:
import argilla as rg
# 简单文本查询
query = rg.Query(query="machine learning")
# 将查询与 dataset.records() 配合使用
results = dataset.records(query=query)
该搜索会返回在其任意文本字段中包含所有指定术语的记录。
来源:argilla/src/argilla/records/_dataset_records.py:195-241, argilla/src/argilla/records/_search.py:149-201
高级查询语法
Argilla 支持基于 Elasticsearch 简单查询字符串语法的高级查询语法:
| 运算符 | 描述 | 示例 |
|---|---|---|
+ 或空格 | AND:搜索两个术语 | argilla + distilabel 或 argilla distilabel |
| | OR:搜索任一术语 | argilla | distilabel |
- | 否定:排除某个术语 | argilla -distilabel |
* | 前缀:搜索某个前缀 | arg* |
" | 短语:搜索一个短语 | "argilla and distilabel" |
( 和 ) | 优先级:对术语分组 | (argilla | distilabel) rules |
~N | 编辑距离:模糊搜索 | argilla~1 |
来源:argilla/docs/how_to_guides/query.md:77-92
使用 筛选 类过滤记录
Filter 类允许你定义条件,根据记录的属性、元数据、响应或建议来过滤记录。
来源:argilla/src/argilla/records/_search.py:38-95, argilla/src/argilla/records/_search.py:128-147
过滤运算符
Argilla 支持以下过滤运算符:
| 运算符 | 描述 |
|---|---|
== | 等于 |
>= | 大于等于 |
<= | 小于等于 |
in | 值在列表中 |
过滤示例
import argilla as rg
# 按单个条件过滤
filter_condition = ("sentiment.suggestion", "==", "positive")
query = rg.Query(filter=filter_condition)
# 按多个条件过滤
filter_conditions = [
("metadata.source", "==", "twitter"),
("sentiment.score", ">=", 0.8)
]
query = rg.Query(filter=filter_conditions)
# 显式创建过滤对象
my_filter = rg.Filter([
("label", "in", ["positive", "neutral"]),
("metadata.processed", "==", True)
])
query = rg.Query(filter=my_filter)
# 使用查询
results = dataset.records(query=query)
来源:argilla/docs/how_to_guides/query.md:94-166, argilla/tests/integration/test_search_records.py:60-99
按记录属性过滤
你可以按记录的核心属性进行过滤:
# 按记录 ID 过滤
query = rg.Query(filter=("id", "==", "1"))
# 按服务器生成的 ID 过滤
query = rg.Query(filter=("_server_id", "==", server_id))
# 按状态过滤
query = rg.Query(filter=("status", "==", "pending"))
# 按时间戳过滤
query = rg.Query(filter=("inserted_at", "<=", datetime.utcnow()))
query = rg.Query(filter=("updated_at", ">=", datetime.utcnow() - timedelta(days=1)))
来源:argilla/tests/integration/test_search_records.py:60-112
按元数据、响应和建议过滤
# 按元数据过滤
query = rg.Query(filter=("metadata.language", "==", "en"))
# 按建议值过滤
query = rg.Query(filter=("label.suggestion", "==", "positive"))
# 按建议分数过滤
query = rg.Query(filter=("label.score", ">=", 0.8))
# 按建议代理过滤
query = rg.Query(filter=("label.agent", "==", "model-1"))
# 按建议类型(人工或模型)过滤
query = rg.Query(filter=("label.type", "==", "human"))
# 按响应状态过滤
query = rg.Query(filter=("response.status", "==", "submitted"))
来源:argilla/tests/integration/test_search_records.py:114-156
使用 Similar 类进行语义搜索
Similar 类支持基于向量的相似度搜索,用于查找与参考向量或记录语义相似的记录。
查找相似记录
import argilla as rg
# 搜索与某个向量相似的记录
vector_values = [0.1, 0.2, 0.3, 0.4, 0.5]
similar = rg.Similar(name="text_vector", value=vector_values)
query = rg.Query(similar=similar)
# 搜索与另一条记录相似的记录
reference_record = dataset.records(limit=1)[0]
similar = rg.Similar(name="text_vector", value=reference_record)
query = rg.Query(similar=similar)
# 搜索最不相似的记录
similar = rg.Similar(name="text_vector", value=vector_values, most_similar=False)
query = rg.Query(similar=similar)
# 使用查询
results = dataset.records(query=query)
使用相似度搜索时,结果以 (record, similarity_score) 元组的形式返回。
来源:argilla/src/argilla/records/_search.py:100-126, argilla/tests/integration/test_search_records.py:158-245
配置记录检索
检索记录时,你可以配置各种参数:
# 限制返回的记录数量
records = dataset.records(
query=my_query,
limit=100
)
# 设置检索的批次大小
records = dataset.records(
query=my_query,
batch_size=256
)
# 从指定偏移量开始
records = dataset.records(
query=my_query,
start_offset=10
)
# 包含或排除建议、响应和向量
records = dataset.records(
query=my_query,
with_suggestions=True,
with_responses=True,
with_vectors=["text_vector"]
)
来源:argilla/src/argilla/records/_dataset_records.py:195-241, argilla/tests/integration/test_list_records.py:43-66
导出搜索结果
Argilla 提供了多种导出搜索结果的方法:
# 导出为字典列表
records_list = dataset.records(query=my_query).to_list(flatten=True)
# 导出为以字段名作为键的字典
records_dict = dataset.records(query=my_query).to_dict(flatten=True)
# 导出为以记录 ID 作为键的字典
records_dict = dataset.records(query=my_query).to_dict(flatten=False, orient="index")
# 导出为 JSON 文件
dataset.records(query=my_query).to_json(path="records.json")
# 导出为 Hugging Face 数据集
hf_dataset = dataset.records(query=my_query).to_datasets()
来源:argilla/src/argilla/records/_dataset_records.py:348-415, argilla/tests/integration/test_export_records.py:94-197
扁平格式与嵌套格式
导出记录时,你可以选择扁平格式或嵌套格式:
- 扁平格式(flatten=True):字段、元数据、建议和响应使用点号表示法展平(例如
label.suggestion、label.response) - 嵌套格式(flatten=False):字段、元数据、建议和响应为嵌套字典
来源:argilla/src/argilla/records/_io/_generic.py:27-140, argilla/tests/unit/test_io/test_generic.py:23-97
完整工作流示例
以下是一个完整的示例,展示了如何搜索、过滤和导出记录:
import argilla as rg
from datetime import datetime, timedelta
# 连接到 Argilla
client = rg.Argilla(api_url="<api_url>", api_key="<api_key>")
# 获取数据集
dataset = client.datasets("my_dataset")
# 创建结合文本搜索、过滤和相似度搜索的复杂查询
one_week_ago = datetime.utcnow() - timedelta(days=7)
# 获取用于相似度搜索的参考记录
reference_record = dataset.records(limit=1)[0]
# 创建查询
query = rg.Query(
query="important customer feedback",
filter=[
("metadata.source", "==", "customer_support"),
("sentiment.suggestion", "in", ["negative", "very_negative"]),
("sentiment.score", ">=", 0.8),
("inserted_at", ">=", one_week_ago)
],
similar=rg.Similar(name="feedback_vector", value=reference_record)
)
# 执行查询并获取结果
results = dataset.records(
query=query,
limit=100,
with_suggestions=True,
with_responses=True
)
# 导出为 Hugging Face 数据集
hf_dataset = results.to_datasets()
# 打印摘要统计信息
print(f"找到 {len(hf_dataset)} 条记录")
来源:argilla/docs/how_to_guides/query.md:1-166, argilla/tests/integration/test_search_records.py:36-245
实现细节
查询 类实现
Query 类在 _search.py 中实现,提供了创建搜索查询的高级接口:
class Query:
def __init__(
self,
*,
query: Union[str, None] = None,
similar: Union[Similar, None] = None,
filter: Union[Filter, Conditions, None] = None,
):
# 实现细节
def has_search(self) -> bool:
return bool(self.query or self.has_similar() or self.filter)
def has_similar(self) -> bool:
return bool(self.similar)
def api_model(self) -> SearchQueryModel:
# 实现细节
来源:argilla/src/argilla/records/_search.py:149-200
DatasetRecordsIterator 实现
DatasetRecordsIterator 类负责实际的查询和记录迭代:
class DatasetRecordsIterator:
def __init__(
self,
dataset: "Dataset",
client: "Argilla",
query: Optional[Query] = None,
start_offset: int = 0,
batch_size: Optional[int] = None,
with_suggestions: bool = False,
with_responses: bool = False,
with_vectors: Optional[Union[str, List[str], bool]] = None,
limit: Optional[int] = None,
):
# 实现细节
def _fetch_from_server_with_search(self) -> List[Tuple[RecordModel, float]]:
search_items, total = self.__client.api.records.search(
dataset_id=self.__dataset.id,
query=self.__query.api_model(),
limit=self.__batch_size,
offset=self.__offset,
with_responses=self.__with_responses,
with_suggestions=self.__with_suggestions,
with_vectors=self.__with_vectors,
)
return search_items
来源:argilla/src/argilla/records/_dataset_records.py:42-162
性能考量
处理大型数据集时,请考虑以下性能建议:
- 高效使用过滤条件:更具体的过滤条件有助于更快地缩小结果范围
- 限制结果集:当不需要所有记录时,始终使用
limit参数 - 批量处理:使用适当的
batch_size以避免处理大型数据集时出现内存问题 - 选择性加载:仅在需要时才包含建议、响应或向量,使用相应的参数
- 分页:对于大型结果集,使用
start_offset进行分页遍历
来源:argilla/src/argilla/records/_dataset_records.py:42-162, argilla/tests/integration/test_list_records.py:43-212