Update app.py

app.py

@@ -1,131 +1,56 @@
-import logging
+import os
 import gradio as gr
-import torch
-from langchain_chroma import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFacePipeline
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import pipeline
+from langchain.llms import HuggingFacePipeline
 
-logging.basicConfig(level=logging.INFO)
+# 设置路径
+VECTOR_STORE_DIR = "./vector_store"
+MODEL_NAME = "uer/gpt2-chinese-cluecorpussmall"
 
-# 1. Load vector store
-embedding_model = HuggingFaceEmbeddings(
-    model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-)
-vector_store = Chroma(
-    persist_directory="vector_store",
-    embedding_function=embedding_model,
-)
+# 设置 LLM 和检索器
+print("🔧 加载生成模型...")
+gen_pipe = pipeline("text-generation", model=MODEL_NAME, max_new_tokens=256)
+llm = HuggingFacePipeline(pipeline=gen_pipe)
 
-# 2. Load lightweight LLM (Phi-2)
-model_id = "microsoft/phi-2"
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    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,
+print("📚 加载向量库...")
+embeddings = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 )
-llm = HuggingFacePipeline(pipeline=gen_pipe)
+vectordb = Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
 
-# 3. Build RAG QA chain
-retriever = vector_store.as_retriever(search_kwargs={"k": 3})
+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
 )
 
-# 4. Business functions
-def simple_qa(user_query: str) -> str:
-    if not user_query.strip():
-        return "⚠️ 请输入学习问题，例如：什么是定积分？"
+def qa_fn(query):
+    if not query.strip():
+        return "❌ 请输入问题内容。"
     try:
-        return qa_chain.run(user_query)
-    except Exception as e:
-        logging.exception("问答失败：%s", e)
-        return f"⚠️ 问答失败，请稍后再试。\n[调试信息] {e}"
-
-def generate_outline(topic: str):
-    if not topic.strip():
-        yield "⚠️ 请输入章节或主题，例如：高等数学 第六章 定积分", ""
-        return
-
-    yield "⌛ 正在检索/生成，请稍候…", ""
-
-    try:
-        docs = retriever.get_relevant_documents(topic)
-        if not docs:
-            yield "⚠️ 没有找到相关内容，请换个关键词试试。", ""
-            return
-
-        snippet = "\n".join(d.page_content for d in docs)
-        prompt = (
-            f"根据以下内容，为“{topic}”生成大学本科层次的结构化学习大纲，格式示例：\n"
-            f"一、章节标题\n  1. 节标题\n    （1）要点描述\n...\n\n"
-            f"文档内容：\n{snippet}\n\n学习大纲："
+        result = qa_chain({"query": query})
+        answer = result["result"]
+        sources = result.get("source_documents", [])
+        sources_text = "\n\n".join(
+            [f"【片段 {i+1}】\n" + doc.page_content for i, doc in enumerate(sources)]
         )
-        raw = gen_pipe(prompt, max_new_tokens=512)[0]["generated_text"]
-        outline = raw.split("学习大纲：")[-1].strip()
-        yield outline, snippet
+        return f"📌 回答：{answer.strip()}\n\n📚 参考：\n{sources_text}"
     except Exception as e:
-        logging.exception("大纲生成失败：%s", e)
-        yield "⚠️ 抱歉，生成失败，请稍后再试。", ""
-
-def placeholder_fn(*args, **kwargs):
-    return "功能尚未实现，请等待后续更新。"
-
-# 5. Gradio UI
-with gr.Blocks(title="智能学习助手", theme=gr.themes.Base()) as demo:
-    gr.Markdown("# 📚 智能学习助手 v2.0\n— 专业课向量问答与大纲生成 —")
-
-    with gr.Tabs():
-        # Chat tab
-        with gr.TabItem("💬 智能问答"):
-            chatbot = gr.Chatbot(show_label=False, height=400)
-            user_msg = gr.Textbox(placeholder="输入学习问题", show_label=False)
-            send_btn = gr.Button("发送", variant="primary")
-
-            def chat_flow(message, history):
-                history.append((message, "🤔 正在思考中，请稍后…"))
-                yield "", history
-                ans = simple_qa(message)
-                history[-1] = (message, ans)
-                yield "", history
-
-            send_btn.click(chat_flow, [user_msg, chatbot], [user_msg, chatbot])
-            user_msg.submit(chat_flow, [user_msg, chatbot], [user_msg, chatbot])
-
-        # Outline tab
-        with gr.TabItem("📝 生成学习大纲"):
-            topic_in = gr.Textbox(label="章节主题", placeholder="例如：定积分")
-            outline_out = gr.Textbox(label="系统生成的大纲", lines=12)
-            snippet_out = gr.Textbox(label="[调试] 检索片段", lines=6, visible=False)
-            gen_btn = gr.Button("生成大纲", variant="primary")
-            gen_btn.click(generate_outline, inputs=topic_in, outputs=[outline_out, snippet_out])
+        return f"❌ 出现错误：{str(e)}"
 
-        # Placeholder tabs
-        with gr.TabItem("❓ 自动出题"):
-            gr.Textbox(label="知识点").render()
-            gr.Dropdown(["简单", "中等", "困难"], label="难度").render()
-            gr.Slider(1, 10, step=1, label="题目数量").render()
-            gr.Button("开始出题").click(placeholder_fn, [], [])
+with gr.Blocks(title="数学知识问答助手") as demo:
+    gr.Markdown("## 📘 数学知识问答助手\n输入教材相关问题，例如：“什么是函数的定义域？”")
+    with gr.Row():
+        query_input = gr.Textbox(label="问题", placeholder="请输入你的问题", lines=2)
+        output_box = gr.Textbox(label="回答", lines=15)
+    submit_btn = gr.Button("提问")
 
-        with gr.TabItem("✅ 答案批改"):
-            gr.Textbox(label="标准答案", lines=4).render()
-            gr.Textbox(label="学生答案", lines=4).render()
-            gr.Button("开始批改").click(placeholder_fn, [], [])
+    submit_btn.click(fn=qa_fn, inputs=query_input, outputs=output_box)
 
-    gr.Markdown("---\n模型：Phi-2 + 向量库检索 | Powered by Hugging Face Spaces")
+demo.launch()
 
-if __name__ == "__main__":
-    demo.launch()