如何用SentenceTransformers替代OpenAI生成向量嵌入？代码报错求助

背景 #

我找到一份用于构建RAG系统的代码，功能是读取文档、分块、生成向量嵌入并保存到Chroma数据库，但原代码依赖OpenAI密钥。因为无法访问OpenAI，我尝试用免费的SentenceTransformers改写，但运行报错。

尝试代码 #

from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
from sentence_transformers import SentenceTransformer
from langchain.vectorstores.chroma import Chroma
import os
import shutil

CHROMA_PATH = "chroma"
DATA_PATH = "data/books"

embedder = SentenceTransformer("all-MiniLM-L6-v2")

def main():
    generate_data_store()

def generate_data_store():
    documents = load_documents()
    chunks = split_text(documents)
    save_to_chroma(chunks)

def load_documents():
    loader = DirectoryLoader(DATA_PATH, glob="*.md")
    documents = loader.load()
    return documents

def split_text(documents: list[Document]):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,
        chunk_overlap=100,
        length_function=len,
        add_start_index=True,
    )
    chunks = text_splitter.split_documents(documents)
    print(f"Split {len(documents)} documents into {len(chunks)} chunks.")

    document = chunks[10]
    print(document.page_content)
    print(document.metadata)

    return chunks

def save_to_chroma(chunks: list[Document]):
    # Clear out the database first.
    if os.path.exists(CHROMA_PATH):
        shutil.rmtree(CHROMA_PATH)

    # Create a new DB from the documents.
    db = Chroma.from_documents(
        chunks, embedder.encode, persist_directory=CHROMA_PATH
    )
    db.persist()
    print(f"Saved {len(chunks)} chunks to {CHROMA_PATH}.")

if __name__ == "__main__":
    main()

错误信息 #

Traceback (most recent call last):
  File "/media/andrew/Simple Tom/Robotics/Crew_AI/langchain-rag-tutorial/create_database.py", line 56, in <module>
    main()
  File "/media/andrew/Simple Tom/Robotics/Crew_AI/langchain-rag-tutorial/create_database.py", line 15, in main
    generate_data_store()
  File "/media/andrew/Simple Tom/Robotics/Crew_AI/langchain-rag-tutorial/create_database.py", line 20, in generate_data_store
    save_to_chroma(chunks)
  File "/media/andrew/Simple Tom/Robotics/Crew_AI/langchain-rag-tutorial/create_database.py", line 49, in save_to_chroma
    db = Chroma.from_documents(
  File "/home/andrew/.local/lib/python3.10/site-packages/langchain_community/vectorstores/chroma.py", line 778, in from_documents
    return cls.from_texts(
  File "/home/andrew/.local/lib/python3.10/site-packages/langchain_community/vectorstores/chroma.py", line 736, in from_texts
    chroma_collection.add_texts(
  File "/home/andrew/.local/lib/python3.10/site-packages/langchain_community/vectorstores/chroma.py", line 275, in add_texts
    embeddings = self._embedding_function.embed_documents(texts)
AttributeError: 'function' object has no attribute 'embed_documents'

解答 #

完全可以不用OpenAI密钥实现这个功能，你的代码错误在于：

Chroma的from_documents方法需要传入LangChain标准的Embeddings类实例，但你直接传了SentenceTransformer的encode方法（一个函数），而LangChain会尝试调用这个对象的embed_documents方法，所以触发了AttributeError。

修正方案 #

使用LangChain提供的HuggingFaceEmbeddings类来包装SentenceTransformer模型，它会自动实现LangChain要求的embed_documents和embed_query方法。

修正后的完整代码 #

from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
# 替换导入：用LangChain的HuggingFaceEmbeddings替代直接导入SentenceTransformer
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores.chroma import Chroma
import os
import shutil

CHROMA_PATH = "chroma"
DATA_PATH = "data/books"

# 初始化LangChain兼容的嵌入模型
embedder = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

def main():
    generate_data_store()

def generate_data_store():
    documents = load_documents()
    chunks = split_text(documents)
    save_to_chroma(chunks)

def load_documents():
    loader = DirectoryLoader(DATA_PATH, glob="*.md")
    documents = loader.load()
    return documents

def split_text(documents: list[Document]):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,
        chunk_overlap=100,
        length_function=len,
        add_start_index=True,
    )
    chunks = text_splitter.split_documents(documents)
    print(f"Split {len(documents)} documents into {len(chunks)} chunks.")

    document = chunks[10]
    print(document.page_content)
    print(document.metadata)

    return chunks

def save_to_chroma(chunks: list[Document]):
    if os.path.exists(CHROMA_PATH):
        shutil.rmtree(CHROMA_PATH)

    # 直接传入embedder实例，而不是encode方法
    db = Chroma.from_documents(
        chunks, embedder, persist_directory=CHROMA_PATH
    )
    db.persist()
    print(f"Saved {len(chunks)} chunks to {CHROMA_PATH}.")

if __name__ == "__main__":
    main()

关键修改点 #

1. 导入替换：去掉from sentence_transformers import SentenceTransformer，改为导入langchain_community.embeddings.HuggingFaceEmbeddings
2. 嵌入模型初始化：用HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")替代SentenceTransformer("all-MiniLM-L6-v2")，得到LangChain兼容的嵌入对象
3. Chroma调用：在Chroma.from_documents中传入embedder实例，而不是embedder.encode函数

这样修改后，代码就能正常调用SentenceTransformer生成嵌入，并保存到Chroma数据库了。

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