Spaces:
Running
Running
Commit
·
fd12765
1
Parent(s):
27b3757
Upload 2 files
Browse files- app.py +311 -0
- 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)
|