跳到主要内容

ChromadbRM

改编自 https://github.com/animtel 提供的文档

ChromadbRM

ChromadbRM具有灵活性,可以使用chromadb嵌入文档中概述的各种嵌入函数。虽然有不同的选项可供选择,但此示例演示了如何专门利用OpenAI嵌入。

设置ChromadbRM客户端

构造函数初始化ChromadbRM类的一个实例,可以选择使用OpenAI的嵌入或任何chromadb支持的替代方案,详细信息请参阅官方chromadb嵌入文档

  • collection_name (str): chromadb集合的名称。
  • persist_directory (str): chromadb数据持久化的目录路径。
  • embedding_function (Optional[EmbeddingFunction[Embeddable]], optional): 用于嵌入文档和查询的函数。如果未指定,默认为DefaultEmbeddingFunction()
  • k (int, optional): 要检索的顶部段落数。默认为7。

ChromadbRM构造函数示例:

ChromadbRM(
collection_name: str,
persist_directory: str,
embedding_function: Optional[EmbeddingFunction[Embeddable]] = OpenAIEmbeddingFunction(),
k: int = 7,
)

内部机制

forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = None) -> dspy.Prediction

参数:

  • query_or_queries (Union[str, List[str]]): 要搜索的查询或查询列表。
  • k (Optional[int], optional): 要检索的结果数。如果未指定,默认为初始化期间设置的值。

返回:

  • dspy.Prediction: 包含检索到的段落,每个段落表示为具有long_text属性的dotdict

搜索chromadb集合,查找与给定查询或查询匹配的前k个段落,使用通过指定的embedding_function生成的嵌入。

通过ChromadbRM客户端发送检索请求

from dspy.retrieve.chromadb_rm import ChromadbRM
import os
import openai
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction

embedding_function = OpenAIEmbeddingFunction(
api_key=os.environ.get('OPENAI_API_KEY'),
model_name="text-embedding-ada-002"
)

retriever_model = ChromadbRM(
'your_collection_name',
'/path/to/your/db',
embedding_function=embedding_function,
k=5
)

results = retriever_model("探索量子计算的重要性", k=5)

for result in results:
print("文档:", result.long_text, "\n")