新增 rag_demo.py 文件

02_rag_homework/rag_demo.py

@@ -0,0 +1,146 @@
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import DashScopeEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_community.chat_models import ChatTongyi
+from langchain_core.output_parsers import StrOutputParser
+from dotenv import load_dotenv
+import re
+import os
+
+load_dotenv()
+api_key = os.getenv('QWEN_API_KEY')
+
+# ============================================================
+# 1. 加载 PDF 文档
+# ============================================================
+loader = PyMuPDFLoader("D:/code/shuheAI/02_RAG/华为擎云W585X 用户指南-(PGUX,KOS&UOS_02,zh-cn).pdf")
+pdf_pages = loader.load()
+
+print(f"文档类型：{type(pdf_pages)}")
+print(f"PDF 共{len(pdf_pages)}页")
+
+# 查看第一页的内容和元数据
+first_page = pdf_pages[0]
+print(f"元数据：{first_page.metadata}")
+print(f"内容预览：{first_page.page_content[:200]}")
+
+# ============================================================
+# 2. 清洗 PDF 文本
+# ============================================================
+def clean_pdf_text(text: str) -> str:
+    """清洗 PDF 解析出的文本，去除常见噪声"""
+    # 删除非中文字符之间的换行符
+    text = re.sub(r'[^一](\n)[^一]', lambda m: m.group(0).replace('\n', ''), text)
+    # 删除项目符号和多余空格
+    text = text.replace('•', '').replace('  ', ' ')
+    # 删除连续的换行符（保留一个）
+    text = re.sub(r'\n{2,}', '\n', text)
+    return text.strip()
+
+# 对所有页面清洗
+for page in pdf_pages:
+    page.page_content = clean_pdf_text(page.page_content)
+
+# ============================================================
+# 3. 文本分块
+# ============================================================
+splitter = RecursiveCharacterTextSplitter(
+    separators=["\n\n", "\n", "。", "！", "？", " ", ""],
+    chunk_size=50,
+    chunk_overlap=10,
+    length_function=len
+)
+
+split_docs = splitter.split_documents(pdf_pages)
+print(f"切分后的文件数量：{len(split_docs)}")
+print(f"切分后的字符数（可以用来大致评估 token 数）：{sum([len(doc.page_content) for doc in split_docs])}")
+
+# 过滤掉 page_content 为空或仅含空白的文档
+valid_docs = [
+    doc for doc in split_docs
+    if doc.page_content and doc.page_content.strip()
+]
+print(f"有效块数量：{len(valid_docs)}")
+print(f"总字符数（可大致评估 Token 数）：{sum(len(d.page_content) for d in valid_docs)}")
+
+# ============================================================
+# 4. 初始化 Embedding 模型
+# ============================================================
+embedding_model = DashScopeEmbeddings(
+    model="text-embedding-v3",
+    dashscope_api_key=api_key
+)
+
+# 单条文本向量化
+text = "RAG系统搭建实战"
+embedding = embedding_model.embed_query(text)
+print(f"向量维度：{len(embedding)}")
+print(f"前5个值：{embedding[:5]}")
+
+# ============================================================
+# 5. 向量化 + 存入向量库
+# ============================================================
+persist_dir = "./my_knowledge_db"
+
+if os.path.exists(persist_dir) and os.listdir(persist_dir):
+    print(f"加载已有向量库：{persist_dir}")
+    vectordb = Chroma(
+        persist_directory=persist_dir,
+        embedding_function=embedding_model
+    )
+else:
+    print(f"创建新向量库：{persist_dir}")
+    vectordb = Chroma.from_documents(
+        documents=valid_docs,
+        embedding=embedding_model,
+        collection_metadata={"hnsw:space": "cosine"},  # 余弦相似度
+        persist_directory=persist_dir
+    )
+
+# 创建检索器，设置返回 Top-3 最相关文档
+retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+
+# ============================================================
+# 6. 提问
+# ============================================================
+query = "如何进入BIOS设置？"
+relevant_docs = retriever.invoke(query)
+
+for i, doc in enumerate(relevant_docs):
+    print(f"--- 结果{i+1} ---")
+    print(f"内容：{doc.page_content[:100]}...")
+    print(f"来源：{doc.metadata}")
+    print()
+
+# 带分数的相似度检索（分数越低越相似，0 表示完全匹配）
+results = vectordb.similarity_search_with_score(query, k=3)
+for doc, score in results:
+    print(f"内容：{doc.page_content[:100]}... | 相似度分数：{score:.4f}")
+
+# ============================================================
+# 7. 生成回答
+# ============================================================
+context = "\n\n---\n\n".join([d.page_content for d in relevant_docs])
+
+prompt = ChatPromptTemplate.from_template("""
+你是一个专业的知识库助手。请根据以下上下文回答问题。
+
+**规则：**
+- 只基于提供的上下文回答，不要编造
+- 如果上下文中没有相关信息，直接说「根据现有资料，我找不到这个问题的答案」
+- 回答要简洁直接，引用原文时用引号
+
+**上下文：**
+{context}
+
+**问题：**
+{question}
+""")
+
+llm = ChatTongyi(model="qwen-plus", dashscope_api_key=api_key)
+chain = prompt | llm | StrOutputParser()
+
+answer = chain.invoke({"context": context, "question": query})
+print(answer)