File size: 3,024 Bytes
d933810 459405c 27e44bf d933810 459405c d933810 459405c d933810 459405c d933810 459405c d933810 459405c d933810 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import os
import streamlit as st
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains import RetrievalQA
# 1. 读取环境变量（在 Hugging Face Space 中添加环境变量 OPENAI_API_KEY）
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
st.set_page_config(page_title="RAG 文档问答 Demo", layout="wide")
st.title("📄 RAG 文档问答 Demo (LangChain + Chroma + Streamlit)")
st.markdown(
"""
说明：
1. 上传一份 PDF 文件（例如论文、说明文档）。
2. 稍等系统构建索引。
3. 在下方输入问题，将基于文档内容进行回答。
"""
)
# 2. 侧边栏：上传文件
uploaded_file = st.sidebar.file_uploader("上传 PDF 文件", type=["pdf"])
if not OPENAI_API_KEY:
st.error("请在环境变量中设置 OPENAI_API_KEY 才能调用 OpenAI 接口。")
st.stop()
if uploaded_file:
# 把上传的文件临时保存到本地（Space 的临时存储）
temp_pdf_path = "temp.pdf"
with open(temp_pdf_path, "wb") as f:
f.write(uploaded_file.getbuffer())
# 3. 加载 PDF 文档
loader = PyPDFLoader(temp_pdf_path)
pages = loader.load()
# 4. 切分文本
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
length_function=len,
)
documents = text_splitter.split_documents(pages)
st.sidebar.write(f"文档页数: {len(pages)}")
st.sidebar.write(f"切分后的文本块数: {len(documents)}")
# 5. 构建向量库（Chroma）
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
vectorstore = Chroma.from_documents(documents, embedding=embeddings)
# 6. 构建 RAG QA 链
llm = ChatOpenAI(
temperature=0.1,
model="gpt-4o-mini", # 或 gpt-4o / gpt-3.5-turbo 等
openai_api_key=OPENAI_API_KEY
)
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
chain_type="stuff", # 简单拼接检索到的文本块
return_source_documents=True
)
# 7. 用户提问
user_question = st.text_input("在此输入关于文档的问题：", "")
if st.button("生成回答") and user_question.strip():
with st.spinner("正在检索并生成回答，请稍候..."):
result = qa_chain(user_question)
answer = result["result"]
source_docs = result["source_documents"]
st.subheader("回答：")
st.write(answer)
with st.expander("查看参考片段（检索到的文档内容）"):
for i, doc in enumerate(source_docs):
st.markdown(f"**片段 {i+1}:**")
st.write(doc.page_content)
st.markdown("---")
else:
st.info("请先在左侧上传一个 PDF 文件。") |