""" Streamlit + Groq API - 8種 RAG 策略 PDF 問答系統 安裝依賴: pip install streamlit groq pypdf sentence-transformers numpy faiss-cpu scikit-learn 執行方式: streamlit run rag_streamlit.py """ import streamlit as st from groq import Groq import numpy as np from sentence_transformers import SentenceTransformer import faiss from pypdf import PdfReader import re from sklearn.feature_extraction.text import TfidfVectorizer import io # ───────────────────────────────────────────── # RAG 核心類別 # ───────────────────────────────────────────── class MultiStrategyRAG: def __init__(self, api_key: str): self.client = Groq(api_key=api_key) self.embedding_model = SentenceTransformer( "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" ) self.chunks: list[str] = [] self.embeddings = None self.index = None self.tfidf_vectorizer = None self.tfidf_matrix = None # ── 文件載入 ────────────────────────────── def load_pdf(self, pdf_bytes: bytes) -> str: try: reader = PdfReader(io.BytesIO(pdf_bytes)) full_text = "".join( (page.extract_text() or "") + "\n" for page in reader.pages ) self.chunks = self._split_text(full_text, chunk_size=800, overlap=150) self.embeddings = self.embedding_model.encode( self.chunks, convert_to_numpy=True ) dim = self.embeddings.shape[1] self.index = faiss.IndexFlatL2(dim) self.index.add(self.embeddings.astype("float32")) self.tfidf_vectorizer = TfidfVectorizer(max_features=1000) self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(self.chunks) return f"✅ 成功載入 PDF！共 {len(reader.pages)} 頁，分割為 {len(self.chunks)} 個片段" except Exception as e: return f"❌ 載入失敗：{e}" def _split_text(self, text: str, chunk_size: int, overlap: int) -> list[str]: chunks, start = [], 0 while start < len(text): chunk = re.sub(r"\s+", " ", text[start : start + chunk_size]).strip() if chunk: chunks.append(chunk) start += chunk_size - overlap return chunks # ── 8 種 RAG 策略 ───────────────────────── def strategy_1_basic_similarity(self, query: str, top_k: int = 3) -> list[str]: """策略1：基礎語意相似度搜尋""" vec = self.embedding_model.encode([query]).astype("float32") _, indices = self.index.search(vec, top_k) return [self.chunks[i] for i in indices[0]] def strategy_2_tfidf(self, query: str, top_k: int = 3) -> list[str]: """策略2：TF-IDF 關鍵詞搜尋""" qvec = self.tfidf_vectorizer.transform([query]) scores = (self.tfidf_matrix * qvec.T).toarray().flatten() top_idx = scores.argsort()[-top_k:][::-1] return [self.chunks[i] for i in top_idx] def strategy_3_hybrid(self, query: str, top_k: int = 3) -> list[str]: """策略3：混合搜尋（語意 + TF-IDF）""" vec = self.embedding_model.encode([query]).astype("float32") _, sem_idx = self.index.search(vec, top_k * 2) qvec = self.tfidf_vectorizer.transform([query]) tfidf_scores = (self.tfidf_matrix * qvec.T).toarray().flatten() tfidf_idx = tfidf_scores.argsort()[-top_k * 2 :][::-1] combined = list(dict.fromkeys(sem_idx[0].tolist() + tfidf_idx.tolist())) return [self.chunks[i] for i in combined[:top_k]] def strategy_4_reranking(self, query: str, top_k: int = 3) -> list[str]: """策略4：重新排序（LLM 評分重排）""" candidates = self.strategy_1_basic_similarity(query, top_k=top_k * 2) reranked = [] for chunk in candidates: prompt = ( f"問題：{query}\n\n文本：{chunk[:200]}...\n\n" "這段文本與問題的相關度(0-10)：" ) try: resp = self.client.chat.completions.create( model="llama-3.1-8b-instant", messages=[{"role": "user", "content": prompt}], max_tokens=10, temperature=0, ) raw = resp.choices[0].message.content.strip() nums = re.findall(r"\d+", raw) score = float(nums[0]) if nums else 0 except Exception: score = 0 reranked.append((chunk, score)) reranked.sort(key=lambda x: x[1], reverse=True) return [c for c, _ in reranked[:top_k]] def strategy_5_multi_query(self, query: str, top_k: int = 3) -> list[str]: """策略5：多查詢擴展""" prompt = f"將以下問題改寫成3個相關但不同角度的問題，用換行分隔：\n{query}" try: resp = self.client.chat.completions.create( model="llama-3.1-8b-instant", messages=[{"role": "user", "content": prompt}], max_tokens=200, temperature=0.7, ) queries = [query] + resp.choices[0].message.content.strip().split("\n")[:3] except Exception: queries = [query] all_chunks: list[str] = [] for q in queries: all_chunks.extend(self.strategy_1_basic_similarity(q, top_k=2)) return list(dict.fromkeys(all_chunks))[:top_k] def strategy_6_contextual_compression(self, query: str, top_k: int = 3) -> list[str]: """策略6：上下文壓縮""" chunks = self.strategy_1_basic_similarity(query, top_k=top_k) compressed = [] for chunk in chunks: prompt = ( f"從以下文本中提取與問題「{query}」最相關的1-2句話：\n\n{chunk}" ) try: resp = self.client.chat.completions.create( model="llama-3.1-8b-instant", messages=[{"role": "user", "content": prompt}], max_tokens=150, temperature=0, ) compressed.append(resp.choices[0].message.content.strip()) except Exception: compressed.append(chunk[:300]) return compressed def strategy_7_parent_child(self, query: str, top_k: int = 3) -> list[str]: """策略7：父子文檔（小片段對應大上下文）""" small_chunks = self._split_text(" ".join(self.chunks), chunk_size=300, overlap=50) small_emb = self.embedding_model.encode(small_chunks, convert_to_numpy=True) small_index = faiss.IndexFlatL2(small_emb.shape[1]) small_index.add(small_emb.astype("float32")) vec = self.embedding_model.encode([query]).astype("float32") _, indices = small_index.search(vec, top_k) results = [] for idx in indices[0]: for big in self.chunks: if small_chunks[idx] in big: results.append(big) break return list(dict.fromkeys(results))[:top_k] def strategy_8_hypothetical_answer(self, query: str, top_k: int = 3) -> list[str]: """策略8：假設性答案（HyDE）""" prompt = f"請對以下問題給出一個假設性的答案（即使不確定）：\n{query}" try: resp = self.client.chat.completions.create( model="llama-3.1-8b-instant", messages=[{"role": "user", "content": prompt}], max_tokens=200, temperature=0.7, ) hypo = resp.choices[0].message.content except Exception: hypo = query vec = self.embedding_model.encode([hypo]).astype("float32") _, indices = self.index.search(vec, top_k) return [self.chunks[i] for i in indices[0]] # ── 答案生成 ────────────────────────────── def generate_answer( self, query: str, strategy: str, top_k: int = 3 ) -> tuple[str, str]: if not self.chunks: return "❌ 請先上傳 PDF 檔案！", "" strategy_map = { "1. 基礎語意搜尋": self.strategy_1_basic_similarity, "2. TF-IDF 關鍵詞": self.strategy_2_tfidf, "3. 混合搜尋": self.strategy_3_hybrid, "4. 重新排序": self.strategy_4_reranking, "5. 多查詢擴展": self.strategy_5_multi_query, "6. 上下文壓縮": self.strategy_6_contextual_compression, "7. 父子文檔": self.strategy_7_parent_child, "8. 假設性答案 (HyDE)": self.strategy_8_hypothetical_answer, } retrieval_fn = strategy_map.get(strategy, self.strategy_1_basic_similarity) relevant_chunks = retrieval_fn(query, top_k) context = "\n\n---\n\n".join(relevant_chunks) prompt = ( "請根據以下上下文回答問題。如果上下文中沒有相關資訊，請說明無法回答。\n\n" f"上下文：\n{context}\n\n問題：{query}\n\n請用繁體中文詳細回答：" ) try: resp = self.client.chat.completions.create( model="llama-3.1-8b-instant", messages=[ {"role": "system", "content": "你是專業的文件分析助手。"}, {"role": "user", "content": prompt}, ], max_tokens=1024, temperature=0.3, ) answer = resp.choices[0].message.content source_info = ( f"📚 使用策略：{strategy}\n" f"📄 檢索片段數：{len(relevant_chunks)}\n\n" + "=" * 50 + "\n相關文本片段：\n" + "=" * 50 + f"\n\n{context}" ) return answer, source_info except Exception as e: return f"❌ 生成答案失敗：{e}", "" # ───────────────────────────────────────────── # Streamlit 頁面 # ───────────────────────────────────────────── STRATEGY_DESCRIPTIONS = { "1. 基礎語意搜尋": "使用向量餘弦相似度，將問題與文件片段對比，找出語意最接近的段落。", "2. TF-IDF 關鍵詞": "基於詞頻-逆文檔頻率（TF-IDF）統計，適合精確關鍵字匹配場景。", "3. 混合搜尋": "同時執行語意搜尋與 TF-IDF，融合兩者結果，兼顧語意與關鍵字。", "4. 重新排序": "先用語意搜尋召回候選片段，再讓 LLM 為每段打分重新排序。", "5. 多查詢擴展": "讓 LLM 將問題改寫為多個角度的問題，再分別搜尋合併結果。", "6. 上下文壓縮": "先語意搜尋，再請 LLM 從每段中萃取與問題最相關的 1-2 句。", "7. 父子文檔": "以更小的子片段搜尋，但回傳包含該子片段的原始大段文本。", "8. 假設性答案 (HyDE)": "先讓 LLM 生成一個假設答案，用此假設答案的向量搜尋文件。", } EXAMPLE_QUESTIONS = [ "這份文件的主要內容是什麼？", "文件中提到哪些重要概念？", "有哪些關鍵數據或統計資料？", "文件的結論是什麼？", ] def get_rag(api_key: str) -> MultiStrategyRAG: """在 session_state 中快取 RAG 實例（避免重複載入模型）""" if "rag" not in st.session_state: st.session_state.rag = MultiStrategyRAG(api_key=api_key) return st.session_state.rag def main(): st.set_page_config( page_title="多策略 RAG PDF 問答系統", page_icon="🤖", layout="wide", ) # ── 標題 ───────────────────────────────── st.title("🤖 多策略 RAG PDF 問答系統") st.markdown( "採用 **8 種不同的 RAG 策略**，為您的 PDF 文件提供智能問答服務！" ) st.divider() # ── API 金鑰（側邊欄） ──────────────────── with st.sidebar: st.header("⚙️ 設定") api_key = st.text_input( "Groq API Key", value="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", type="password", help="請輸入您的 Groq API 金鑰", ) st.divider() st.subheader("📖 策略說明") for name, desc in STRATEGY_DESCRIPTIONS.items(): with st.expander(name): st.write(desc) # ── 初始化 RAG ──────────────────────────── if not api_key: st.warning("⚠️ 請在側邊欄輸入 Groq API Key 後繼續。") st.stop() rag = get_rag(api_key) # ── 步驟 1：上傳 PDF ────────────────────── st.subheader("📤 步驟 1：上傳 PDF") uploaded_file = st.file_uploader("選擇 PDF 檔案", type=["pdf"]) if uploaded_file is not None: # 只在檔名改變時重新載入 if st.session_state.get("loaded_filename") != uploaded_file.name: with st.spinner("🔄 正在載入並建立索引，請稍候…"): status = rag.load_pdf(uploaded_file.read()) st.session_state["loaded_filename"] = uploaded_file.name st.session_state["load_status"] = status status_msg = st.session_state.get("load_status", "") if "✅" in status_msg: st.success(status_msg) else: st.error(status_msg) st.divider() # ── 步驟 2 & 3：策略選擇 + 提問 ────────── col_left, col_right = st.columns([1, 2]) with col_left: st.subheader("⚙️ 步驟 2：RAG 策略") strategy = st.selectbox( "選擇策略", list(STRATEGY_DESCRIPTIONS.keys()), index=0, label_visibility="collapsed", ) st.caption(STRATEGY_DESCRIPTIONS[strategy]) top_k = st.slider( "檢索片段數量（Top-K）", min_value=1, max_value=10, value=3, step=1, ) with col_right: st.subheader("💬 步驟 3：提問") # 範例問題快速填入 st.caption("💡 快速填入範例問題：") example_cols = st.columns(len(EXAMPLE_QUESTIONS)) for col, q in zip(example_cols, EXAMPLE_QUESTIONS): if col.button(q[:10] + "…", key=f"ex_{q}", use_container_width=True, help=q): st.session_state["question_input"] = q question = st.text_area( "輸入您的問題", value=st.session_state.get("question_input", ""), placeholder="例如：這份文件的主要內容是什麼？", height=100, key="question_input", ) ask_clicked = st.button( "🔍 提問", type="primary", use_container_width=True, disabled=(not rag.chunks), ) if not rag.chunks: st.info("ℹ️ 請先上傳 PDF 文件，才能開始提問。") # ── 答案輸出 ────────────────────────────── if ask_clicked: if not question.strip(): st.warning("⚠️ 請輸入問題後再送出。") else: with st.spinner("🧠 正在思考中，請稍候…"): answer, source_info = rag.generate_answer(question, strategy, top_k) st.divider() st.subheader("💡 AI 回答") if answer.startswith("❌"): st.error(answer) else: st.markdown(answer) if source_info: with st.expander("📚 查看檢索到的文本片段"): st.text(source_info) if __name__ == "__main__": main()