Upload 2 files

app.py

@@ -0,0 +1,445 @@
+import os
+import gradio as gr
+from langchain_community.document_loaders import TextLoader, DirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_openai import ChatOpenAI
+from langchain.prompts import PromptTemplate
+import numpy as np
+import faiss
+from collections import deque
+from langchain_core.embeddings import Embeddings
+import threading
+import queue
+from langchain_core.messages import HumanMessage, AIMessage
+from sentence_transformers import SentenceTransformer
+import pickle
+import torch
+from langchain_core.documents import Document
+import time
+from tqdm import tqdm
+import jieba  # 引入中文分词库
+
+# 获取 OPENROUTER_API_KEY 环境变量
+os.environ["OPENROUTER_API_KEY"] = os.getenv("OPENROUTER_API_KEY", "")
+if not os.environ["OPENROUTER_API_KEY"]:
+    raise ValueError("OPENROUTER_API_KEY 未设置，请在环境变量中配置或在 .env 文件中添加")
+
+# 自定义 SentenceTransformerEmbeddings 类（使用 BAAI/bge-m3 模型，适配 CPU）
+class SentenceTransformerEmbeddings(Embeddings):
+    def __init__(self, model_name="BAAI/bge-m3"):
+        self.model = SentenceTransformer(model_name, device="cpu")
+        self.batch_size = 64
+        self.query_cache = {}
+
+    def embed_documents(self, texts):
+        total_chunks = len(texts)
+        embeddings_list = []
+        batch_size = 1000
+
+        print(f"开始生成嵌入（共 {total_chunks} 个分片，每批 {batch_size} 个分片）")
+        start_time = time.time()
+        with torch.no_grad():
+            for i in tqdm(range(0, total_chunks, batch_size), desc="生成嵌入进度"):
+                batch_start = i
+                batch_end = min(i + batch_size, total_chunks)
+                batch_texts = [text.page_content for text in texts[batch_start:batch_end]]
+                
+                batch_start_time = time.time()
+                batch_emb = self.model.encode(
+                    batch_texts,
+                    normalize_embeddings=True,
+                    batch_size=self.batch_size,
+                    show_progress_bar=True
+                )
+                batch_time = time.time() - batch_start_time
+                
+                if isinstance(batch_emb, torch.Tensor):
+                    embeddings_list.append(batch_emb.cpu().numpy())
+                else:
+                    embeddings_list.append(batch_emb)
+                print(f"完成批次 {i//batch_size + 1}/{total_chunks//batch_size + 1}，处理了 {batch_end - batch_start} 个分片，耗时 {batch_time:.2f} 秒")
+        
+        embeddings_array = np.vstack(embeddings_list)
+        total_time = time.time() - start_time
+        print(f"嵌入生成完成，总耗时 {total_time:.2f} 秒，平均每 1000 个分片耗时 {total_time/total_chunks*1000:.2f} 秒")
+        
+        np.save("embeddings.npy", embeddings_array)
+        return embeddings_array
+
+    def embed_query(self, text):
+        if text in self.query_cache:
+            return self.query_cache[text]
+        with torch.no_grad():
+            emb = self.model.encode([text], normalize_embeddings=True, batch_size=1, show_progress_bar=False)[0]
+            self.query_cache[text] = emb
+        return emb
+
+# 按权重混合检索函数（优化得分和多样性）
+def hybrid_retrieval(query, vector_store, documents, top_n=15, bm25_weight=0.4, semantic_weight=0.6):
+    try:
+        if not documents or not query:
+            raise ValueError("查询或文档列表为空")
+        
+        # 创建文档到 ID 的映射
+        doc_to_id = {id(doc): str(i) for i, doc in enumerate(documents)}
+        id_to_doc = {str(i): doc for i, doc in enumerate(documents)}
+        
+        # 验证索引与文档一致性
+        index_ids = set(vector_store.index_to_docstore_id.values())
+        doc_ids = set(str(i) for i in range(len(documents)))
+        if index_ids != doc_ids:
+            print("警告：索引与文档 ID 不匹配！index_ids:", index_ids, "doc_ids:", doc_ids)
+        
+        # 语义搜索（FAISS）
+        vector_store.index.hnsw.efSearch = 300
+        query_embedding = vector_store.embedding_function.embed_query(query)
+        D, I = vector_store.index.search(np.array([query_embedding]), min(top_n * 2, len(documents)))
+        print(f"FAISS 搜索结果 - 距离 (D): {D[0][:5]}... (前5个)")
+        print(f"FAISS 搜索结果 - 索引 (I): {I[0][:5]}... (前5个)")
+        semantic_results = []
+        if D.size > 0:
+            max_dist = np.max(D) if np.max(D) > 0 else 1.0
+            for idx, dist in zip(I[0], D[0]):
+                if idx == -1:
+                    continue
+                doc_id = vector_store.index_to_docstore_id.get(idx)
+                if doc_id is None:
+                    continue
+                doc = vector_store.docstore.search(doc_id)
+                if doc:
+                    # 归一化距离为相似度（0到1，1为最相似），设置默认值
+                    similarity = 1.0 - (dist / max_dist if max_dist > 0 else 0.5)
+                    semantic_results.append((doc, similarity))
+        else:
+            print("警告：FAISS 距离数组为空，可能索引异常")
+
+        # 使用 doc_id 存储语义得分
+        semantic_scores = {}
+        for doc, score in semantic_results:
+            doc_id = doc_to_id.get(id(doc))
+            if doc_id is not None:
+                semantic_scores[doc_id] = score
+        print(f"语义得分 (semantic_scores): {dict(list(semantic_scores.items())[:5])}... (前5个)")
+
+        # 关键字搜索（BM25，使用 jieba 分词）
+        tokenized_corpus = [list(jieba.cut(doc.page_content)) for doc in documents]
+        bm25 = BM25Okapi(tokenized_corpus)
+        tokenized_query = list(jieba.cut(query))
+        bm25_scores = bm25.get_scores(tokenized_query)
+        print(f"BM25 得分 (bm25_scores): {bm25_scores[:5]}... (前5个)")
+        
+        # 归一化 BM25 得分
+        max_bm25 = max(bm25_scores) if bm25_scores.size > 0 and max(bm25_scores) > 0 else 1.0
+        normalized_bm25_scores = bm25_scores / max_bm25 if max_bm25 > 0 else bm25_scores
+        print(f"归一化 BM25 得分 (normalized_bm25_scores): {normalized_bm25_scores[:5]}... (前5个)")
+
+        # 合并得分
+        combined_scores = {}
+        for i, doc in enumerate(documents):
+            doc_id = str(i)
+            semantic_score = semantic_scores.get(doc_id, 0.0)
+            bm25_score = normalized_bm25_scores[i] if i < len(normalized_bm25_scores) else 0.0
+            combined_score = (bm25_weight * bm25_score) + (semantic_weight * semantic_score)
+            combined_scores[doc_id] = combined_score
+        
+        # 按组合得分排序
+        ranked_ids = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)[:top_n]
+        ranked_docs = [(id_to_doc[doc_id], score) for doc_id, score in ranked_ids]
+        
+        print(f"Query: {query[:100]}... (长度: {len(query)})")
+        print(f"混合检索结果 (数量: {len(ranked_docs)}):")
+        for i, (doc, score) in enumerate(ranked_docs):
+            print(f"  Doc {i}: {doc.page_content[:100]}... (来源: {doc.metadata.get('book', '未知来源')}, 得分: {score:.4f})")
+        
+        return ranked_docs
+    
+    except Exception as e:
+        error_msg = str(e)
+        print(f"错误详情: {error_msg}")
+        raise Exception(f"混合检索失败: {error_msg}")
+
+# 构建 HNSW 索引
+def build_hnsw_index(knowledge_base_path, index_path):
+    print("开始加载文档...")
+    start_time = time.time()
+    loader = DirectoryLoader(knowledge_base_path, glob="*.txt", loader_cls=lambda path: TextLoader(path, encoding="utf-8"), use_multithreading=False)
+    documents = loader.load()
+    load_time = time.time() - start_time
+    print(f"加载完成，共 {len(documents)} 个文档，耗时 {load_time:.2f} 秒")
+    
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    if not os.path.exists("chunks.pkl"):
+        print("开始分片...")
+        start_time = time.time()
+        texts = []
+        total_chars = 0
+        total_bytes = 0
+        for i, doc in enumerate(documents):
+            doc_chunks = text_splitter.split_documents([doc])
+            for chunk in doc_chunks:
+                content = chunk.page_content
+                file_path = chunk.metadata.get("source", "")
+                book_name = os.path.basename(file_path).replace(".txt", "").replace("_", "·")
+                texts.append(Document(page_content=content, metadata={"book": book_name or "未知来源"}))
+                total_chars += len(content)
+                total_bytes += len(content.encode('utf-8'))
+            if i < 5:
+                print(f"文件 {i} 字符数: {len(doc.page_content)}, 字节数: {len(doc.page_content.encode('utf-8'))}, 来源: {file_path}")
+            if (i + 1) % 10 == 0:
+                print(f"分片进度: 已处理 {i + 1}/{len(documents)} 个文件，当前分片总数: {len(texts)}")
+        with open("chunks.pkl", "wb") as f:
+            pickle.dump(texts, f)
+        split_time = time.time() - start_time
+        print(f"分片完成，共 {len(texts)} 个 chunk，总字符数: {total_chars}，总字节数: {total_bytes}，耗时 {split_time:.2f} 秒")
+    else:
+        with open("chunks.pkl", "rb") as f:
+            texts = pickle.load(f)
+        print(f"加载已有分片，共 {len(texts)} 个 chunk")
+    
+    if not os.path.exists("embeddings.npy"):
+        print("开始生成嵌入（使用 BAAI/bge-m3，CPU 模式，分批处理）...")
+        embeddings_array = embeddings.embed_documents(texts)
+        if os.path.exists("embeddings_temp.npy"):
+            os.remove("embeddings_temp.npy")
+        print(f"嵌入生成完成，维度: {embeddings_array.shape}")
+    else:
+        embeddings_array = np.load("embeddings.npy")
+        print(f"加载已有嵌入，维度: {embeddings_array.shape}")
+    
+    dimension = embeddings_array.shape[1]
+    index = faiss.IndexHNSWFlat(dimension, 16)
+    index.hnsw.efConstruction = 100
+    print("开始构建 HNSW 索引...")
+    
+    batch_size = 5000
+    total_vectors = embeddings_array.shape[0]
+    for i in range(0, total_vectors, batch_size):
+        batch = embeddings_array[i:i + batch_size]
+        index.add(batch)
+        print(f"索引构建进度: {min(i + batch_size, total_vectors)} / {total_vectors}")
+    
+    text_embeddings = [(text.page_content, embeddings_array[i]) for i, text in enumerate(texts)]
+    vector_store = FAISS.from_embeddings(text_embeddings, embeddings, normalize_L2=True)
+    vector_store.index = index
+    vector_store.docstore._dict.clear()
+    vector_store.index_to_docstore_id.clear()
+    
+    for i, text in enumerate(texts):
+        doc_id = str(i)
+        vector_store.docstore._dict[doc_id] = text
+        vector_store.index_to_docstore_id[i] = doc_id
+    
+    print("开始保存索引...")
+    vector_store.save_local(index_path)
+    print(f"HNSW 索引已生成并保存到 '{index_path}'")
+    return vector_store, texts  # 返回 vector_store 和分片文本
+
+# 初始化嵌入模型
+embeddings = SentenceTransformerEmbeddings(model_name="BAAI/bge-m3")
+print("已初始化 BAAI/bge-m3 嵌入模型，用于知识库检索（CPU 模式）")
+
+# 加载或生成索引
+index_path = "faiss_index_hnsw_new"
+knowledge_base_path = "knowledge_base"
+
+if not os.path.exists(index_path):
+    if os.path.exists(knowledge_base_path):
+        print("检测到 knowledge_base，正在生成 HNSW 索引...")
+        vector_store, all_documents = build_hnsw_index(knowledge_base_path, index_path)
+    else:
+        raise FileNotFoundError("未找到 'knowledge_base'，请提供知识库数据")
+else:
+    vector_store = FAISS.load_local(index_path, embeddings=embeddings, allow_dangerous_deserialization=True)
+    vector_store.index.hnsw.efSearch = 300
+    print("已加载 HNSW 索引 'faiss_index_hnsw_new'，efSearch 设置为 300")
+    with open("chunks.pkl", "rb") as f:
+        all_documents = pickle.load(f)
+    # 验证 all_documents 内容
+    book_counts = {}
+    for doc in all_documents:
+        book = doc.metadata.get("book", "未知来源")
+        book_counts[book] = book_counts.get(book, 0) + 1
+    print(f"all_documents 书籍分布: {book_counts}")
+
+# 初始化 ChatOpenAI
+llm = ChatOpenAI(
+    model="deepseek/deepseek-r1:free",
+    api_key=os.environ["OPENROUTER_API_KEY"],
+    base_url="https://openrouter.ai/api/v1",
+    timeout=60,
+    temperature=0.3,
+    max_tokens=130000,
+    streaming=True
+)
+
+# 定义提示词模板
+prompt_template = PromptTemplate(
+    input_variables=["context", "question", "chat_history"],
+    template="""  
+    你是一个研究李敖的专家，根据用户提出的问题{question}、最近10轮对话历史{chat_history}以及从李敖相关书籍和评论中检索的至少10篇文本内容{context}回答问题。  
+    在回答时，请注意以下几点：  
+    - 结合李敖的写作风格和思想，筛选出与问题和对话历史最相关的检索内容，避免无关信息。  
+    - 必须在回答中引用至少10篇不同的文本内容，引用格式为[引用: 文本序号]，例如[引用: 1][引用: 2]，并确保每篇文本在回答中都有明确使用。  
+    - 在回答的末尾，必须以“引用文献”标题列出所有引用的文本序号及其内容摘要（每篇不超过50字）以及具体的书目信息（例如书名和章节），格式为：  
+      - 引用文献：  
+        1. [文本 1] 摘要... 出自：书名，第X页/章节。  
+        2. [文本 2] 摘要... 出自：书名，第X页/章节。  
+        （依此类推，至少10篇）  
+    - 如果问题涉及李敖对某人或某事的评价，优先引用李敖的直接言论或文字，并说明出处。  
+    - 回答应结构化、分段落，确保逻辑清晰，语言生动，类似李敖的犀利风格。  
+    - 如果检索内容和历史不足以直接回答问题，可根据李敖的性格和观点推测其可能的看法，但需说明这是推测。  
+    - 只能基于提供的知识库内容{context}和对话历史{chat_history}回答，不得引入外部信息。  
+    - 对于列举类问题，控制在10个要点以内，并优先提供最相关项。  
+    - 如果回答较长，结构化分段总结，分点作答控制在8个点以内。  
+    - 根据对话历史调整回答，避免重复或矛盾。  
+    - 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。  
+    - 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。  
+    - 对于创作类的问题（如写论文），请务必在正文的段落中引用对应的参考编号，例如[引用:3][引用:5]，不能只在文章末尾引用。你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。  
+    - 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在8个点以内，并合并相关的内容。  
+    - 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。  
+    - 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。  
+    - 你的回答应该综合多个相关知识库内容来回答，不能重复引用一个知识库内容。  
+    - 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
+    """
+)
+
+# 对话历史管理类
+class ConversationHistory:
+    def __init__(self, max_length=10):
+        self.history = deque(maxlen=max_length)
+
+    def add_turn(self, question, answer):
+        self.history.append((question, answer))
+
+    def get_history(self):
+        return [(turn[0], turn[1]) for turn in self.history]
+
+    def clear(self):
+        self.history.clear()
+
+# 用户会话状态类
+class UserSession:
+    def __init__(self):
+        self.conversation = ConversationHistory()
+        self.output_queue = queue.Queue()
+        self.stop_flag = threading.Event()
+
+# 生成回答的线程函数
+def generate_answer_thread(question, session):
+    stop_flag = session.stop_flag
+    output_queue = session.output_queue
+    conversation = session.conversation
+    
+    stop_flag.clear()
+    try:
+        history_list = conversation.get_history()
+        history_text = "\n".join([f"问: {q}\n答: {a}" for q, a in history_list]) if history_list else ""
+        query_with_context = f"{history_text}\n当前问题: {question}" if history_text else question
+        
+        # 1. 使用 BAAI/bge-m3 生成查询嵌入
+        start_time = time.time()
+        query_embedding = embeddings.embed_query(query_with_context)
+        embed_time = time.time() - start_time
+        output_queue.put(f"嵌入耗时 (BAAI/bge-m3): {embed_time:.2f} 秒\n")
+        
+        if stop_flag.is_set():
+            output_queue.put("生成已停止")
+            return
+        
+        # 2. 使用混合检索（BM25 + FAISS）
+        start_time = time.time()
+        retrieved_docs_with_scores = hybrid_retrieval(
+            query_with_context,
+            vector_store,
+            all_documents,
+            top_n=15,
+            bm25_weight=0.4,
+            semantic_weight=0.6
+        )
+        retrieval_time = time.time() - start_time
+        output_queue.put(f"混合检索耗时: {retrieval_time:.2f} 秒\n")
+        
+        if stop_flag.is_set():
+            output_queue.put("生成已停止")
+            return
+        
+        # 调整 final_docs 数量，取前 10 篇
+        final_docs = [doc for doc, _ in retrieved_docs_with_scores][:10]
+        if len(final_docs) < 10:
+            output_queue.put(f"警告：仅检索到 {len(final_docs)} 篇文本，可能无法满足引用 10 篇的要求")
+        
+        # 构造 context，包含文本内容和书目信息
+        context = "\n\n".join([f"[文本 {i+1}] {doc.page_content} (出处: {doc.metadata.get('book', '未知来源')})" for i, doc in enumerate(final_docs)])
+        chat_history = [HumanMessage(content=q) if i % 2 == 0 else AIMessage(content=a) 
+                        for i, (q, a) in enumerate(history_list)]
+        prompt = prompt_template.format(context=context, question=question, chat_history=history_text)
+        
+        # 3. 使用 LLM 生成回答
+        answer = ""
+        start_time = time.time()
+        for chunk in llm.stream([HumanMessage(content=prompt)]):
+            if stop_flag.is_set():
+                output_queue.put(answer + "\n\n(生成已停止)")
+                return
+            answer += chunk.content
+            output_queue.put(answer)
+        llm_time = time.time() - start_time
+        output_queue.put(f"\nLLM 生成耗时: {llm_time:.2f} 秒")
+        
+        conversation.add_turn(question, answer)
+        output_queue.put(answer)
+
+    except Exception as e:
+        output_queue.put(f"Error: {str(e)}")
+
+# Gradio 接口函数
+def answer_question(question, session_state):
+    if session_state is None:
+        session_state = UserSession()
+    
+    thread = threading.Thread(target=generate_answer_thread, args=(question, session_state))
+    thread.start()
+    
+    while thread.is_alive() or not session_state.output_queue.empty():
+        try:
+            output = session_state.output_queue.get(timeout=0.1)
+            yield output, session_state
+        except queue.Empty:
+            continue
+    
+    while not session_state.output_queue.empty():
+        yield session_state.output_queue.get(), session_state
+
+def stop_generation(session_state):
+    if session_state is not None:
+        session_state.stop_flag.set()
+    return "生成已停止，正在中止..."
+
+def clear_conversation():
+    return "对话历史已清空，请开始新的对话。", UserSession()
+
+# 创建 Gradio 界面
+with gr.Blocks(title="AI李敖助手") as interface:
+    gr.Markdown("### AI李敖助手")
+    gr.Markdown("基于李敖163本相关书籍构建的知识库，支持上下文关联，记住最近10轮对话，输入问题以获取李敖风格的回答。")
+    
+    session_state = gr.State(value=None)
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            question_input = gr.Textbox(label="请输入您的问题", placeholder="输入您的问题...")
+            submit_button = gr.Button("提交")
+        with gr.Column(scale=1):
+            clear_button = gr.Button("新建对话")
+            stop_button = gr.Button("停止生成")
+    
+    output_text = gr.Textbox(label="回答", interactive=False)
+    
+    submit_button.click(fn=answer_question, inputs=[question_input, session_state], outputs=[output_text, session_state])
+    clear_button.click(fn=clear_conversation, inputs=None, outputs=[output_text, session_state])
+    stop_button.click(fn=stop_generation, inputs=[session_state], outputs=output_text)
+
+# 启动应用
+if __name__ == "__main__":
+    interface.launch(share=True)