Update app.py

app.py

@@ -1,37 +1,94 @@
 import os
 import gradio as gr
-from langchain.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceEmbeddings
+import torch # 确保导入 torch，因为 phi-2 模型需要
+import logging
+
+# LangChain 新版导入方式
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+# 注意：HuggingFacePipeline 从 0.2.x 开始推荐从 langchain_huggingface 导入
+# 如果遇到问题，也可以尝试从 langchain_community.llms 导入
+from langchain_huggingface.llms import HuggingFacePipeline # 或者 from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import RetrievalQA
-from transformers import pipeline
-from langchain.llms import HuggingFacePipeline
+
+# Transformers 库
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+logging.basicConfig(level=logging.INFO) # 更好的日志级别
 
 # 设置路径
+# 确保这些路径与您的预下载模型和向量库文件夹名称匹配
 VECTOR_STORE_DIR = "./vector_store"
-MODEL_NAME = "uer/gpt2-chinese-cluecorpussmall"
+# 将模型名称指向本地预下载的路径
+MODEL_NAME = "./hf_models_cache/models--uer--gpt2-chinese-cluecorpussmall"
+EMBEDDING_MODEL_NAME = "./hf_models_cache/models--sentence-transformers--paraphrase-multilingual-mpnet-base-v2"
+
 
-# 设置 LLM 和检索器
+# 1. 轻量 LLM（uer/gpt2-chinese-cluecorpussmall）
 print("🔧 加载生成模型...")
-gen_pipe = pipeline("text-generation", model=MODEL_NAME, max_new_tokens=256)
-llm = HuggingFacePipeline(pipeline=gen_pipe)
-
-print("📚 加载向量库...")
-embeddings = HuggingFaceEmbeddings(
-    model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-)
-vectordb = Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
-
-retriever = vectordb.as_retriever(search_kwargs={"k": 3})
-qa_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    chain_type="stuff",
-    retriever=retriever,
-    return_source_documents=True
-)
+try:
+    # 确保 tokenizer 和 model 是从正确的本地路径加载
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto",
+    )
+    gen_pipe = pipeline(
+        task="text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=256,
+        temperature=0.5,
+        top_p=0.9,
+        do_sample=True,
+    )
+    llm = HuggingFacePipeline(pipeline=gen_pipe)
+    print("✅ 生成模型加载成功。")
+except Exception as e:
+    logging.error(f"加载生成模型失败: {e}", exc_info=True)
+    # 如果加载失败，可以考虑在这里退出或者给一个友好的错误提示
+    llm = None # 确保 llm 为 None，避免后续报错
+    print("❌ 生成模型加载失败，应用可能无法正常工作。")
+
+
+# 2. 向量库和嵌入模型
+print("📚 加载向量库和嵌入模型...")
+try:
+    embeddings = HuggingFaceEmbeddings(
+        model_name=EMBEDDING_MODEL_NAME # 指向本地预下载的嵌入模型路径
+    )
+    vectordb = Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
+    print("✅ 向量库加载成功。")
+except Exception as e:
+    logging.error(f"加载向量库失败: {e}", exc_info=True)
+    vectordb = None # 确保 vectordb 为 None
+    print("❌ 向量库加载失败，RAG功能将无法使用。")
+
 
+# 3. RAG 问答链
+qa_chain = None
+if llm and vectordb: # 只有当LLM和向量库都成功加载时才构建RAG链
+    try:
+        retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=retriever,
+            return_source_documents=True
+        )
+        print("✅ RAG问答链构建成功。")
+    except Exception as e:
+        logging.error(f"构建RAG问答链失败: {e}", exc_info=True)
+        print("❌ RAG问答链构建失败。")
+
+
+# 4. 业务函数
 def qa_fn(query):
     if not query.strip():
         return "❌ 请输入问题内容。"
+    if not qa_chain: # 检查 qa_chain 是否已成功构建
+        return "⚠️ 问答系统未完全加载，请稍后再试或检查日志。"
     try:
         result = qa_chain({"query": query})
         answer = result["result"]
@@ -41,9 +98,11 @@ def qa_fn(query):
         )
         return f"📌 回答：{answer.strip()}\n\n📚 参考：\n{sources_text}"
     except Exception as e:
-        return f"❌ 出现错误：{str(e)}"
+        logging.exception("问答失败：%s", e)
+        return f"❌ 出现错误：{str(e)}\n请检查日志获取更多信息。"
 
-with gr.Blocks(title="数学知识问答助手") as demo:
+# 5. Gradio UI
+with gr.Blocks(title="数学知识问答助手", theme=gr.themes.Base()) as demo:
     gr.Markdown("## 📘 数学知识问答助手\n输入教材相关问题，例如：“什么是函数的定义域？”")
     with gr.Row():
         query_input = gr.Textbox(label="问题", placeholder="请输入你的问题", lines=2)
@@ -52,5 +111,9 @@ with gr.Blocks(title="数学知识问答助手") as demo:
 
     submit_btn.click(fn=qa_fn, inputs=query_input, outputs=output_box)
 
-demo.launch()
+    gr.Markdown("---\n模型：uer/gpt2-chinese-cluecorpussmall + Chroma RAG | Powered by Hugging Face Spaces")
+
+
+if __name__ == "__main__":
+    demo.launch()