import streamlit as st import torch import io import re import numpy as np import pandas as pd from threading import Thread # LLM from transformers.models.auto.tokenization_auto import AutoTokenizer from transformers.models.auto.modeling_auto import AutoModelForCausalLM from transformers.generation.streamers import TextIteratorStreamer # Embedding + Reranker from langchain_huggingface import HuggingFaceEmbeddings from sentence_transformers import CrossEncoder # Document processing from langchain_community.vectorstores import FAISS from pypdf import PdfReader # BM25 (keyword search) from rank_bm25 import BM25Okapi # ========================================== # ⚡ CPU CONFIG — free HuggingFace Spaces # ========================================== # ⚡ CPU CONFIG — call this before any parallel work try: torch.set_num_threads(2) torch.set_num_interop_threads(1) except RuntimeError: pass # Already initialized by Streamlit at startup; ignore. # Layout st.set_page_config(page_title="Hybrid RAG + Reranker", page_icon="🔍", layout="wide") st.title("🔍 Hybrid RAG + Reranker") # Sidebar text with st.sidebar.expander("⚙️ How does this system work?", expanded=False): st.markdown(""" **Step 1 - Dual retrieval (hybrid)** When you ask a question, two search engines work in parallel over your chunks: - **BM25**: classic keyword search (like Google in the 1990s). Strong on names, dates, and exact terms. Very fast, no AI model. - **FAISS**: semantic vector search. Understands meaning even when the words differ. Each returns its own **top 5 chunks** -> we merge them -> **10 candidate chunks**. --- **Step 2 - Reranking (cross-encoder)** A small model (~70M params) reads each *(question, chunk)* pair together and assigns a true relevance score. Unlike embeddings that encode the question and chunk **separately**, the cross-encoder reads them **together** -> it actually understands the relationship. Result: we keep the **2 best chunks** out of the 10. --- **Step 3 - Generation** The LLM receives only these 2 highly relevant chunks and generates the answer. One LLM call, clean context = precise and fast response. """) # Load the models once (embedding + reranker + LLM + quantization) @st.cache_resource(show_spinner="Loading models... (first run ~2 min)") def load_models(): # ⚡ Embedding: MiniLM 120M - fast and sufficient embeddings = HuggingFaceEmbeddings( model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", model_kwargs={'device': 'cpu'}, encode_kwargs={'normalize_embeddings': True, 'batch_size': 16} ) # ⚡ Reranker: lightweight cross-encoder, ~70M params # Reads (question + chunk) together -> real relevance score # Much more accurate than simple cosine similarity reranker = CrossEncoder( "cross-encoder/ms-marco-MiniLM-L-6-v2", max_length=512, device="cpu" ) # LLM: Qwen 1.5B Instruct model_id = "Qwen/Qwen2.5-0.5B-Instruct" # instead of 1.5B tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True) llm = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=torch.bfloat16, # instead of float32 low_cpu_mem_usage=True ) # Warm-up: JIT-compile the kernels and avoids lag on the first request dummy = tokenizer("warmup", return_tensors="pt") with torch.inference_mode(): llm.generate(**dummy, max_new_tokens=1, pad_token_id=tokenizer.eos_token_id) return embeddings, reranker, tokenizer, llm embeddings, reranker, tokenizer, llm = load_models() # Qwen stop tokens EOS_IDS = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>")] MAX_TOKENS_ANSWER = 384 # ========================================== # SESSION STATE # ========================================== for key, default in [ ("vector_db", None), ("bm25", None), ("chunks_text", None), ("current_file_name", None) ]: if key not in st.session_state: st.session_state[key] = default # PDF upload section (allow multiple files) uploaded_files = st.sidebar.file_uploader("📂 Upload your PDFs (multiple allowed)", type=["pdf"], accept_multiple_files=True) if not uploaded_files: for key in ["vector_db", "bm25", "chunks_text", "current_file_names", "chunks_source"]: st.session_state[key] = None else: # Collect filenames filenames = [f.name for f in uploaded_files] # Only re-index if uploaded set changed if st.session_state.get("current_file_names") != filenames: for key in ["vector_db", "bm25", "chunks_text", "chunks_source"]: st.session_state[key] = None st.session_state.current_file_names = filenames if st.session_state.vector_db is None: with st.spinner("Indexing documents... ⏳"): all_chunks = [] all_sources = [] # Extract and chunk each uploaded PDF for f in uploaded_files: try: reader = PdfReader(io.BytesIO(f.getvalue())) full_text = "\n".join(page.extract_text() or "" for page in reader.pages) normalized_text = re.sub(r"\s+", " ", full_text).strip() def split_text_into_chunks(text: str, chunk_size: int = 500, chunk_overlap: int = 200) -> list[str]: if not text: return [] chunks = [] start = 0 text_length = len(text) while start < text_length: end = min(start + chunk_size, text_length) chunks.append(text[start:end].strip()) if end >= text_length: break start = max(end - chunk_overlap, 0) return [chunk for chunk in chunks if chunk] chunks = split_text_into_chunks(normalized_text) all_chunks.extend(chunks) all_sources.extend([f.name] * len(chunks)) except Exception as e: st.sidebar.warning(f"Could not process {f.name}: {e}") # --- FAISS index (semantic search) --- st.session_state.vector_db = FAISS.from_texts(all_chunks, embeddings) # --- BM25 index (keyword search) --- tokenized = [t.lower().split() for t in all_chunks] st.session_state.bm25 = BM25Okapi(tokenized) st.session_state.chunks_text = all_chunks # raw text for BM25 st.session_state.chunks_source = all_sources st.sidebar.success(f"✅ Indexed {len(all_chunks)} chunks from {len(filenames)} files") # Sidebar: scope selector (All files or specific uploaded file) scope_options = ["All files"] if st.session_state.get("current_file_names"): scope_options += st.session_state.current_file_names selected_scope = st.sidebar.selectbox("Index scope", scope_options, index=0) selected_file = None if selected_scope == "All files" else selected_scope # Fonctions def hybrid_search(question: str, k: int = 10, source_filter: str | None = None) -> list[str]: """ Hybrid search: BM25 + FAISS -> merge the results. Why k=5 for each? We retrieve broadly on purpose (10 candidates total) because the reranker will sort and keep the true best matches. BM25 catches what FAISS misses (exact terms) and vice versa. """ question_lower = question.lower() # -- BM25: term-frequency score -- bm25_scores = st.session_state.bm25.get_scores(question_lower.split()) top_bm25_idx = np.argsort(bm25_scores)[::-1][:k] bm25_results = [] for i in top_bm25_idx: text = st.session_state.chunks_text[i] if source_filter is None or st.session_state.chunks_source[i] == source_filter: bm25_results.append(text) # -- FAISS: semantic similarity score -- faiss_docs = st.session_state.vector_db.similarity_search(question, k=k) faiss_results = [] for d in faiss_docs: # d may be a Document-like object text = getattr(d, 'page_content', str(d)) # try to find index to check source try: idx = st.session_state.chunks_text.index(text) source = st.session_state.chunks_source[idx] except Exception: source = None if source_filter is None or source == source_filter: faiss_results.append(text) # -- Merge with deduplication -- seen = set() merged = [] for text in bm25_results + faiss_results: if text not in seen: seen.add(text) merged.append(text) return merged # up to 10 unique chunks def rerank(question: str, candidates: list[str], top_n: int = 2): pairs = [(question, chunk) for chunk in candidates] scores = reranker.predict(pairs) ranked = sorted(zip(scores, candidates), reverse=True) best_score = ranked[0][0] best_chunks = [chunk for _, chunk in ranked[:top_n]] return best_chunks, float(best_score) def build_prompt(system: str, user: str) -> str: return ( f"<|im_start|>system\n{system}\n<|im_end|>\n" f"<|im_start|>user\n{user}\n<|im_end|>\n" f"<|im_start|>assistant\n" ) def stream_llm(prompt: str, placeholder) -> str: """Streamed generation with an animated cursor.""" inputs = tokenizer(prompt, return_tensors="pt") streamer = TextIteratorStreamer( tokenizer, skip_prompt=True, skip_special_tokens=True ) def _generate(): with torch.inference_mode(): llm.generate( **inputs, streamer=streamer, max_new_tokens=MAX_TOKENS_ANSWER, do_sample=False, repetition_penalty=1.3, pad_token_id=tokenizer.eos_token_id, eos_token_id=EOS_IDS ) Thread(target=_generate).start() full_text = "" for token in streamer: full_text += token placeholder.markdown(full_text + "▌") placeholder.markdown(full_text) return full_text def clean_chunk(text: str) -> str: """Mark truncated chunks instead of sending them raw.""" text = text.strip() # If the chunk starts lowercase, it is likely truncated. if text and text[0].islower(): text = "(...) " + text return text def classify_question_style(question: str) -> str: question_lower = question.lower() value_markers = [ "value", "values", "calculate", "calculation", "compute", "how many", "how much", "what is the amount", "report the number", "report the value", "percentage", "percent", "%", "ratio", "emissions", "ghg", "co2", "scope 1", "scope 2", "scope 3", "tonnes", "tons" ] structured_markers = [ "list", "provide", "state", "report", "identify", "name", "names", "describe", "what are", "which", "give the", "enumerate", "include" ] if any(marker in question_lower for marker in value_markers): return "value_only" if any(marker in question_lower for marker in structured_markers): return "structured" return "paragraph" def build_answer_system_prompt(answer_style: str) -> str: if answer_style == "value_only": return ( "You are a precise extraction engine. " "Return only the exact final value, number, percentage, or unit requested by the question. " "Do not add explanation, labels, or narrative text. " "If the answer requires a calculation, compute it from the context and return the final value only. " "If the answer is not present, say: Information not found." ) if answer_style == "structured": return ( "You are a structured extractor. " "Return a concise answer using short sentences or bullet points, preserving all requested facts. " "Do not compress multiple requested fields into one vague paragraph. " "If the question asks for several items, answer each item explicitly. " "Do not omit numbers, names, dates, locations, or amounts that are relevant. " "If the answer is not present, say: Information not found." ) return ( "You are a fact extractor. " "Use the context to answer fully and faithfully, without compressing away important facts. " "If the answer contains multiple facts, include them all in a clear paragraph. " "Do not omit numbers, names, dates, locations, or amounts that are relevant. " "Never generate URLs, emails, menus, or language lists. " "If the context contains '(...)', it is an incomplete sentence, so reconstruct the meaning. " "If you cannot find the answer, say: Information not found." ) def build_answer_instruction(answer_style: str) -> str: if answer_style == "value_only": return ( "Answer format: output only the final value, number, percentage, or unit. " "Do not explain your reasoning." ) if answer_style == "structured": return ( "Answer format: use short structured points or a compact list, one item per requested field." ) return ( "Answer format: provide a full, faithful paragraph with all relevant facts." ) def generate_answer(question: str, context: str, batch_mode: bool = False) -> str: answer_style = classify_question_style(question) prompt = build_prompt( system=build_answer_system_prompt(answer_style), user=( f"CONTEXT:\n{context}\n\n" f"QUESTION: {question}\n\n" f"{build_answer_instruction(answer_style)}" ) ) inputs = tokenizer(prompt, return_tensors="pt") with torch.inference_mode(): output = llm.generate( **inputs, max_new_tokens=MAX_TOKENS_ANSWER if batch_mode else MAX_TOKENS_ANSWER, do_sample=False, repetition_penalty=1.05 if answer_style == "value_only" else 1.15, pad_token_id=tokenizer.eos_token_id, eos_token_id=EOS_IDS ) generated_tokens = output[0][inputs["input_ids"].shape[-1]:] answer_text = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip() return answer_text # ========================================== # MAIN INTERFACE # ========================================== if st.session_state.vector_db is not None: question = st.text_input( "Ask your question:", placeholder="Type a question" ) if st.button("Run analysis", type="primary") and question: st.markdown("---") # --- STEP 1: HYBRID RETRIEVAL --- with st.spinner("🔎 Hybrid retrieval (BM25 + vector search)..."): candidates = hybrid_search(question, k=5, source_filter=selected_file) # Optional debug display with st.expander(f"📚 {len(candidates)} candidate chunks retrieved (before reranking)"): for i, c in enumerate(candidates): st.caption(f"**Chunk {i+1}**: {c[:150]}...") # --- STEP 2: RERANKING --- with st.spinner("⚖️ Reranking (selecting the 2 best chunks)..."): best_chunks, confidence = rerank(question, candidates, top_n=2) st.caption(f"Relevance score: `{confidence:.2f}`") if confidence < -2.0: # very negative score = nothing relevant st.warning("❌ No relevant passage found for this question.") st.stop() with st.expander("✅ 2 chunks kept after reranking"): for i, c in enumerate(best_chunks): st.info(f"**Chunk {i+1}**: {c}") # --- STEP 3: LLM GENERATION --- context = "\n\n".join([clean_chunk(c) for c in best_chunks]) st.markdown("### 🎯 Answer:") placeholder = st.empty() placeholder.markdown(generate_answer(question, context)) # ========================================== # INDICATORS & BATCH ANALYSIS # ========================================== st.markdown("---") st.markdown("## 📋 Batch Analysis: All Indicators") @st.cache_resource(show_spinner=False) def load_indicators(): """Parse the indicators and questions from the file.""" indicators = [] try: with open("all_indicators_with_questions - Copy.txt", "r", encoding="utf-8") as f: lines = f.readlines() i = 0 while i < len(lines): line = lines[i].strip() # Match indicator lines (e.g., "1. IMP-M01-I01 | ...") if line and line[0].isdigit() and "|" in line: parts = line.split("|") if len(parts) >= 3: id_part = parts[0].strip().split(".")[-1].strip() category = parts[1].strip() indicator_name = parts[2].strip() # Next line should be the question i += 1 if i < len(lines) and lines[i].strip().startswith("Question:"): question_text = lines[i].strip()[10:].strip() indicators.append({ "id": id_part, "category": category, "indicator": indicator_name, "question": question_text, "answer": "" }) i += 1 except Exception as e: st.error(f"Error loading indicators: {str(e)}") return indicators indicators = load_indicators() if indicators: col1, col2 = st.columns([3, 1]) with col1: st.write(f"**Found {len(indicators)} indicators**") with col2: run_batch_analysis = st.button("🚀 Analyze All", type="primary", use_container_width=True) if run_batch_analysis: # Create a progress bar and results storage progress_bar = st.progress(0) status_text = st.empty() results_placeholder = st.empty() results = [] for idx, ind in enumerate(indicators): # Update progress progress = (idx) / len(indicators) progress_bar.progress(progress) status_text.write(f"Processing {idx + 1}/{len(indicators)}: **{ind['id']}** - {ind['indicator'][:50]}...") # Run the analysis for this indicator question = ind["question"] try: # Hybrid search candidates = hybrid_search(question, k=5, source_filter=selected_file) if candidates: # Reranking best_chunks, confidence = rerank(question, candidates, top_n=2) # Generate answer if confidence >= -2.0: context = "\n\n".join([clean_chunk(c) for c in best_chunks]) answer_text = generate_answer(question, context, batch_mode=True) results.append({ **ind, "answer": answer_text, "confidence": f"{confidence:.2f}" }) else: results.append({**ind, "answer": "No relevant passage found.", "confidence": f"{confidence:.2f}"}) else: results.append({**ind, "answer": "No relevant passages retrieved.", "confidence": "N/A"}) except Exception as e: results.append({**ind, "answer": f"Error: {str(e)[:50]}", "confidence": "N/A"}) # Complete progress progress_bar.progress(1.0) status_text.write(f"✅ Completed: {len(results)} indicators analyzed") # Display results in a table if results: st.subheader("📊 Analysis Results") # Create a DataFrame for display df_results = pd.DataFrame([ { "ID": r["id"], "Category": r["category"], "Indicator": r["indicator"], "Question": r["question"], "Answer": r["answer"], "Score": r["confidence"] } for r in results ]) df_display = df_results.copy() df_display["Indicator"] = df_display["Indicator"].apply(lambda value: value[:40] + "..." if len(value) > 40 else value) df_display["Question"] = df_display["Question"].apply(lambda value: value[:60] + "..." if len(value) > 60 else value) df_display["Answer"] = df_display["Answer"].apply(lambda value: value[:120] + "..." if len(value) > 120 else value) st.dataframe(df_display, use_container_width=True, height=500) # Option to export results csv_output = df_results.to_csv(index=False) st.download_button( label="📥 Download Results as CSV", data=csv_output, file_name="indicators_analysis_results.csv", mime="text/csv" ) else: st.warning("⚠️ Could not load indicators. Make sure 'all_indicators_with_questions - Copy.txt' is in the current directory.") else: st.info("⬅️ Upload a PDF file in the sidebar to begin.")