rag_deepseek / app.py
ikun520's picture
Update app.py
d12a2ae verified
from openai import OpenAI
from docx import Document
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
import os
import gradio as gr
# 配置参数
WORD_DOC_PATH = "知识库.docx"
VECTOR_INDEX_PATH = "faiss_index.index"
TEXT_DATA_PATH = "text_data.npy"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
# 初始化模型和客户端
client = OpenAI(
base_url='https://api-inference.modelscope.cn/v1/',
api_key='7ed44f86-e2c6-4b85-9c4a-26eacfc2e5ee',
)
embedder = SentenceTransformer(EMBEDDING_MODEL)
def process_word_document():
"""处理Word文档并分块"""
doc = Document(WORD_DOC_PATH)
chunks = []
current_chunk = []
for para in doc.paragraphs:
text = para.text.strip()
if text:
# 检查是否是新的条款
if text.startswith("第") and "条" in text:
# 如果当前块有内容,则保存当前块
if current_chunk:
chunks.append(" ".join(current_chunk))
current_chunk = []
# 将当前段落加入当前块
current_chunk.append(text)
# 保存最后一个块
if current_chunk:
chunks.append(" ".join(current_chunk))
return chunks
def create_vector_store():
"""创建并保存向量存储"""
if os.path.exists(VECTOR_INDEX_PATH):
return
chunks = process_word_document()
embeddings = embedder.encode(chunks, convert_to_tensor=False)
embeddings = np.array(embeddings).astype('float32')
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)
faiss.write_index(index, VECTOR_INDEX_PATH)
np.save(TEXT_DATA_PATH, np.array(chunks))
def search_knowledge(query, top_k=6):
"""知识检索"""
index = faiss.read_index(VECTOR_INDEX_PATH)
text_data = np.load(TEXT_DATA_PATH, allow_pickle=True)
query_embedding = embedder.encode([query], convert_to_tensor=False)
query_embedding = np.array(query_embedding).astype('float32')
distances, indices = index.search(query_embedding, top_k)
return "\n".join([text_data[i] for i in indices[0]])
def respond(message, history, max_tokens, temperature, top_p=3):
"""Gradio响应函数(包含深度思考过程)"""
# 检索相关知识
context = search_knowledge(message)
# 构建对话消息
messages = [
{"role": "system", "content": f"基于以下知识回答问题,如果不知道就说不知道:\n{context}"},
{"role": "user", "content": message}
]
# 初始化响应构建
full_response = f"🔍 知识库相关内容:\n{context}\n\n💡 深度思考:\n"
done_reasoning = False
yield full_response # 首次返回知识库内容
# 流式生成响应
response = client.chat.completions.create(
model='deepseek-ai/DeepSeek-R1',
messages=messages,
stream=True,
max_tokens=max_tokens,
temperature=temperature,
top_p=top_p
)
for chunk in response:
# 获取推理内容和回答内容
reasoning_chunk = getattr(chunk.choices[0].delta, "reasoning_content", "") or ""
answer_chunk = chunk.choices[0].delta.content or ""
if reasoning_chunk:
full_response += reasoning_chunk
yield full_response
if answer_chunk and not reasoning_chunk:
if not done_reasoning:
full_response += "\n\n✅ 最终答案:\n"
done_reasoning = True
full_response += answer_chunk
yield full_response
# 最终返回完整响应
yield full_response
# 初始化向量存储
create_vector_store()
# 创建Gradio界面
demo = gr.ChatInterface(
fn=respond,
additional_inputs=[
gr.Slider(512, 2048, value=2048, step=1, label="最大Token数"),
gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="温度参数"),
gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p采样"),
],
title="制度文档问答系统",
description="输入关于广西警察学院制度的问题进行问答",
)
if __name__ == "__main__":
demo.launch()