aihuashanying commited on
Commit
1a893e2
·
1 Parent(s): 6c3a484

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -433
app.py DELETED
@@ -1,433 +0,0 @@
1
- import os
2
- import gradio as gr
3
- from langchain_community.document_loaders import TextLoader, DirectoryLoader
4
- from langchain.text_splitter import RecursiveCharacterTextSplitter
5
- from langchain_community.vectorstores import FAISS
6
- from langchain_openai import ChatOpenAI
7
- from langchain.prompts import PromptTemplate
8
- import numpy as np
9
- import faiss
10
- from collections import deque
11
- from langchain_core.embeddings import Embeddings
12
- import threading
13
- import queue
14
- from langchain_core.messages import HumanMessage, AIMessage
15
- from sentence_transformers import SentenceTransformer
16
- import pickle
17
- import torch
18
- from langchain_core.documents import Document
19
- import time
20
- from tqdm import tqdm
21
- from rank_bm25 import BM25Okapi # 新增 BM25 库
22
-
23
- # 获取 OPENROUTER_API_KEY 环境变量
24
- os.environ["OPENROUTER_API_KEY"] = os.getenv("OPENROUTER_API_KEY", "")
25
- if not os.environ["OPENROUTER_API_KEY"]:
26
- raise ValueError("OPENROUTER_API_KEY 未设置,请在环境变量中配置或在 .env 文件中添加")
27
-
28
- # 自定义 SentenceTransformerEmbeddings 类(使用 BAAI/bge-m3 模型,适配 CPU)
29
- class SentenceTransformerEmbeddings(Embeddings):
30
- def __init__(self, model_name="BAAI/bge-m3"):
31
- self.model = SentenceTransformer(model_name, device="cpu")
32
- self.batch_size = 64
33
- self.query_cache = {}
34
-
35
- def embed_documents(self, texts):
36
- total_chunks = len(texts)
37
- embeddings_list = []
38
- batch_size = 1000
39
-
40
- print(f"开始生成嵌入(共 {total_chunks} 个分片,每批 {batch_size} 个分片)")
41
- start_time = time.time()
42
- with torch.no_grad():
43
- for i in tqdm(range(0, total_chunks, batch_size), desc="生成嵌入进度"):
44
- batch_start = i
45
- batch_end = min(i + batch_size, total_chunks)
46
- batch_texts = [text.page_content for text in texts[batch_start:batch_end]]
47
-
48
- batch_start_time = time.time()
49
- batch_emb = self.model.encode(
50
- batch_texts,
51
- normalize_embeddings=True,
52
- batch_size=self.batch_size,
53
- show_progress_bar=True
54
- )
55
- batch_time = time.time() - batch_start_time
56
-
57
- if isinstance(batch_emb, torch.Tensor):
58
- embeddings_list.append(batch_emb.cpu().numpy())
59
- else:
60
- embeddings_list.append(batch_emb)
61
- print(f"完成批次 {i//batch_size + 1}/{total_chunks//batch_size + 1},处理了 {batch_end - batch_start} 个分片,耗时 {batch_time:.2f} 秒")
62
-
63
- embeddings_array = np.vstack(embeddings_list)
64
- total_time = time.time() - start_time
65
- print(f"嵌入生成完成,总耗时 {total_time:.2f} 秒,平均每 1000 个分片耗时 {total_time/total_chunks*1000:.2f} 秒")
66
-
67
- np.save("embeddings.npy", embeddings_array)
68
- return embeddings_array
69
-
70
- def embed_query(self, text):
71
- if text in self.query_cache:
72
- return self.query_cache[text]
73
- with torch.no_grad():
74
- emb = self.model.encode([text], normalize_embeddings=True, batch_size=1, show_progress_bar=False)[0]
75
- self.query_cache[text] = emb
76
- return emb
77
-
78
- # 按权重混合检索函数(优化得分和多样性)
79
- def hybrid_retrieval(query, vector_store, documents, top_n=15, bm25_weight=0.4, semantic_weight=0.6):
80
- try:
81
- if not documents or not query:
82
- raise ValueError("查询或文档列表为空")
83
-
84
- # 创建文档到 ID 的映射
85
- doc_to_id = {id(doc): str(i) for i, doc in enumerate(documents)}
86
- id_to_doc = {str(i): doc for i, doc in enumerate(documents)}
87
-
88
- # 语义搜索(FAISS,恢复 efSearch 为 300)
89
- vector_store.index.hnsw.efSearch = 300 # 增加搜索范围
90
- query_embedding = vector_store.embedding_function.embed_query(query)
91
- D, I = vector_store.index.search(np.array([query_embedding]), min(top_n * 2, len(documents)))
92
- print(f"FAISS 搜索结果 - 距离 (D): {D[0][:5]}... (前5个)")
93
- print(f"FAISS 搜索结果 - 索引 (I): {I[0][:5]}... (前5个)")
94
- semantic_results = []
95
- if D.size > 0 and np.max(D) > 0:
96
- for idx, dist in zip(I[0], D[0]):
97
- if idx == -1:
98
- continue
99
- doc_id = vector_store.index_to_docstore_id.get(idx)
100
- if doc_id is None:
101
- continue
102
- doc = vector_store.docstore.search(doc_id)
103
- if doc:
104
- # 归一化距离为相似度(0到1,1为最相似)
105
- similarity = 1.0 - (dist / np.max(D) if np.max(D) > 0 else 0.0)
106
- semantic_results.append((doc, similarity))
107
- else:
108
- print("警告:FAISS 距离数组为空或全为0,可能索引异常")
109
-
110
- # 使用 doc_id 存储语义得分
111
- semantic_scores = {}
112
- for doc, score in semantic_results:
113
- doc_id = doc_to_id.get(id(doc))
114
- if doc_id is not None:
115
- semantic_scores[doc_id] = score
116
- print(f"语义得分 (semantic_scores): {dict(list(semantic_scores.items())[:5])}... (前5个)")
117
-
118
- # 关键字搜索(BM25)
119
- tokenized_corpus = [doc.page_content.split() for doc in documents]
120
- bm25 = BM25Okapi(tokenized_corpus)
121
- tokenized_query = query.split()
122
- bm25_scores = bm25.get_scores(tokenized_query)
123
- print(f"BM25 得分 (bm25_scores): {bm25_scores[:5]}... (前5个)")
124
-
125
- # 归一化 BM25 得分
126
- max_bm25 = max(bm25_scores) if bm25_scores.size > 0 and max(bm25_scores) > 0 else 1.0
127
- normalized_bm25_scores = bm25_scores / max_bm25 if max_bm25 > 0 else bm25_scores
128
- print(f"归一化 BM25 得分 (normalized_bm25_scores): {normalized_bm25_scores[:5]}... (前5个)")
129
-
130
- # 合并得分
131
- combined_scores = {}
132
- for i, doc in enumerate(documents):
133
- doc_id = str(i)
134
- semantic_score = semantic_scores.get(doc_id, 0.0)
135
- bm25_score = normalized_bm25_scores[i] if i < len(normalized_bm25_scores) else 0.0
136
- combined_score = (bm25_weight * bm25_score) + (semantic_weight * semantic_score)
137
- combined_scores[doc_id] = combined_score
138
-
139
- # 按组合得分排序
140
- ranked_ids = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)[:top_n]
141
- ranked_docs = [(id_to_doc[doc_id], score) for doc_id, score in ranked_ids]
142
-
143
- print(f"Query: {query[:100]}... (长度: {len(query)})")
144
- print(f"混合检索结果 (数量: {len(ranked_docs)}):")
145
- for i, (doc, score) in enumerate(ranked_docs):
146
- print(f" Doc {i}: {doc.page_content[:100]}... (来源: {doc.metadata.get('book', '未知来源')}, 得分: {score:.4f})")
147
-
148
- return ranked_docs
149
-
150
- except Exception as e:
151
- error_msg = str(e)
152
- print(f"错误详情: {error_msg}")
153
- raise Exception(f"混合检索失败: {error_msg}")
154
-
155
- # 构建 HNSW 索引
156
- def build_hnsw_index(knowledge_base_path, index_path):
157
- print("开始加载文档...")
158
- start_time = time.time()
159
- loader = DirectoryLoader(knowledge_base_path, glob="*.txt", loader_cls=lambda path: TextLoader(path, encoding="utf-8"), use_multithreading=False)
160
- documents = loader.load()
161
- load_time = time.time() - start_time
162
- print(f"加载完成,共 {len(documents)} 个文档,耗时 {load_time:.2f} 秒")
163
-
164
- text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
165
- if not os.path.exists("chunks.pkl"):
166
- print("开始分片...")
167
- start_time = time.time()
168
- texts = []
169
- total_chars = 0
170
- total_bytes = 0
171
- for i, doc in enumerate(documents):
172
- doc_chunks = text_splitter.split_documents([doc])
173
- for chunk in doc_chunks:
174
- content = chunk.page_content
175
- file_path = chunk.metadata.get("source", "")
176
- book_name = os.path.basename(file_path).replace(".txt", "").replace("_", "·")
177
- texts.append(Document(page_content=content, metadata={"book": book_name or "未知来源"}))
178
- total_chars += len(content)
179
- total_bytes += len(content.encode('utf-8'))
180
- if i < 5:
181
- print(f"文件 {i} 字符数: {len(doc.page_content)}, 字节数: {len(doc.page_content.encode('utf-8'))}, 来源: {file_path}")
182
- if (i + 1) % 10 == 0:
183
- print(f"分片进度: 已处理 {i + 1}/{len(documents)} 个文件,当前分片总数: {len(texts)}")
184
- with open("chunks.pkl", "wb") as f:
185
- pickle.dump(texts, f)
186
- split_time = time.time() - start_time
187
- print(f"分片完成,共 {len(texts)} 个 chunk,总字符数: {total_chars},总字节数: {total_bytes},耗时 {split_time:.2f} 秒")
188
- else:
189
- with open("chunks.pkl", "rb") as f:
190
- texts = pickle.load(f)
191
- print(f"加载已有分片,共 {len(texts)} 个 chunk")
192
-
193
- if not os.path.exists("embeddings.npy"):
194
- print("开始生成嵌入(使用 BAAI/bge-m3,CPU 模式,分批处理)...")
195
- embeddings_array = embeddings.embed_documents(texts)
196
- if os.path.exists("embeddings_temp.npy"):
197
- os.remove("embeddings_temp.npy")
198
- print(f"嵌入生成完成,维度: {embeddings_array.shape}")
199
- else:
200
- embeddings_array = np.load("embeddings.npy")
201
- print(f"加载已有嵌入,维度: {embeddings_array.shape}")
202
-
203
- dimension = embeddings_array.shape[1]
204
- index = faiss.IndexHNSWFlat(dimension, 16)
205
- index.hnsw.efConstruction = 100
206
- print("开始构建 HNSW 索引...")
207
-
208
- batch_size = 5000
209
- total_vectors = embeddings_array.shape[0]
210
- for i in range(0, total_vectors, batch_size):
211
- batch = embeddings_array[i:i + batch_size]
212
- index.add(batch)
213
- print(f"索引构建进度: {min(i + batch_size, total_vectors)} / {total_vectors}")
214
-
215
- text_embeddings = [(text.page_content, embeddings_array[i]) for i, text in enumerate(texts)]
216
- vector_store = FAISS.from_embeddings(text_embeddings, embeddings, normalize_L2=True)
217
- vector_store.index = index
218
- vector_store.docstore._dict.clear()
219
- vector_store.index_to_docstore_id.clear()
220
-
221
- for i, text in enumerate(texts):
222
- doc_id = str(i)
223
- vector_store.docstore._dict[doc_id] = text
224
- vector_store.index_to_docstore_id[i] = doc_id
225
-
226
- print("开始保存索引...")
227
- vector_store.save_local(index_path)
228
- print(f"HNSW 索引已生成并保存到 '{index_path}'")
229
- return vector_store, texts # 返回 vector_store 和分片文本
230
-
231
- # 初始化嵌入模型
232
- embeddings = SentenceTransformerEmbeddings(model_name="BAAI/bge-m3")
233
- print("已初始化 BAAI/bge-m3 嵌入模型,用于知识库检索(CPU 模式)")
234
-
235
- # 加载或生成索引
236
- index_path = "faiss_index_hnsw_new"
237
- knowledge_base_path = "knowledge_base"
238
-
239
- if not os.path.exists(index_path):
240
- if os.path.exists(knowledge_base_path):
241
- print("检测到 knowledge_base,正在生成 HNSW 索引...")
242
- vector_store, all_documents = build_hnsw_index(knowledge_base_path, index_path)
243
- else:
244
- raise FileNotFoundError("未找到 'knowledge_base',请提供知识库数据")
245
- else:
246
- vector_store = FAISS.load_local(index_path, embeddings=embeddings, allow_dangerous_deserialization=True)
247
- vector_store.index.hnsw.efSearch = 300 # 初始设置为300,可根据需要调整
248
- print("已加载 HNSW 索引 'faiss_index_hnsw_new',efSearch 设置为 300")
249
- with open("chunks.pkl", "rb") as f:
250
- all_documents = pickle.load(f)
251
- print(f"加载已有分片,共 {len(all_documents)} 个 chunk")
252
-
253
- # 初始化 ChatOpenAI
254
- llm = ChatOpenAI(
255
- model="deepseek/deepseek-r1:free",
256
- api_key=os.environ["OPENROUTER_API_KEY"],
257
- base_url="https://openrouter.ai/api/v1",
258
- timeout=60,
259
- temperature=0.3,
260
- max_tokens=130000,
261
- streaming=True
262
- )
263
-
264
- # 定义提示词模板
265
- prompt_template = PromptTemplate(
266
- input_variables=["context", "question", "chat_history"],
267
- template="""
268
- 你是一个研究李敖的专家,根据用户提出的问题{question}、最近10轮对话历史{chat_history}以及从李敖相关书籍和评论中检索的至少10篇文本内容{context}回答问题。
269
- 在回答时,请注意以下几点:
270
- - 结合李敖的写作风格和思想,筛选出与问题和对话历史最相关的检索内容,避免无关信息。
271
- - 必须在回答中引用至少10篇不同的文本内容,引用格式为[引用: 文本序号],例如[引用: 1][引用: 2],并确保每篇文本在回答中都有明确使用。
272
- - 在回答的末尾,必须以“引用文献”标题列出所有引用的文本序号及其内容摘要(每篇不超过50字)以及具体的书目信息(例如书名和章节),格式为:
273
- - 引用文献:
274
- 1. [文本 1] 摘要... 出自:书名,第X页/章节。
275
- 2. [文本 2] 摘要... 出自:书名,第X页/章节。
276
- (依此类推,至少10篇)
277
- - 如果问题涉及李敖对某人或某事的评价,优先引用李敖的直接言论或文字,并说明出处。
278
- - 回答应结构化、分段落,确保逻辑清晰,语言生动,类似李敖的犀利风格。
279
- - 如果检索内容和历史不足以直接回答问题,可根据李敖的性格和观点推测其可能的看法,但需说明这是推测。
280
- - 只能基于提供的知识库内容{context}和对话历史{chat_history}回答,不得引入外部信息。
281
- - 对于列举类问题,控制在10个要点以内,并优先提供最相关项。
282
- - 如果回答较长,结构化分段总结,分点作答控制在8个点以内。
283
- - 根据对话历史调整回答,避免重复或矛盾。
284
- - 并非搜索结果的所有内容都与用户的问题密切相关,你需要结合问题,对搜索结果进行甄别、筛选。
285
- - 对于列举类的问题(如列举所有航班信息),尽量将答案控制在10个要点以内,并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项;如非必要,不要主动告诉用户搜索结果未提供的内容。
286
- - 对于创作类的问题(如写论文),请务必在正文的段落中引用对应的参考编号,例如[引用:3][引用:5],不能只在文章末尾引用。你需要解读并概括用户的题目要求,选择合适的格式,充分利用搜索结果并抽取重要信息,生成符合用户要求、极具思想深度��富有创造力与专业性的答案。你的创作篇幅需要尽可能延长,对于每一个要点的论述要推测用户的意图,给出尽可能多角度的回答要点,且务必信息量大、论述详尽。
287
- - 如果回答很长,请尽量结构化、分段落总结。如果需要分点作答,尽量控制在8个点以内,并合并相关的内容。
288
- - 对于客观类的问答,如果问题的答案非常简短,可以适当补充一到两句相关信息,以丰富内容。
289
- - 你需要根据用户要求和回答内容选择合适、美观的回答格式,确保可读性强。
290
- - 你的回答应该综合多个相关知识库内容来回答,不能重复引用一个知识库内容。
291
- - 除非用户要求,否则你回答的语言需要和用户提问的语言保持一致。
292
- """
293
- )
294
-
295
- # 对话历史管理类
296
- class ConversationHistory:
297
- def __init__(self, max_length=10):
298
- self.history = deque(maxlen=max_length)
299
-
300
- def add_turn(self, question, answer):
301
- self.history.append((question, answer))
302
-
303
- def get_history(self):
304
- return [(turn[0], turn[1]) for turn in self.history]
305
-
306
- def clear(self):
307
- self.history.clear()
308
-
309
- # 用户会话状态类
310
- class UserSession:
311
- def __init__(self):
312
- self.conversation = ConversationHistory()
313
- self.output_queue = queue.Queue()
314
- self.stop_flag = threading.Event()
315
-
316
- # 生成回答的线程函数
317
- def generate_answer_thread(question, session):
318
- stop_flag = session.stop_flag
319
- output_queue = session.output_queue
320
- conversation = session.conversation
321
-
322
- stop_flag.clear()
323
- try:
324
- history_list = conversation.get_history()
325
- history_text = "\n".join([f"问: {q}\n答: {a}" for q, a in history_list]) if history_list else ""
326
- query_with_context = f"{history_text}\n当前问题: {question}" if history_text else question
327
-
328
- # 1. 使用 BAAI/bge-m3 生成查询嵌入
329
- start_time = time.time()
330
- query_embedding = embeddings.embed_query(query_with_context)
331
- embed_time = time.time() - start_time
332
- output_queue.put(f"嵌入耗时 (BAAI/bge-m3): {embed_time:.2f} 秒\n")
333
-
334
- if stop_flag.is_set():
335
- output_queue.put("生成已停止")
336
- return
337
-
338
- # 2. 使用混合检索(BM25 + FAISS)
339
- start_time = time.time()
340
- retrieved_docs_with_scores = hybrid_retrieval(
341
- query_with_context,
342
- vector_store,
343
- all_documents,
344
- top_n=15,
345
- bm25_weight=0.4,
346
- semantic_weight=0.6
347
- )
348
- retrieval_time = time.time() - start_time
349
- output_queue.put(f"混合检索耗时: {retrieval_time:.2f} 秒\n")
350
-
351
- if stop_flag.is_set():
352
- output_queue.put("生成已停止")
353
- return
354
-
355
- # 调整 final_docs 数量,取前 10 篇
356
- final_docs = [doc for doc, _ in retrieved_docs_with_scores][:10]
357
- if len(final_docs) < 10:
358
- output_queue.put(f"警告:仅检索到 {len(final_docs)} 篇文本,可能无法满足引用 10 篇的要求")
359
-
360
- # 构造 context,包含文本内容和书目信息
361
- context = "\n\n".join([f"[文本 {i+1}] {doc.page_content} (出处: {doc.metadata.get('book', '未知来源')})" for i, doc in enumerate(final_docs)])
362
- chat_history = [HumanMessage(content=q) if i % 2 == 0 else AIMessage(content=a)
363
- for i, (q, a) in enumerate(history_list)]
364
- prompt = prompt_template.format(context=context, question=question, chat_history=history_text)
365
-
366
- # 3. 使用 LLM 生成回答
367
- answer = ""
368
- start_time = time.time()
369
- for chunk in llm.stream([HumanMessage(content=prompt)]):
370
- if stop_flag.is_set():
371
- output_queue.put(answer + "\n\n(生成已停止)")
372
- return
373
- answer += chunk.content
374
- output_queue.put(answer)
375
- llm_time = time.time() - start_time
376
- output_queue.put(f"\nLLM 生成耗时: {llm_time:.2f} 秒")
377
-
378
- conversation.add_turn(question, answer)
379
- output_queue.put(answer)
380
-
381
- except Exception as e:
382
- output_queue.put(f"Error: {str(e)}")
383
-
384
- # Gradio 接口函数
385
- def answer_question(question, session_state):
386
- if session_state is None:
387
- session_state = UserSession()
388
-
389
- thread = threading.Thread(target=generate_answer_thread, args=(question, session_state))
390
- thread.start()
391
-
392
- while thread.is_alive() or not session_state.output_queue.empty():
393
- try:
394
- output = session_state.output_queue.get(timeout=0.1)
395
- yield output, session_state
396
- except queue.Empty:
397
- continue
398
-
399
- while not session_state.output_queue.empty():
400
- yield session_state.output_queue.get(), session_state
401
-
402
- def stop_generation(session_state):
403
- if session_state is not None:
404
- session_state.stop_flag.set()
405
- return "生成已停止,正在中止..."
406
-
407
- def clear_conversation():
408
- return "对话历史已清空,请开始新的对话。", UserSession()
409
-
410
- # 创建 Gradio 界面
411
- with gr.Blocks(title="AI李敖助手") as interface:
412
- gr.Markdown("### AI李敖助手")
413
- gr.Markdown("基于李敖163本相关书籍构建的知识库,支持上下文关联,记住最近10轮对话,输入问题以获取李敖风格的回答。")
414
-
415
- session_state = gr.State(value=None)
416
-
417
- with gr.Row():
418
- with gr.Column(scale=3):
419
- question_input = gr.Textbox(label="请输入您的问题", placeholder="输入您的问题...")
420
- submit_button = gr.Button("提交")
421
- with gr.Column(scale=1):
422
- clear_button = gr.Button("新建对话")
423
- stop_button = gr.Button("停止生成")
424
-
425
- output_text = gr.Textbox(label="回答", interactive=False)
426
-
427
- submit_button.click(fn=answer_question, inputs=[question_input, session_state], outputs=[output_text, session_state])
428
- clear_button.click(fn=clear_conversation, inputs=None, outputs=[output_text, session_state])
429
- stop_button.click(fn=stop_generation, inputs=[session_state], outputs=output_text)
430
-
431
- # 启动应用
432
- if __name__ == "__main__":
433
- interface.launch(share=True)