如何使用SentenceTransformers创建向量嵌入?

1 投票
1 回答
35 浏览
提问于 2025-04-12 01:24

我发现了这段代码:https://github.com/pixegami/langchain-rag-tutorial/blob/main/create_database.py

这段代码的功能是把文档分成小块,然后为每一块创建一个向量表示,并把这些向量保存到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密钥的情况下实现这个功能,如果可以的话,能指出我哪里出错了吗,那就太好了。

1 个回答

1
  1. 首先,定义你的 embedding 模型,使用 HuggingFaceEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
embedder = HuggingFaceEmbeddings(
    model_name = "sentence-transformers/all-MiniLM-L6-v2"
)
  1. 接着,把这些小块数据嵌入到向量数据库中。
    db = Chroma.from_documents(
        documents=chunks, 
        embedding=embedder, 
        persist_directory=CHROMA_PATH
    )
    db.persist()

撰写回答