【AI学习路线 07】RAG系统 - 构建知识增强应用

from sentence_transformers import SentenceTransformer

# 加载嵌入模型
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')

# 生成向量
texts = ["机器学习是人工智能的核心技术", "深度学习是机器学习的分支"]
embeddings = model.encode(texts)

print(f"向量维度: {embeddings.shape}")  # (2, 1024)

# 计算相似度
from numpy import dot
from numpy.linalg import norm
similarity = dot(embeddings[0], embeddings[1]) / (norm(embeddings[0]) * norm(embeddings[1]))
print(f"相似度: {similarity:.4f}")

import chromadb

# 创建客户端
client = chromadb.Client()
collection = client.create_collection("documents")

# 添加文档
collection.add(
    documents=["机器学习是AI的核心", "深度学习是机器学习的分支"],
    metadatas=[{"source": "doc1"}, {"source": "doc2"}],
    ids=["id1", "id2"]
)

# 查询
results = collection.query(
    query_texts=["什么是AI的核心技术"],
    n_results=2
)
print(results)

import faiss
import numpy as np

# 创建索引
dimension = 768
index = faiss.IndexFlatL2(dimension)

# 添加向量
vectors = np.random.randn(1000, dimension).astype('float32')
index.add(vectors)

# 搜索
query = np.random.randn(1, dimension).astype('float32')
distances, indices = index.search(query, k=5)

from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType

# 连接
connections.connect("default", host="localhost", port="19530")

# 创建集合
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768)
]
schema = CollectionSchema(fields, "documents")
collection = Collection("docs", schema)

# 创建索引
index_params = {"metric_type": "L2", "index_type": "IVF_FLAT", "params": {"nlist": 1024}}
collection.create_index("embedding", index_params)

from langchain.document_loaders import PyPDFLoader, TextLoader, DirectoryLoader

# PDF加载
pdf_loader = PyPDFLoader("document.pdf")
pdf_docs = pdf_loader.load()

# 文本加载
text_loader = TextLoader("document.txt")
text_docs = text_loader.load()

# 目录批量加载
dir_loader = DirectoryLoader("./docs", glob="**/*.pdf", loader_cls=PyPDFLoader)
all_docs = dir_loader.load()

from langchain.text_splitter import RecursiveCharacterTextSplitter

# 创建切分器
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,          # 每块大小
    chunk_overlap=50,        # 重叠大小
    separators=["\n\n", "\n", "。", "！", "？", "；", " ", ""]
)

# 切分文档
chunks = text_splitter.split_documents(pdf_docs)

print(f"原始文档数: {len(pdf_docs)}")
print(f"切分后块数: {len(chunks)}")

from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

# 初始化嵌入模型
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-large-zh-v1.5")

# 创建向量存储
vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./chroma_db"
)

# 创建检索器
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 4}
)

# 创建Prompt模板
prompt_template = """
请根据以下上下文回答问题。如果上下文中没有相关信息，请说"我不知道"。

上下文：
{context}

问题：{question}

回答：
"""

PROMPT = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "question"]
)

# 创建RAG链
qa_chain = RetrievalQA.from_chain_type(
    llm=ChatOpenAI(model="gpt-4", temperature=0),
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": PROMPT}
)

# 查询
result = qa_chain({"query": "什么是机器学习？"})
print(result["result"])
print("\n来源文档:", result["source_documents"])

from langchain.retrievers import EnsembleRetriever
from langchain.retrievers import BM25Retriever

# BM25检索器 (关键词)
bm25_retriever = BM25Retriever.from_documents(chunks)
bm25_retriever.k = 4

# 向量检索器 (语义)
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

# 混合检索
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6]
)

from sentence_transformers import CrossEncoder

# 加载重排序模型
reranker = CrossEncoder('BAAI/bge-reranker-large')

def rerank_results(query, documents, top_k=5):
    # 计算重排序分数
    pairs = [[query, doc.page_content] for doc in documents]
    scores = reranker.predict(pairs)
    
    # 按分数排序
    scored_docs = list(zip(documents, scores))
    scored_docs.sort(key=lambda x: x[1], reverse=True)
    
    return [doc for doc, score in scored_docs[:top_k]]

class RAGApplication:
    def __init__(self, docs_path, persist_directory="./chroma_db"):
        self.embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-large-zh-v1.5")
        self.vectorstore = None
        self.qa_chain = None
        
        # 加载或创建向量存储
        if os.path.exists(persist_directory):
            self.load_vectorstore(persist_directory)
        else:
            self.build_index(docs_path, persist_directory)
    
    def build_index(self, docs_path, persist_directory):
        """构建索引"""
        # 加载文档
        loader = DirectoryLoader(docs_path, glob="**/*.pdf", loader_cls=PyPDFLoader)
        docs = loader.load()
        
        # 切分
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
        chunks = text_splitter.split_documents(docs)
        
        # 创建向量存储
        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=self.embeddings,
            persist_directory=persist_directory
        )
    
    def query(self, question, k=4):
        """查询"""
        retriever = self.vectorstore.as_retriever(search_kwargs={"k": k})
        docs = retriever.get_relevant_documents(question)
        
        # 重排序
        reranked_docs = rerank_results(question, docs, top_k=k)
        
        # 生成回答
        context = "\n\n".join([doc.page_content for doc in reranked_docs])
        
        return self.generate_answer(question, context)
    
    def generate_answer(self, question, context):
        """生成回答"""
        # 使用LLM生成
        pass

# Prompt优化
prompt = """
你是一个专业的问答助手。请根据提供的上下文准确回答问题。

要求：
1. 只使用上下文中的信息
2. 如果上下文不足以回答，明确说明
3. 回答要简洁准确
4. 引用相关的上下文片段

上下文：
{context}

问题：{question}
"""