Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,20 +1,90 @@
+import os
 import streamlit as st
-from transformers import pipeline
-from PIL import Image
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
+from langchain.text_splitter 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
 
-st.title("Hot Dog? Or Not?")
+# 1. 读取环境变量（在 Hugging Face Space 中添加环境变量 OPENAI_API_KEY）
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
-file_name = st.file_uploader("Upload a hot dog candidate image")
+st.set_page_config(page_title="RAG 文档问答 Demo", layout="wide")
+st.title("📄 RAG 文档问答 Demo (LangChain + Chroma + Streamlit)")
 
-if file_name is not None:
-    col1, col2 = st.columns(2)
+st.markdown(
+    """
+    说明：
+    1. 上传一份 PDF 文件（例如论文、说明文档）。
+    2. 稍等系统构建索引。
+    3. 在下方输入问题，将基于文档内容进行回答。
+    """
+)
 
-    image = Image.open(file_name)
-    col1.image(image, use_column_width=True)
-    predictions = pipeline(image)
+# 2. 侧边栏：上传文件
+uploaded_file = st.sidebar.file_uploader("上传 PDF 文件", type=["pdf"])
 
-    col2.header("Probabilities")
-    for p in predictions:
-        col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")
+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 文件。")