|
|
@@ -0,0 +1,131 @@
|
|
|
+import os
|
|
|
+from dotenv import load_dotenv
|
|
|
+from langchain_community.document_loaders import PyMuPDFLoader
|
|
|
+from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
+from langchain_openai import OpenAIEmbeddings
|
|
|
+# from langchain_community.vectorstores import Chroma
|
|
|
+from langchain_chroma import Chroma
|
|
|
+from langchain_core.prompts import ChatPromptTemplate
|
|
|
+from langchain_community.chat_models import ChatTongyi
|
|
|
+from langchain_core.output_parsers import StrOutputParser
|
|
|
+from langchain_openai import ChatOpenAI
|
|
|
+
|
|
|
+load_dotenv(dotenv_path=".env")
|
|
|
+emb_base_url = os.getenv("EMB_MODEL_URL")
|
|
|
+emb_model_name = os.getenv("EMB_MODEL_NAME")
|
|
|
+chat_base_url = os.getenv("CHAT_MODEL_URL")
|
|
|
+chat_model_name = os.getenv("CHAT_MODEL_NAME")
|
|
|
+
|
|
|
+
|
|
|
+# ========== 第一步：加载文档 ==========
|
|
|
+loader = PyMuPDFLoader("中国近代史.pdf")
|
|
|
+pages = loader.load()
|
|
|
+
|
|
|
+# ========== 第二步：清洗数据（可选，根据文档质量决定）==========
|
|
|
+# clean_pages = [clean_pdf_text(page) for page in pages]
|
|
|
+
|
|
|
+# ========== 第三步：分块 ==========
|
|
|
+# splitter = RecursiveCharacterTextSplitter(
|
|
|
+# # 分隔符优先级：段落 → 换行 → 句号 → 空格 → 硬切
|
|
|
+# separators=["\n\n", "\n", "。", "！", "？", " ", ""],
|
|
|
+# # 每个块最大 50 字符
|
|
|
+# chunk_size=180,
|
|
|
+# # 相邻块重叠 10 字符（chunk_size 的 20%）
|
|
|
+# chunk_overlap=10,
|
|
|
+# # 长度计算函数
|
|
|
+# length_function=len
|
|
|
+# )
|
|
|
+# docs = splitter.split_documents(pages)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+# ========== 第四步：向量化 + 存入向量库 ==========
|
|
|
+embedding_model = OpenAIEmbeddings(
|
|
|
+ model=emb_model_name, # 本地模型名称（根据实际部署填写）
|
|
|
+ api_key="not-needed", # 本地服务通常不需要真实 key，填任意值即可
|
|
|
+ base_url=emb_base_url # 本地服务地址，根据实际端口修改
|
|
|
+)
|
|
|
+
|
|
|
+# embedding_model = OpenAIEmbeddings(
|
|
|
+# model="BAAI/bge-m3", # 本地模型名称（根据实际部署填写）
|
|
|
+# api_key="sk-hynxocghcwhevjapdsccungulfggjlaphlxiyenghcotcwux", # 本地服务通常不需要真实 key，填任意值即可
|
|
|
+# base_url="https://api.siliconflow.cn/v1" # 本地服务地址，根据实际端口修改
|
|
|
+# )
|
|
|
+
|
|
|
+#
|
|
|
+# vectorstore = Chroma.from_documents(
|
|
|
+# documents=docs,
|
|
|
+# embedding=embedding_model,
|
|
|
+# collection_metadata={"hnsw:space": "cosine"},
|
|
|
+# persist_directory="./knowledge_db"
|
|
|
+# )
|
|
|
+#
|
|
|
+# exit()
|
|
|
+# ========== 加载已有向量库 ==========
|
|
|
+vectorstore = Chroma(
|
|
|
+ persist_directory="./knowledge_db", # 之前保存的目录
|
|
|
+ embedding_function=embedding_model # 加载时也需要指定 embedding 模型
|
|
|
+)
|
|
|
+
|
|
|
+# ========== 查看向量库前十条chunk ==========
|
|
|
+all_data = vectorstore.get(limit=10)
|
|
|
+
|
|
|
+# 查看结构
|
|
|
+# print(all_data.keys())
|
|
|
+# 输出: dict_keys(['ids', 'embeddings', 'documents', 'metadatas'])
|
|
|
+
|
|
|
+# 查看所有 chunk 的文本内容
|
|
|
+for i, doc in enumerate(all_data['documents']):
|
|
|
+ print(f"\n--- Chunk {i} ---")
|
|
|
+ print(f"ID: {all_data['ids'][i]}")
|
|
|
+ print(f"文本: {doc}")
|
|
|
+ # print(f"元数据: {all_data['metadatas'][i]}")
|
|
|
+
|
|
|
+
|
|
|
+# # ========== 第五步：创建检索器 ==========
|
|
|
+# retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
|
|
|
+
|
|
|
+
|
|
|
+# # ========== 第六步：提问 ==========
|
|
|
+query = "徐中约是谁"
|
|
|
+
|
|
|
+relevant_docs = vectorstore.similarity_search_with_score(query, k=3)
|
|
|
+
|
|
|
+for doc, score in relevant_docs:
|
|
|
+ print(f"{doc.page_content}")
|
|
|
+ print("-"* 30)
|
|
|
+
|
|
|
+
|
|
|
+# ========== 第七步：生成回答 ==========
|
|
|
+context = "\n\n---\n\n".join([content.page_content for content, score in relevant_docs])
|
|
|
+# print(context)
|
|
|
+#
|
|
|
+prompt = ChatPromptTemplate.from_template("""
|
|
|
+你是一个专业的知识库助手。请根据以下上下文回答问题。
|
|
|
+
|
|
|
+**规则：**
|
|
|
+- 只基于提供的上下文回答，不要编造
|
|
|
+- 如果上下文中没有相关信息，直接说「根据现有资料，我找不到这个问题的答案」
|
|
|
+- 回答要简洁直接，引用原文时用引号
|
|
|
+
|
|
|
+**上下文：**
|
|
|
+{context}
|
|
|
+
|
|
|
+**问题：**
|
|
|
+{question}
|
|
|
+""")
|
|
|
+#
|
|
|
+llm = ChatOpenAI(
|
|
|
+ model=chat_model_name,
|
|
|
+ api_key="api_key",
|
|
|
+ base_url=chat_base_url,
|
|
|
+ temperature=0.5,
|
|
|
+ streaming=True
|
|
|
+)
|
|
|
+chain = prompt | llm | StrOutputParser()
|
|
|
+# print(chain)
|
|
|
+#
|
|
|
+# answer = chain.invoke({"context": context, "question": query})
|
|
|
+for chunk in chain.stream({"context": context, "question": query}):
|
|
|
+ print(chunk, end="", flush=True)
|
|
|
+# print(answer)
|
zhouche