使用Chroma构建RAG遇SentenceTransformer无embed_documents属性错误求助

使用Chroma数据库构建RAG系统，目标是将生成的索引用作聊天机器人查询引擎，但执行代码时触发错误：AttributeError: 'SentenceTransformer' object has no attribute 'embed_documents'。因环境限制无法修改Chroma库，寻求可行的解决办法。

尝试的代码 #

# 加载文本片段并指定要嵌入的列
chunks = DataFrameLoader(final_df_for_chroma_injection,
                        page_content_column='TEXT').load()

# 创建开源嵌入函数
embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2')
# 加载存储了之前嵌入结果的持久化目录，并添加新的片段/嵌入结果
index = Chroma.from_documents(chunks,
                             embedding_model,
                             persist_directory="./chroma_db")

错误信息 #

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[47], line 3
      1 #-Load the persist directory on which are stored the previous embeddings
      2 #-And add the new ones from chunks/embeddings
----> 3 index = Chroma.from_documents(chunks,
      4                              embedding_model,
      5                              persist_directory="./chroma_db")

File /opt/anaconda3_envs/abeille_pytorch_p310/lib/python3.10/site-packages/langchain_community/vectorstores/chroma.py:778, in Chroma.from_documents(cls, documents, embedding, ids, collection_name, persist_directory, client_settings, client, collection_metadata, **kwargs)
    776 texts = [doc.page_content for doc in documents]
    777 metadatas = [doc.metadata for doc in documents]
--> 778 return cls.from_texts(
    779     texts=texts,
    780     embedding=embedding,
    781     metadatas=metadatas,
    782     ids=ids,
    783     collection_name=collection_name,
    784     persist_directory=persist_directory,
    785     client_settings=client_settings,
    786     client=client,
    787     collection_metadata=collection_metadata,
    788     **kwargs,
    789 )

File /opt/anaconda3_envs/abeille_pytorch_p310/lib/python3.10/site-packages/langchain_community/vectorstores/chroma.py:736, in Chroma.from_texts(cls, texts, embedding, metadatas, ids, collection_name, persist_directory, client_settings, client, collection_metadata, **kwargs)
    728     from chromadb.utils.batch_utils import create_batches
    730     for batch in create_batches(
    731         api=chroma_collection._client,
    732         ids=ids,
    733         metadatas=metadatas,
    734         documents=texts,
    735     ):
--> 736         chroma_collection.add_texts(
    737             texts=batch[3] if batch[3] else [],
    738             metadatas=batch[2] if batch[2] else None,
    739             ids=batch[0],
    740         )
    741 else:
    742     chroma_collection.add_texts(texts=texts, metadatas=metadatas, ids=ids)

File /opt/anaconda3_envs/abeille_pytorch_p310/lib/python3.10/site-packages/langchain_community/vectorstores/chroma.py:275, in Chroma.add_texts(self, texts, metadatas, ids, **kwargs)
    273 texts = list(texts)
    274 if self._embedding_function is not None:
--> 275     embeddings = self._embedding_function.embed_documents(texts)
    276 if metadatas:
    277     # fill metadatas with empty dicts if somebody
    278     # did not specify metadata for all texts
    279     length_diff = len(texts) - len(metadatas)

File /opt/anaconda3_envs/abeille_pytorch_p310/lib/python3.10/site-packages/torch/nn/modules/module.py:1688, in Module.__getattr__(self, name)
   1686     if name in modules:
   1687         return modules[name]
--> 1688 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

AttributeError: 'SentenceTransformer' object has no attribute 'embed_documents'

解决建议 #

错误原因是LangChain的Chroma接口要求嵌入函数必须实现embed_documents和embed_query方法，但原生SentenceTransformer对象仅提供encode方法，不符合接口要求。无需修改Chroma库，可通过以下两种方式解决：

方案1：使用LangChain封装的SentenceTransformerEmbeddings #

直接替换原生SentenceTransformer为LangChain提供的适配类，该类已实现所需接口：

# 加载文本片段并指定要嵌入的列
chunks = DataFrameLoader(final_df_for_chroma_injection,
                        page_content_column='TEXT').load()

# 导入LangChain封装的嵌入类
from langchain_community.embeddings import SentenceTransformerEmbeddings
# 创建符合LangChain接口的嵌入函数
embedding_model = SentenceTransformerEmbeddings(model_name='sentence-transformers/all-MiniLM-L12-v2')

# 初始化Chroma索引
index = Chroma.from_documents(chunks,
                             embedding_model,
                             persist_directory="./chroma_db")

方案2：自定义适配类包装原生SentenceTransformer #

如果不想依赖LangChain的封装，可手动实现一个适配类，将原生encode方法映射为embed_documents和embed_query：

from sentence_transformers import SentenceTransformer

class CustomSentenceTransformerEmbeddings:
    def __init__(self, model_name):
        self.model = SentenceTransformer(model_name)
    
    def embed_documents(self, texts):
        # 批量处理文本嵌入，返回numpy数组或列表
        return self.model.encode(texts, convert_to_numpy=True).tolist()
    
    def embed_query(self, text):
        # 单个查询文本嵌入
        return self.model.encode(text, convert_to_numpy=True).tolist()

# 使用自定义类创建嵌入函数
embedding_model = CustomSentenceTransformerEmbeddings('sentence-transformers/all-MiniLM-L12-v2')

# 后续代码保持不变
chunks = DataFrameLoader(final_df_for_chroma_injection,
                        page_content_column='TEXT').load()
index = Chroma.from_documents(chunks,
                             embedding_model,
                             persist_directory="./chroma_db")

内容的提问来源于stack exchange，提问作者Amin543