在 chromadb 中创建自定义嵌入函数以进行语义搜索

Question

我有下面的 python 3 代码。 我有 chromadb 矢量数据库，我正在尝试使用自定义嵌入函数为文本块创建嵌入，如下例所示。 我的最终目标是对我从这些文本块创建的集合进行语义搜索。 因此，我将文本块以及嵌入和元数据更新插入到 chromadb 集合中，然后查询该集合。 我想知道我是否正确创建了自定义嵌入函数。

创建自定义嵌入函数的文档可以在这里找到：

https://docs.trychroma.com/embeddings

我想使用名为 all-MiniLM-L12-v2 的 Huggingface 句子转换器模型。 它的文档可以在这里找到：

https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2

（我知道可以在 SentenceTransformerEmbeddingFunction 中切换到 all_MiniLM-L12-v2，但我正在用它构建一个函数，以更好地了解如何创建自定义嵌入函数。）

我想知道是否需要将文本拆分为逗号分隔的句子列表并在 MyEmbeddingFunction 中传递 model.encode？ 我通过结合上面列出的文档中提到的两个代码源创建了下面的自定义函数。 在 all-MiniLM-L12-v2 文档页面示例代码中，传递给 model.encode 的输入是逗号分隔的句子列表。 但在我的代码中，我传递了整个段落，如下所示。 如果我将下面的示例拆分为逗号分隔的句子列表并将该列表传递给 model.encode，查询集合是否会更好？

下面还有我的代码和查询结果。

示例：

“一封标题为“紧急 || 数据科学家/工程师 || 地点 - 内华达州拉斯维加斯”的电子邮件已于太平洋夏令时 2023 年 8 月 22 日星期二上午 06:54 发送给求职者杰罗姆·鲍威尔。这是针对数据职位的科学家/工程师。工作地点是内华达州拉斯维加斯。工作类型是合同制。它拥有所需的技能：统计编程语言、R、Python、sql、hive、pig、scala、java、C++、统计、统计测试、分布。 、回归、最大似然估计、机器学习、k-最近邻、朴素贝叶斯、SVM、决策森林、数据整理、数据可视化、matplotlib、ggplot、d3.js、Tableau、沟通技巧、软件工程、问题解决、分析，程度。”

代码：

# creating custom embeddings with non-default embedding model

from chromadb import Documents, EmbeddingFunction, Embeddings

class MyEmbeddingFunction(EmbeddingFunction):
    def __call__(self, input: Documents) -> Embeddings:
        # embed the documents
        
        from sentence_transformers import SentenceTransformer

        sentences = input
    
        model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2',
                                   device='cuda',
                           use_auth_token=hf_auth,
                           cache_folder='/home/username/stuff/username_storage/LLM/weights/huggingface/hub/')
        embeddings = model.encode(sentences)


        # Convert embeddings to a list of lists
        embeddings_as_list = [embedding.tolist() for embedding in embeddings]
        
        return embeddings_as_list
        
        
        
custom_embeddings=MyEmbeddingFunction()

test_collection = chroma_client\
.get_or_create_collection(name="test_custom_embeddings",
                          embedding_function=custom_embeddings
                         
                         )
                         
                         
# inserting data

test_collection.upsert(
    ids=[f"{x}" for x in summary_df['id'].tolist()],
    documents=summary_df['summary'].tolist(),
    metadatas=summary_df['meta'].tolist()    
)


qry_str = """Title contains Data Scientist"""


db_query_results=test_collection.query(query_texts=qry_str, n_results=2)

result_summaries=[x['summary'] for x in db_query_results['metadatas'][0]]

result_summaries  

输出：

[“一封标题为“紧急 || 数据科学家/工程师 || 地点 - 内华达州拉斯维加斯”的电子邮件已于太平洋夏令时间 2023 年 8 月 22 日星期二上午 06:54 发送给求职者 Jerome Powell。该职位的职位是：数据科学家/工程师。它的位置是拉斯维加斯，内华达州。雇佣类型是合同。它具有所需的技能：统计编程语言、R、Python、sql、hive、pig、scala、java、C++、统计、统计测试、分布、回归、最大似然估计、机器学习、k 最近邻、朴素贝叶斯、SVM、决策森林、数据整理、数据可视化、matplotlib、ggplot、d3.js、Tableau、沟通技巧、软件工程、问题解决，分析，程度。", “一封标题为：首席数据科学家 - 密苏里州奥法伦（混合型）的电子邮件于太平洋夏令时间 2023 年 8 月 22 日星期二上午 07:16 发送给求职者杰罗姆·鲍威尔。这是针对首席数据科学家的职位。地点是密苏里州奥法伦（混合） 就业类型是合同 它具有所需的技能：数学、统计学、计算机科学或相关领域的硕士或博士学位，领导大型数据科学项目、研究、沟通技巧、预测能力。 、batch、streaming、python、R、hadoop、spark、MySQL、异常检测、监督学习、无监督学习、时间序列、自然语言处理、Numpy、SciPy、Pandas、Scikit-learn、Tensorflow、Keras、NLTK、Gensim、 BERT、NetworkX、有组织、自我激励、数据可视化。”]

Answer 1

所以我相信对于你的用例，你应该尝试几种不同的分块选项，例如，如果你使用 langchain，它会像 -

from langchain_text_splitters import RecursiveCharacterTextSplitter
default_text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100,
    length_function=len,
    is_separator_regex=False,
)

您可以尝试更改 chunk_size 和 chunk_overlap 参数，看看哪种分块策略提供更好的结果。
您还可以研究不同的检索方法，例如句子窗口检索和自动合并检索。