跳转到主要内容

​
retriever_init

签名 
async def retriever_init(
    model_name_or_path: str,
    backend_configs: Dict[str, Any],
    batch_size: int,
    corpus_path: str,
    gpu_ids: Optional[object] = None,
    is_multimodal: bool = False,
    backend: str = "sentence_transformers",
    index_backend: str = "faiss",
    index_backend_configs: Optional[Dict[str, Any]] = None,
    is_demo: bool = False,
    collection_name: str = "",
) -> None
功能
  • 初始化检索服务。
  • Embedding Backend (backend): 负责将文本/图像转换为向量 (Infinity, SentenceTransformers, OpenAI, BM25)。
  • Index Backend (index_backend): 负责向量的存储与检索 (FAISS, Milvus)。
  • Demo Mode: 若 is_demo=True，强制使用 OpenAI + Milvus 配置，忽略部分参数。

​
retriever_embed

签名 
async def retriever_embed(
    embedding_path: Optional[str] = None,
    overwrite: bool = False,
    is_multimodal: bool = False,
) -> None
功能
  • (非 Demo 模式) 批量计算语料库的向量表示，并保存为 .npy 文件。
  • 仅适用于 Dense Retriever 后端（BM25 不支持）。

​
retriever_index

签名 
async def retriever_index(
    embedding_path: str,
    overwrite: bool = False,
    collection_name: str = "",
    corpus_path: str = ""
) -> None
功能
  • 构建检索索引。
  • FAISS: 读取 embedding_path (.npy) 构建本地索引文件。
  • Milvus / Demo: 读取 corpus_path (.jsonl)，生成向量并插入到指定的 collection_name 中。
签名 
async def retriever_search(
    query_list: List[str],
    top_k: int = 5,
    query_instruction: str = "",
    collection_name: str = "",
) -> Dict[str, List[List[str]]]
功能
  • 对单条或多条查询进行检索。
  • 自动处理查询向量化（添加 query_instruction）并在指定 collection_name (针对 Milvus) 或默认索引中查找 Top-K。
输出格式（JSON） 
{"ret_psg": [["passage 1", "passage 2"], ["..." ]]}
签名 
async def retriever_batch_search(
    batch_query_list: List[List[str]],
    top_k: int = 5,
    query_instruction: str = "",
    collection_name: str = "",
) -> Dict[str, List[List[List[str]]]]
功能
  • etriever_search 的批处理版本，接受嵌套列表输入。
输出格式（JSON） 
{"ret_psg_ls": [[["psg 1-1"], ["psg 1-2"]], [["psg 2-1"]]]}

​
bm25_index

签名 
async def bm25_index(
    overwrite: bool = False,
) -> None
功能
  • 当 backend="bm25" 时，构建 BM25 稀疏索引并保存。
签名 
async def bm25_search(
    query_list: List[str],
    top_k: int = 5,
) -> Dict[str, List[List[str]]]
功能
  • 基于 BM25 算法进行关键词检索。
输出格式（JSON） 
{"ret_psg": [["passage 1", "passage 2"], ["..." ]]}
签名 
async def retriever_deploy_search(
    retriever_url: str,
    query_list: List[str],
    top_k: int = 5,
    query_instruction: str = "",
) -> Dict[str, List[List[str]]]
功能
  • 作为客户端，调用部署在 retriever_url 的远程检索服务进行查询。
输出格式（JSON） 
{"ret_psg": [["passage 1", "passage 2"], ["..." ]]}
签名 
async def retriever_exa_search(
    query_list: List[str],
    top_k: Optional[int] | None = 5,
    retrieve_thread_num: Optional[int] | None = 1,
) -> Dict[str, List[List[str]]]
功能
  • 调用 Exa Web 检索（需要 EXA_API_KEY）。
输出格式（JSON） 
{"ret_psg": [["snippet 1", "snippet 2"], ["..." ]]}
签名 
async def retriever_tavily_search(
    query_list: List[str],
    top_k: Optional[int] | None = 5,
    retrieve_thread_num: Optional[int] | None = 1,
) -> Dict[str, List[List[str]]]
功能
  • 调用 Tavily Web 检索（需要 TAVILY_API_KEY）。
输出格式（JSON） 
{"ret_psg": [["snippet 1", "snippet 2"], ["..." ]]}
签名 
async def retriever_zhipuai_search(
    query_list: List[str],
    top_k: Optional[int] | None = 5,
    retrieve_thread_num: Optional[int] | None = 1,
) -> Dict[str, List[List[str]]]
功能
  • 调用 智谱AI web_search（需要 ZHIPUAI_API_KEY）。
输出格式（JSON） 
{"ret_psg": [["snippet 1", "snippet 2"], ["..." ]]}

​
参数配置

https://mintcdn.com/ultrarag/T7GffHzZitf6TThi/images/yaml.svg?fit=max&auto=format&n=T7GffHzZitf6TThi&q=85&s=69b41e79144bc908039c2ee3abbb1c3bservers/retriever/parameter.yaml
model_name_or_path: openbmb/MiniCPM-Embedding-Light
corpus_path: data/corpus_example.jsonl
embedding_path: embedding/embedding.npy
collection_name: wiki

# Embedding Backend Configuration
backend: sentence_transformers # options: infinity, sentence_transformers, openai, bm25
backend_configs:
  infinity:
    bettertransformer: false
    pooling_method: auto
    model_warmup: false
    trust_remote_code: true
  sentence_transformers:
    trust_remote_code: true
    sentence_transformers_encode:
      normalize_embeddings: false
      encode_chunk_size: 256
      q_prompt_name: query
      psg_prompt_name: document
      psg_task: null
      q_task: null
  openai:
    model_name: text-embedding-3-small
    base_url: "https://api.openai.com/v1"
    api_key: "abc"
  bm25:
    lang: en
    save_path: index/bm25

# Index Backend Configuration
index_backend: faiss # options: faiss, milvus
index_backend_configs:
  faiss:
    index_use_gpu: True
    index_chunk_size: 10000
    index_path: index/index.index
  milvus:
    uri: index/milvus_demo.db # Local file for Lite, or http://host:port
    token: null
    id_field_name: id
    vector_field_name: vector
    text_field_name: contents
    index_params:
      index_type: AUTOINDEX
      metric_type: IP

batch_size: 16
top_k: 5
gpu_ids: "1"
query_instruction: ""
is_multimodal: false
overwrite: false
retrieve_thread_num: 1
retriever_url: "http://127.0.0.1:64501"
is_demo: false
参数说明：
参数类型说明
model_name_or_pathstr检索模型路径或名称（如 HuggingFace 模型 ID）
corpus_pathstr输入语料 JSONL 文件路径
embedding_pathstr向量文件保存路径（.npy）
collection_namestrMilvus 集合名称
backendstr选择检索后端：infinity、sentence_transformers、openai、bm25
index_backendstr索引后端：faiss, milvus
backend_configsdict各后端的参数配置（见下表）
index_backend_configsdict各索引后端的参数配置（见下表）
batch_sizeint向量生成或检索的批大小
top_kint返回的候选段落数量
gpu_idsstr指定可见 GPU 设备，如 "0,1"
query_instructionstr查询前缀（instruction-tuning 模型使用）
is_multimodalbool是否启用多模态嵌入（如图像）
overwritebool若已存在嵌入或索引文件是否覆盖
retrieve_thread_numint外部 Web 检索（Exa/Tavily/Zhipu）并发线程数
retriever_urlstr部署 retriever server 的 url
is_demobool演示模式开关（强制使用 OpenAI+Milvus，简化配置）
backend_configs 子项：
后端参数类型说明
infinitybettertransformerbool是否启用高效推理优化
pooling_methodstr池化方式（如 auto, mean）
model_warmupbool是否预加载模型到显存
trust_remote_codebool是否信任远程代码（适用于自定义模型）
sentence_transformerstrust_remote_codebool是否信任远程模型代码
sentence_transformers_encodedict编码详细参数，见下表
openaimodel_namestrOpenAI 模型名称（如 text-embedding-3-small）
base_urlstrAPI 基地址
api_keystrOpenAI API 密钥
bm25langstr语言（决定停用词与分词器）
save_pathstrBM25 稀疏索引的保存目录
sentence_transformers_encode 参数：
参数类型说明
normalize_embeddingsbool是否归一化向量
encode_chunk_sizeint编码块大小（避免显存溢出）
q_prompt_namestr查询模板名
psg_prompt_namestr段落模板名
q_taskstr任务描述（针对特定模型需要指定 Task 的情况）
psg_taskstr任务描述（针对特定模型需要指定 Task 的情况）
index_backend_configs 参数：
后端参数类型说明
faissindex_use_gpubool是否使用 GPU 构建和检索索引
index_chunk_sizeint构建索引时的分批大小
index_pathstrFAISS 索引文件的保存路径（.index）
milvusuristrMilvus 连接地址（本地文件路径即启用 Milvus Lite）
tokenstr认证 Token（如需要）
id_field_namestr主键字段名（默认 id）
vector_field_namestr向量字段名（默认 vector）
text_field_namestr文本内容字段名（默认 contents）
id_max_lengthint字符串主键的最大长度
text_max_lengthint文本字段的最大长度（超过截断）
metric_typestr距离度量方式（如 IP 内积, L2 欧式距离）
index_paramsDict索引构建参数（如 index_type: AUTOINDEX）
search_paramsDict检索参数（如 nprobe 等）
index_chunk_sizeint插入数据时的批处理大小