app.py

import os
import gradio as gr
import torch
import logging

# LangChain 0.1.x 系列的导入方式
from langchain_chroma import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFacePipeline
from langchain.chains import RetrievalQA

# Transformers 库
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

logging.basicConfig(level=logging.INFO)

# 设置路径
# 直接使用 Hugging Face Hub 上的模型 ID，而不是本地缓存路径
VECTOR_STORE_DIR = "./vector_store" # 这个目录用于 ChromaDB，我们保留
MODEL_NAME = "uer/gpt2-chinese-cluecorpussmall"  # <--- 修改这里！
EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2" # <--- 修改这里！

# 1. 轻量 LLM（uer/gpt2-chinese-cluecorpussmall）
print("🔧 加载生成模型...")
try:
    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
    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
    print("❌ 向量库加载失败，RAG功能将无法使用。")

# 3. RAG 问答链
qa_chain = None
if llm and vectordb:
    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:
        return "⚠️ 问答系统未完全加载，请稍后再试或检查日志。"
    try:
        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)]
        )
        return f"📌 回答：{answer.strip()}\n\n📚 参考：\n{sources_text}"
    except Exception as e:
        logging.exception("问答失败：%s", e)
        return f"❌ 出现错误：{str(e)}\n请检查日志获取更多信息。"

# 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)
        output_box = gr.Textbox(label="回答", lines=15)
    submit_btn = gr.Button("提问")

    submit_btn.click(fn=qa_fn, inputs=query_input, outputs=output_box)

    gr.Markdown("---\n模型：uer/gpt2-chinese-cluecorpussmall + Chroma RAG | Powered by Hugging Face Spaces")

if __name__ == "__main__":
    demo.launch()