Update app.py

app.py

@@ -10,6 +10,8 @@ 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
+
 logging.basicConfig(level=logging.INFO)
 
 # ─── 配置 ─────────────────────────────────────────────────────
@@ -17,8 +19,13 @@ VECTOR_STORE_DIR = "./vector_store"
 MODEL_NAME = "uer/gpt2-chinese-cluecorpussmall"
 EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 
+# 如果向量库不存在，自动构建
+if not os.path.exists(VECTOR_STORE_DIR) or not os.listdir(VECTOR_STORE_DIR):
+    logging.info("向量库不存在，启动自动构建……")
+    build_index_if_needed()
+
 # ─── 1. 加载 LLM ────────────────────────────────────────────────
-print("🔧 加载生成模型…")
+logging.info("🔧 加载生成模型…")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
@@ -33,22 +40,22 @@ gen_pipe = pipeline(
     temperature=0.5,
     top_p=0.9,
     do_sample=True,
+    trust_remote_code=True,
 )
 llm = HuggingFacePipeline(pipeline=gen_pipe)
-print("✅ 生成模型加载成功。")
+logging.info("✅ 生成模型加载成功。")
 
 # ─── 2. 加载向量库 ─────────────────────────────────────────────
-print("📚 加载向量库…")
+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})
-print("✅ 向量库加载成功。")
+logging.info("✅ 向量库加载成功。")
 
 # ─── 3. 自定义 Prompt ─────────────────────────────────────────
 prompt_template = PromptTemplate.from_template(
     """你是一位专业的数学助教，请根据以下参考资料回答用户的问题。
 如果资料中没有相关内容，请直接回答“我不知道”或“资料中未提及”，不要编造答案。
-
 参考资料：
 {context}
 
@@ -67,19 +74,17 @@ qa_chain = RetrievalQA.from_chain_type(
     chain_type_kwargs={"prompt": prompt_template},
     return_source_documents=True,
 )
-print("✅ RAG 问答链构建成功。")
+logging.info("✅ RAG 问答链构建成功。")
 
 # ─── 5. 业务函数 ───────────────────────────────────────────────
 def qa_fn(query: str):
     if not query.strip():
         return "❌ 请输入问题内容。"
-    # 执行检索与问答
     result = qa_chain({"query": query})
     answer = result["result"].strip()
     sources = result.get("source_documents", [])
     if not sources:
         return "📌 回答：未在知识库中找到相关内容，请尝试更换问题或补充教材。"
-    # 拼接参考片段
     sources_text = "\n\n".join(
         [f"【片段 {i+1}】\n{doc.page_content}" for i, doc in enumerate(sources)]
     )
@@ -103,3 +108,4 @@ if __name__ == "__main__":
 
 
 
+