如何在使用LangChain创建Vector Store时添加进度条？

使用LangChain创建Vector Store时，若文档数量较多，FAISS.from_documents()方法会因批量生成嵌入耗时较长，我们可以通过手动拆分嵌入生成与向量库构建步骤，结合tqdm库添加进度条，直观展示处理进度。

原代码示例 #

import pprint
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.docstore.document import Document

model = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"
embeddings = HuggingFaceEmbeddings(model_name = model)

def main():
    doc1 = Document(page_content="The sky is blue.",    metadata={"document_id": "10"})
    doc2 = Document(page_content="The forest is green", metadata={"document_id": "62"})
    docs = []
    docs.append(doc1)
    docs.append(doc2)

    for doc in docs:
        doc.metadata['summary'] = 'hello'

    pprint.pprint(docs)
    db = FAISS.from_documents(docs, embeddings)
    db.save_local("faiss_index")
    new_db = FAISS.load_local("faiss_index", embeddings)

    query = "Which color is the sky?"
    docs = new_db.similarity_search_with_score(query)
    print('Retrieved docs:', docs)
    print('Metadata of the most relevant document:', docs[0][0].metadata)

if __name__ == '__main__':
    main()

测试环境与依赖安装 #

pip install langchain==0.1.1 langchain_openai==0.0.2.post1 sentence-transformers==2.2.2 langchain_community==0.0.13 faiss-cpu==1.7.4

添加进度条的解决方案 #

1. 安装进度条依赖

pip install tqdm

2. 修改核心代码

FAISS.from_documents()本质是先对每个文档生成嵌入向量，再构建向量库。我们可以手动拆分这两个步骤，用tqdm包裹嵌入生成过程：

import pprint
from tqdm import tqdm
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.docstore.document import Document

model = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"
embeddings = HuggingFaceEmbeddings(model_name = model)

def main():
    doc1 = Document(page_content="The sky is blue.",    metadata={"document_id": "10"})
    doc2 = Document(page_content="The forest is green", metadata={"document_id": "62"})
    docs = []
    docs.append(doc1)
    docs.append(doc2)

    for doc in docs:
        doc.metadata['summary'] = 'hello'

    pprint.pprint(docs)
    
    # 手动生成嵌入并添加进度条
    texts = [doc.page_content for doc in docs]
    metadatas = [doc.metadata for doc in docs]
    # 遍历生成嵌入，显示进度条
    embeddings_list = []
    for text in tqdm(texts, desc="生成文档嵌入"):
        embedding = embeddings.embed_query(text)
        embeddings_list.append(embedding)
    
    # 用生成好的嵌入创建FAISS向量库
    db = FAISS.from_embeddings(list(zip(embeddings_list, metadatas)), embeddings)
    
    db.save_local("faiss_index")
    new_db = FAISS.load_local("faiss_index", embeddings)

    query = "Which color is the sky?"
    docs = new_db.similarity_search_with_score(query)
    print('Retrieved docs:', docs)
    print('Metadata of the most relevant document:', docs[0][0].metadata)

if __name__ == '__main__':
    main()

说明

• tqdm(texts, desc="生成文档嵌入")会在控制台显示动态进度条，实时展示当前处理的文档数量、耗时和预估剩余时间。
• 该方式完全兼容原有逻辑，仅拆分from_documents的内部步骤，增加进度可视化效果。

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