Update app.py

app.py

@@ -1,30 +1,33 @@
+# app.py
+
 import os
+import logging
 import gradio as gr
 import torch
-import logging
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-from build_index import main as build_index_if_needed  # 确保提交了 build_index.py
+from build_index import main as build_index_if_needed  # 需确保 build_index.py 在同目录
 
 logging.basicConfig(level=logging.INFO)
 
 # ─── 配置 ─────────────────────────────────────────────────────
-VECTOR_STORE_DIR = "./vector_store"
-MODEL_NAME = "uer/gpt2-chinese-cluecorpussmall"
-EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
+VECTOR_STORE_DIR        = "./vector_store"
+MODEL_NAME              = "uer/gpt2-chinese-cluecorpussmall"
+EMBEDDING_MODEL_NAME    = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 
-# 如果向量库不存在，自动构建
+# 容器启动时自动构建向量库（如果还没提交 vector_store）
 if not os.path.exists(VECTOR_STORE_DIR) or not os.listdir(VECTOR_STORE_DIR):
     logging.info("向量库不存在，启动自动构建……")
     build_index_if_needed()
 
-# ─── 1. 加载 LLM ────────────────────────────────────────────────
+# ─── 1. 加载生成模型 ──────────────────────────────────────────────
 logging.info("🔧 加载生成模型…")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
@@ -48,8 +51,8 @@ logging.info("✅ 生成模型加载成功。")
 # ─── 2. 加载向量库 ─────────────────────────────────────────────
 logging.info("📚 加载向量库…")
 embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
-vectordb = Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
-retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+vectordb   = Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
+retriever  = vectordb.as_retriever(search_kwargs={"k": 3})
 logging.info("✅ 向量库加载成功。")
 
 # ─── 3. 自定义 Prompt ─────────────────────────────────────────
@@ -66,25 +69,34 @@ prompt_template = PromptTemplate.from_template(
 """
 )
 
-# ─── 4. 构建 RAG 问答链 ───────────────────────────────────────
+# ─── 4. 构建 RAG 问答链（map_reduce） ───────────────────────────
 qa_chain = RetrievalQA.from_chain_type(
     llm=llm,
-    chain_type="stuff",
+    chain_type="map_reduce",                             # map_reduce 避免超长
     retriever=retriever,
     chain_type_kwargs={"prompt": prompt_template},
     return_source_documents=True,
 )
-logging.info("✅ RAG 问答链构建成功。")
+logging.info("✅ RAG 问答链（map_reduce）构建成功。")
 
 # ─── 5. 业务函数 ───────────────────────────────────────────────
 def qa_fn(query: str):
-    if not query.strip():
+    if not query or not query.strip():
         return "❌ 请输入问题内容。"
-    result = qa_chain({"query": query})
-    answer = result["result"].strip()
+    try:
+        result = qa_chain({"query": query})
+    except Exception as e:
+        logging.error(f"问答链运行出错：{e}")
+        return "抱歉，问答过程中出现错误，请稍后重试。"
+
+    answer  = result.get("result", "").strip()
     sources = result.get("source_documents", [])
+    if not answer:
+        return "📌 回答：未生成答案，请稍后再试。"
     if not sources:
-        return "📌 回答：未在知识库中找到相关内容，请尝试更换问题或补充教材。"
+        return f"📌 回答：{answer}\n\n（未检索到参考片段）"
+
+    # 拼接参考片段
     sources_text = "\n\n".join(
         [f"【片段 {i+1}】\n{doc.page_content}" for i, doc in enumerate(sources)]
     )
@@ -99,7 +111,7 @@ with gr.Blocks(title="智能学习助手") as demo:
     gr.Button("提问").click(fn=qa_fn, inputs=query, outputs=answer)
     gr.Markdown(
         "---\n"
-        "模型：UER/GPT2-Chinese-ClueCorpus + Sentence-Transformers RAG  \n"
+        "模型：UER/GPT2-Chinese-ClueCorpus + Sentence-Transformers RAG (map_reduce)  \n"
         "由 Hugging Face Spaces 提供算力支持"
     )
 
@@ -109,3 +121,4 @@ if __name__ == "__main__":
 
 
 
+