Update app.py

app.py

@@ -1,10 +1,10 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModel, pipeline
-import faiss
-import numpy as np
-import torch
 import os
+import torch
+import numpy as np
+import faiss
+import gradio as gr
 from PyPDF2 import PdfReader
+from transformers import AutoTokenizer, AutoModel, pipeline
 
 # ===== 嵌入模型 =====
 embed_model = AutoModel.from_pretrained(
@@ -20,26 +20,23 @@ def embed_text(text):
         embeddings = embed_model(**inputs).last_hidden_state[:, 0, :]
     return embeddings[0].numpy()
 
-# ===== 生成模型（轻量LLM） =====
+# ===== 生成模型（Qwen 1.8B） =====
 generator = pipeline(
     "text-generation",
-    model="Qwen/Qwen1.5-1.8B-Chat",  # 改成可用的公开模型
+    model="Qwen/Qwen1.5-1.8B-Chat",
     device=-1
 )
 
-# ===== 全局变量存储索引和文档 =====
+# ===== 全局变量 =====
 index = None
 docs = []
 
-# ===== 文件解析函数 =====
-
-
+# ===== 文件解析 =====
 def load_file(file_obj):
     global index, docs
     docs = []
     text_data = ""
 
-    # 获取文件路径
     file_path = file_obj.name if hasattr(file_obj, "name") else file_obj
     ext = os.path.splitext(file_path)[1].lower()
 
@@ -61,38 +58,59 @@ def load_file(file_obj):
     if not text_data.strip():
         return "未能从文件中提取到文本", None
 
-    # 切块
-    chunks = [text_data[i:i+500] for i in range(0, len(text_data), 500)]
+    # 分块（350字 + 100字重叠）
+    chunk_size = 350
+    overlap = 100
+    start = 0
+    chunks = []
+    while start < len(text_data):
+        end = min(start + chunk_size, len(text_data))
+        chunks.append(text_data[start:end])
+        start += chunk_size - overlap
+
     docs = [{"text": chunk, "source": f"chunk_{i}"} for i, chunk in enumerate(chunks)]
 
-    # 向量化并建索引
+    # 向量化 & 建索引
     doc_embeddings = np.array([embed_text(d["text"]) for d in docs])
     index = faiss.IndexFlatL2(doc_embeddings.shape[1])
     index.add(doc_embeddings)
 
     return f"已加载 {len(docs)} 个文本块", None
 
-# ===== RAG 查询函数 =====
+# ===== RAG 查询 =====
 def rag_query(query):
-    if index is None:
-        return "请先上传文件构建知识库"
+    if index is None or not docs:
+        return "请先上传文件并构建知识库"
     q_emb = embed_text(query).reshape(1, -1)
-    D, I = index.search(q_emb, k=3)
+    D, I = index.search(q_emb, k=5)  # Top-K=5
     retrieved = [docs[i]["text"] for i in I[0]]
-    context = "\n".join(retrieved)
-    prompt = f"已知信息：\n{context}\n\n问题：{query}\n请基于已知信息回答，并引用来源。"
-    result = generator(prompt, max_length=200, do_sample=False)
-    return result[0]["generated_text"]
+    context = "\n".join([f"[{idx+1}] {txt}" for idx, txt in enumerate(retrieved)])
+
+    prompt = f"""已知信息：
+{context}
+
+问题：{query}
+
+要求：
+1. 仅依据已知信息回答
+2. 无法回答时直接说“我不知道”
+3. 在回答中标注引用的片段编号
+"""
+
+    result = generator(prompt, max_length=300, do_sample=False)
+    answer = result[0]["generated_text"]
+
+    return f"回答：\n{answer}\n\n参考片段：\n{context}"
 
 # ===== Gradio 界面 =====
 with gr.Blocks() as demo:
-    gr.Markdown("## 📚 轻量 RAG 原型（上传 PDF/TXT）")
+    gr.Markdown("## 📚 加强版 RAG（Qwen 1.8B + 引用显示）")
     with gr.Row():
         file_input = gr.File(label="上传 PDF 或 TXT 文件")
         load_btn = gr.Button("构建知识库")
     status = gr.Textbox(label="状态")
     query_input = gr.Textbox(label="输入你的问题")
-    answer_output = gr.Textbox(label="回答")
+    answer_output = gr.Textbox(label="回答", lines=10)
     load_btn.click(load_file, inputs=file_input, outputs=status)
     query_input.submit(rag_query, inputs=query_input, outputs=answer_output)