test.py

<<<<<<< HEAD
import torch
print(torch.__version__)  # 如 2.4.0+cu118
print(torch.cuda.is_available())  # 应返回 True
print(torch.cuda.get_device_name(0))  # 应返回 GPU 型号
=======
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.chains import RetrievalQA
from langchain_core.embeddings import Embeddings
from langchain.prompts import PromptTemplate
import requests
import numpy as np
import json
import faiss
from langchain_community.embeddings import OllamaEmbeddings

# 自定义 SiliconFlow 嵌入类
class SiliconFlowEmbeddings(Embeddings):
    def __init__(self, model="BAAI/bge-m3", api_key=None):
        self.model = model
        self.api_key = api_key

    def embed_documents(self, texts):
        return self._get_embeddings(texts)

    def embed_query(self, text):
        return self._get_embeddings([text])[0]

    def _get_embeddings(self, texts):
        url = "https://api.siliconflow.cn/v1/embeddings"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.model,
            "input": texts
        }
        response = requests.post(url, json=payload, headers=headers, timeout=30)
        if response.status_code == 200:
            data = response.json()
            return np.array([item["embedding"] for item in data["data"]])
        else:
            raise Exception(f"API 调用失败: {response.status_code}, {response.text}")

# SiliconFlow 重排序函数
def rerank_documents(query, documents, api_key, top_n=10):
    url = "https://api.siliconflow.cn/v1/rerank"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    doc_texts = [doc.page_content for doc in documents]
    payload = {
        "model": "BAAI/bge-reranker-v2-m3",
        "query": query,
        "documents": doc_texts,
        "top_n": top_n
    }
    response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=30)
    if response.status_code == 200:
        result = response.json()
        reranked_results = result.get("results", [])
        if not reranked_results:
            raise Exception("重排序结果为空")
        reranked_docs_with_scores = [
            (documents[res["index"]], res["relevance_score"])
            for res in reranked_results
        ]
        return reranked_docs_with_scores
    else:
        raise Exception(f"重排序失败: {response.status_code}, {response.text}")

# 设置 API Keys
os.environ["SILICONFLOW_API_KEY"] = os.getenv("SILICONFLOW_API_KEY", "sk-cigytzyzghoziznvniugfihuicjcgmborusgodktydremtvd")
os.environ["OPENROUTER_API_KEY"] = os.getenv("OPENROUTER_API_KEY", "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")

# 初始化嵌入模型
embeddings = OllamaEmbeddings(model="bge-m3", base_url="http://localhost:11434")

# 从 knowledge_base 生成 HNSW 索引
def build_hnsw_index(knowledge_base_path, index_path):
    loader = DirectoryLoader(
        knowledge_base_path,
        glob="*.txt",
        loader_cls=lambda path: TextLoader(path, encoding="utf-8")
    )
    documents = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    texts = text_splitter.split_documents(documents)
    
    # 使用 FAISS.from_documents 创建向量存储
    vector_store = FAISS.from_documents(texts, embeddings)
    
    # 获取嵌入并转换为 HNSW
    embeddings_array = np.array(embeddings.embed_documents([doc.page_content for doc in texts]))
    dimension = embeddings_array.shape[1]
    index = faiss.IndexHNSWFlat(dimension, 16)  # M=16
    index.hnsw.efConstruction = 100
    index.hnsw.efSearch = 50
    index.add(embeddings_array)
    
    # 更新 FAISS 的索引
    vector_store.index = index
    vector_store.save_local(index_path)
    print(f"HNSW 索引已生成并保存到 '{index_path}'")
    return vector_store

# 将已有 faiss_index 转为 HNSW
def convert_to_hnsw(existing_index_path, new_index_path):
    # 加载现有索引
    old_vector_store = FAISS.load_local(existing_index_path, embeddings=embeddings, allow_dangerous_deserialization=True)
    
    # 获取文档内容
    if hasattr(old_vector_store, 'docstore') and hasattr(old_vector_store.docstore, '_dict'):
        docs = list(old_vector_store.docstore._dict.values())
        doc_texts = [doc.page_content if hasattr(doc, 'page_content') else str(doc) for doc in docs]
    else:
        doc_ids = list(old_vector_store.index_to_docstore_id.keys())
        doc_texts = [old_vector_store.docstore._dict[old_vector_store.index_to_docstore_id[i]].page_content 
                     if hasattr(old_vector_store.docstore._dict[old_vector_store.index_to_docstore_id[i]], 'page_content') 
                     else str(old_vector_store.docstore._dict[old_vector_store.index_to_docstore_id[i]]) 
                     for i in doc_ids]
    
    # 使用全局 embeddings 对象生成嵌入
    embeddings_array = np.array(embeddings.embed_documents(doc_texts))
    
    # 创建 HNSW 索引
    dimension = embeddings_array.shape[1]
    index = faiss.IndexHNSWFlat(dimension, 16)  # M=16
    index.hnsw.efConstruction = 100
    index.hnsw.efSearch = 50
    index.add(embeddings_array)
    
    # 创建新的 FAISS 向量存储，注意不直接传递 index，而是稍后赋值
    new_vector_store = FAISS.from_texts(doc_texts, embeddings)
    new_vector_store.index = index  # 直接替换索引
    new_vector_store.save_local(new_index_path)
    print(f"已将 '{existing_index_path}' 转换为 HNSW 并保存到 '{new_index_path}'")
    return new_vector_store

# 加载或生成索引
index_path = "faiss_index_hnsw"
knowledge_base_path = "knowledge_base"

if not os.path.exists(index_path):
    if os.path.exists("faiss_index"):
        print("检测到已有 faiss_index，正在转换为 HNSW...")
        vector_store = convert_to_hnsw("faiss_index", index_path)
    elif os.path.exists(knowledge_base_path):
        print("检测到 knowledge_base，正在生成 HNSW 索引...")
        vector_store = build_hnsw_index(knowledge_base_path, index_path)
    else:
        raise FileNotFoundError("未找到 'faiss_index' 或 'knowledge_base'，请提供知识库数据")
else:
    vector_store = FAISS.load_local(index_path, embeddings=embeddings, allow_dangerous_deserialization=True)
    print("已加载 HNSW 索引 'faiss_index_hnsw'")

# 初始化 ChatOpenAI 使用 OpenRouter
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=88888,
)

# 定义提示词模板
prompt_template = PromptTemplate(
    input_variables=["context", "question"],
    template="""  
    你是一个研究李敖的专家，根据用户提出的问题{question}以及从李敖相关书籍和评论中检索的内容{context}回答问题。  

    在回答时，请注意以下几点：  
    - 结合李敖的写作风格和思想，筛选出与问题最相关的检索内容，避免无关信息。  
    - 如果问题涉及李敖对某人或某事的评价，优先引用李敖的直接言论或文字，并说明出处。  
    - 回答应结构化、分段落，确保逻辑清晰，语言生动，类似李敖的犀利风格。  
    - 如果检索内容不足以直接回答问题，可根据李敖的性格和观点推测其可能的看法，但需说明这是推测。  
    - 列出引用的书籍或文章名称及章节（如有），如《李敖大全集》第X卷或具体书名。  
    - 只能基于提供的知识库内容{context}回答，不得引入外部信息。  
    - 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
    - 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。
    - 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
    - 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
    - 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
    - 你的回答应该综合多个相关知识库内容来回答，不能重复引用一个知识库内容。
    - 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
    """
)

# 创建检索问答链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vector_store.as_retriever(search_kwargs={"k": 30}),
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt_template}
)

# 定义 Gradio 接口函数
def answer_question(question):
    try:
        # Step 1: FAISS 初始检索
        initial_docs_with_scores = vector_store.similarity_search_with_score(question, k=30)
        print(f"初始检索数量: {len(initial_docs_with_scores)}")
        
        # FAISS 返回的是距离，转换为相似度
        similarities = [1 - score for _, score in initial_docs_with_scores]
        print(f"相似度范围: {min(similarities):.4f} - {max(similarities):.4f}")
        
        # 打印前 5 个文档内容和相似度
        for i, (doc, score) in enumerate(initial_docs_with_scores[:5]):
            print(f"Top {i+1} - 相似度: {1 - score:.4f}, 内容: {doc.page_content[:100]}")

        # Step 2: 动态阈值过滤
        similarity_threshold = max(similarities) * 0.8
        filtered_docs_with_scores = [
            (doc, 1 - score)
            for doc, score in initial_docs_with_scores
            if (1 - score) >= similarity_threshold
        ]
        if len(filtered_docs_with_scores) < 5:
            filtered_docs_with_scores = initial_docs_with_scores[:10]
            print(f"过滤后数量不足，保留前 10 个文档")
        else:
            print(f"过滤后数量: {len(filtered_docs_with_scores)}")

        initial_docs = [doc for doc, _ in filtered_docs_with_scores]
        vector_similarities = [sim for _, sim in filtered_docs_with_scores]

        # Step 3: 重排序
        reranked_docs_with_scores = rerank_documents(question, initial_docs, os.environ["SILICONFLOW_API_KEY"], top_n=10)
        reranked_docs = [doc for doc, score in reranked_docs_with_scores]
        rerank_scores = [score for _, score in reranked_docs_with_scores]

        # Step 4: 融合得分并排序
        combined_scores = [
            0.2 * vector_similarities[i] + 0.8 * rerank_scores[i]
            for i in range(len(reranked_docs))
        ]
        sorted_docs_with_scores = sorted(
            zip(reranked_docs, combined_scores),
            key=lambda x: x[1],
            reverse=True
        )
        final_docs = [doc for doc, _ in sorted_docs_with_scores][:5]

        # Step 5: 生成回答
        context = "\n\n".join([doc.page_content for doc in final_docs])
        response = qa_chain.invoke({"query": question, "context": context})
        
        return response["result"]
    except Exception as e:
        return f"Error: {str(e)}"

# 创建 Gradio 界面
interface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(label="请输入您的问题"),
    outputs=gr.Textbox(label="回答"),
    title="AI李敖助手",
    description="基于李敖163本相关书籍构建的知识库，输入问题以获取李敖风格的回答。"
)

# 启动应用
if __name__ == "__main__":
    interface.launch(share=True)
>>>>>>> 921dc7e73a28368974490d7eba946303cf2129ba