aihuashanying commited on
Commit
fd12765
·
1 Parent(s): 27b3757

Upload 2 files

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