ZhouYI/03_RAG优化/main.py

from agent.config import load_config

from langchain_text_splitters import RecursiveCharacterTextSplitter  
from langchain_experimental.text_splitter import SemanticChunker
from langchain_community.embeddings import DashScopeEmbeddings
from agent.config import load_config

from langchain_community.document_loaders import PyMuPDFLoader


from agent.config import load_config
from langchain_community.vectorstores import Chroma
import os
def main():
    loader=PyMuPDFLoader("car_info.pdf")

    embedding_model=DashScopeEmbeddings(
        model="text-embedding-v4",
        dashscope_api_key=load_config().api_key,
    )
    text_splitter = RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", "。", "！", "？", ".", " ",""],
        chunk_size=500,
        chunk_overlap=100,
        length_function=len,
    )
    chuncks=text_splitter.split_documents(loader.load())

    vectorstore=Chroma.from_documents(
        documents=chuncks,
        embedding=embedding_model,
        collection_metadata={"hnsw:space": "cosine"},
        persist_directory="./chroma_db"
    )

    retrever=vectorstore.as_retriever(search_kwargs={"k": 3})

    results=retrever.invoke("汽车的主要组成部分有哪些？")
    
    print(results.page_content)





if __name__ == "__main__":
    main()