Spaces:
Running
Running
| import os | |
| import time | |
| import pandas as pd | |
| import gradio as gr | |
| import config as cfg | |
| from cache import LRUCache | |
| from data_loader import load_documents, validate_document | |
| from metadata_extractor import extract_metadata | |
| from chunker import chunk_documents | |
| from embedder import Embedder | |
| from duplicate_detector import DuplicateDetector | |
| from vector_store import VectorStore | |
| from bm25_retriever import BM25Retriever | |
| from dense_retriever import DenseRetriever | |
| from hybrid_retriever import HybridRetriever | |
| from reranker import Reranker | |
| from context_manager import ContextManager | |
| from llm_handler import LLMHandler | |
| from evaluator import Evaluator | |
| from query_rewriter import QueryRewriter | |
| class RagPipeline: | |
| def __init__(self): | |
| self.chunks = [] | |
| self.embedder = None | |
| self.vector_store = None | |
| self.bm25 = None | |
| self.dense = None | |
| self.hybrid = None | |
| self.reranker = None | |
| self.context_manager = None | |
| self.llm = None | |
| self.evaluator = None | |
| self.rewriter = None | |
| self.index_built = False | |
| self.answer_cache = LRUCache(max_size=cfg.CACHE_MAX_SIZE, ttl_seconds=cfg.CACHE_TTL_SECONDS) | |
| self.retrieval_cache = LRUCache(max_size=cfg.CACHE_MAX_SIZE, ttl_seconds=cfg.CACHE_TTL_SECONDS) | |
| self.embedding_cache = LRUCache(max_size=cfg.CACHE_MAX_SIZE, ttl_seconds=cfg.CACHE_TTL_SECONDS) | |
| def build_index(self, progress=None): | |
| if progress: | |
| progress(0, desc="Loading documents...") | |
| raw = load_documents(cfg.DATA_PATH) | |
| if not validate_document(raw): | |
| raise ValueError("Invalid document format") | |
| if progress: | |
| progress(0.15, desc="Extracting metadata...") | |
| entries = extract_metadata(raw) | |
| if progress: | |
| progress(0.30, desc="Chunking documents...") | |
| self.chunks = chunk_documents( | |
| entries, | |
| chunk_size=cfg.CHUNK_SIZE, | |
| overlap=cfg.CHUNK_OVERLAP, | |
| ) | |
| if progress: | |
| progress(0.40, desc="Removing duplicates...") | |
| dedup = DuplicateDetector( | |
| threshold=cfg.DUP_SIM_THRESHOLD, | |
| num_perm=cfg.DUP_NUM_PERM, | |
| ) | |
| self.chunks = dedup.filter_duplicates(self.chunks) | |
| if progress: | |
| progress(0.50, desc="Loading embedding model...") | |
| self.embedder = Embedder( | |
| model_name=cfg.EMBEDDING_MODEL, | |
| device=cfg.EMBEDDING_DEVICE, | |
| ) | |
| if progress: | |
| progress(0.60, desc="Embedding chunks...") | |
| for c in self.chunks: | |
| c["search_text"] = " | ".join(p for p in (c.get("title", ""), c.get("summary", ""), c["text"]) if p) | |
| texts = [c["search_text"] for c in self.chunks] | |
| embeddings = self.embedder.embed(texts) | |
| if progress: | |
| progress(0.70, desc="Indexing vector store...") | |
| self.vector_store = VectorStore( | |
| persist_dir=cfg.CHROMA_DB_PATH, | |
| embedding_dim=cfg.EMBEDDING_DIM, | |
| ) | |
| self.vector_store.create_collection( | |
| overwrite=True, | |
| ef_construction=cfg.HNSW_EF_CONSTRUCTION, | |
| m=cfg.HNSW_M, | |
| ef_search=cfg.HNSW_EF_SEARCH, | |
| ) | |
| ids = [c["chunk_id"] for c in self.chunks] | |
| metadatas = [] | |
| for c in self.chunks: | |
| meta = {} | |
| for k, v in c.items(): | |
| if k in ("text", "chunk_id"): | |
| continue | |
| if isinstance(v, list): | |
| meta[k] = " | ".join(str(x) for x in v if x) | |
| else: | |
| meta[k] = str(v) | |
| metadatas.append(meta) | |
| self.vector_store.add(ids, embeddings, texts, metadatas) | |
| if progress: | |
| progress(0.85, desc="Building BM25 index...") | |
| self.bm25 = BM25Retriever() | |
| self.bm25.fit(self.chunks) | |
| self.dense = DenseRetriever(self.vector_store) | |
| self.hybrid = HybridRetriever(self.bm25, self.dense, alpha=cfg.HYBRID_ALPHA, fusion_method=cfg.HYBRID_FUSION_METHOD) | |
| if progress: | |
| progress(0.90, desc="Loading reranker...") | |
| self.reranker = Reranker( | |
| model_name=cfg.RERANKER_MODEL, | |
| device=cfg.RERANKER_DEVICE, | |
| ) | |
| self.context_manager = ContextManager(max_tokens=cfg.MAX_CONTEXT_TOKENS) | |
| if cfg.GROQ_API_KEY: | |
| if progress: | |
| progress(0.95, desc="Initializing LLM...") | |
| self.llm = LLMHandler( | |
| api_key=cfg.GROQ_API_KEY, | |
| model=cfg.GROQ_MODEL, | |
| temperature=cfg.GROQ_TEMPERATURE, | |
| max_tokens=cfg.GROQ_MAX_TOKENS, | |
| timeout=cfg.GROQ_TIMEOUT, | |
| ) | |
| self.rewriter = QueryRewriter(llm=self.llm) | |
| self.evaluator = Evaluator( | |
| self.hybrid, self.reranker, self.llm, self.context_manager, self.embedder | |
| ) | |
| self.answer_cache.clear() | |
| self.retrieval_cache.clear() | |
| self.embedding_cache.clear() | |
| self.index_built = True | |
| def _retrieve(self, search_query: str, alpha: float, top_k: int): | |
| cache_key = (search_query, alpha) | |
| cached = self.retrieval_cache.get(cache_key) | |
| if cached: | |
| return cached | |
| cached_emb = self.embedding_cache.get(search_query) | |
| if cached_emb is not None: | |
| query_emb = cached_emb | |
| else: | |
| query_emb = self.embedder.embed_query(search_query) | |
| self.embedding_cache.put(search_query, query_emb) | |
| self.hybrid.alpha = alpha | |
| results = self.hybrid.search(search_query, query_emb, n_results=cfg.HYBRID_TOP_K) | |
| reranked = self.reranker.rerank(search_query, results[:cfg.RERANK_CANDIDATES], top_k=top_k) | |
| self.retrieval_cache.put(cache_key, reranked) | |
| return reranked | |
| def _build_sources_df(self, reranked): | |
| sources = [] | |
| for i, r in enumerate(reranked): | |
| sources.append({ | |
| "Rank": i + 1, | |
| "Score": round(r.get("rerank_score", r.get("score", 0)), 4), | |
| "Title": r["metadata"].get("title", ""), | |
| "Doc ID": r["id"], | |
| "Preview": r["text"][:200] + ("..." if len(r["text"]) > 200 else ""), | |
| }) | |
| return pd.DataFrame(sources) | |
| def query(self, question: str, alpha: float, top_k: int, use_llm: bool, use_reformulation: bool = False): | |
| if not self.index_built: | |
| return "Index not built yet. Click 'Build Index' first.", pd.DataFrame() | |
| cache_key = (question, alpha, top_k, use_llm, use_reformulation) | |
| cached = self.answer_cache.get(cache_key) | |
| if cached: | |
| return cached["answer"], cached["sources"] | |
| start = time.perf_counter() | |
| search_query = question | |
| if use_reformulation and self.rewriter: | |
| search_query = self.rewriter.rewrite(question) | |
| reranked = self._retrieve(search_query, alpha, top_k) | |
| retrieval_time = (time.perf_counter() - start) * 1000 | |
| if use_llm and self.llm: | |
| prompt = self.context_manager.assemble_prompt(question, reranked) | |
| answer = self.llm.generate(prompt) | |
| gen_time = (time.perf_counter() - start) * 1000 - retrieval_time | |
| answer_text = answer or "(LLM returned no response)" | |
| else: | |
| answer_text = "(LLM disabled — retrieval only)" | |
| gen_time = 0 | |
| total_time = (time.perf_counter() - start) * 1000 | |
| df = self._build_sources_df(reranked) | |
| result = f"{answer_text}\n\n---\n*Retrieval: {retrieval_time:.1f}ms | Generation: {gen_time:.1f}ms | Total: {total_time:.1f}ms*" | |
| self.answer_cache.put(cache_key, {"answer": result, "sources": df}) | |
| return result, df | |
| def query_stream(self, question: str, alpha: float, top_k: int, use_llm: bool, use_reformulation: bool = False): | |
| if not question.strip(): | |
| yield "Please enter a question.", pd.DataFrame() | |
| return | |
| if not self.index_built: | |
| yield "Index not built yet. Click 'Build Index' first.", pd.DataFrame() | |
| return | |
| cache_key = (question, alpha, top_k, use_llm, use_reformulation) | |
| cached = self.answer_cache.get(cache_key) | |
| if cached: | |
| yield cached["answer"], cached["sources"] | |
| return | |
| start = time.perf_counter() | |
| search_query = question | |
| if use_reformulation and self.rewriter: | |
| search_query = self.rewriter.rewrite(question) | |
| reranked = self._retrieve(search_query, alpha, top_k) | |
| retrieval_time = (time.perf_counter() - start) * 1000 | |
| df = self._build_sources_df(reranked) | |
| if use_llm and self.llm: | |
| yield "**Retrieval complete. Generating answer...**", df | |
| prompt = self.context_manager.assemble_prompt(question, reranked) | |
| full_answer = "" | |
| for chunk in self.llm.generate_stream(prompt): | |
| if chunk: | |
| full_answer += chunk | |
| yield f"{full_answer}\n\n---\n*Generating...*", df | |
| gen_time = (time.perf_counter() - start) * 1000 - retrieval_time | |
| total_time = (time.perf_counter() - start) * 1000 | |
| answer_text = full_answer or "(LLM returned no response)" | |
| else: | |
| answer_text = "(LLM disabled — retrieval only)" | |
| gen_time = 0 | |
| total_time = (time.perf_counter() - start) * 1000 | |
| final = f"{answer_text}\n\n---\n*Retrieval: {retrieval_time:.1f}ms | Generation: {gen_time:.1f}ms | Total: {total_time:.1f}ms*" | |
| self.answer_cache.put(cache_key, {"answer": final, "sources": df}) | |
| yield final, df | |
| def run_evaluation(self, num_questions: int): | |
| if not self.index_built: | |
| return "Index not built yet.", pd.DataFrame() | |
| results = self.evaluator.full_evaluation(self.chunks, num_questions) | |
| rows = [] | |
| for k, v in results["retrieval"].items(): | |
| val_str = f"{v:.4f}" if isinstance(v, float) else str(v) | |
| rows.append({"Metric": k, "Value": val_str}) | |
| rows.append({ | |
| "Metric": "Num Eval Questions", | |
| "Value": results["num_questions"], | |
| }) | |
| gen = results["generation"] | |
| rows.append({ | |
| "Metric": "Avg End-to-End Latency (ms)", | |
| "Value": gen.get("Avg End-to-End Latency (ms)", "N/A"), | |
| }) | |
| for m in ("Avg Faithfulness", "Avg Answer Accuracy"): | |
| if m in gen: | |
| v = gen[m] | |
| rows.append({"Metric": m, "Value": f"{v:.4f}" if isinstance(v, float) else str(v)}) | |
| df = pd.DataFrame(rows) | |
| details = "" | |
| for a in gen.get("answers", []): | |
| faith = a.get("faithfulness", "N/A") | |
| acc = a.get("accuracy", "N/A") | |
| faith_str = f"{faith:.3f}" if isinstance(faith, float) else str(faith) | |
| acc_str = f"{acc:.3f}" if isinstance(acc, float) else str(acc) | |
| details += f"**Q:** {a['question']}\n**A:** {a['answer'][:300]}...\n**Latency:** {a.get('latency_ms', 'N/A')}ms | **Faithfulness:** {faith_str} | **Accuracy:** {acc_str}\n\n---\n\n" | |
| return details, df | |
| def get_chunk_preview(self, max_rows: int = 50): | |
| if not self.chunks: | |
| return pd.DataFrame() | |
| data = [] | |
| for c in self.chunks[:max_rows]: | |
| data.append({ | |
| "Chunk ID": c.get("chunk_id", ""), | |
| "Title": c.get("title", ""), | |
| "Level": c.get("level", ""), | |
| "Tokens": c.get("num_tokens", 0), | |
| "Preview": c["text"][:150] + ("..." if len(c["text"]) > 150 else ""), | |
| }) | |
| return pd.DataFrame(data) | |
| pipeline = RagPipeline() | |
| def build_index_fn(progress=gr.Progress()): | |
| if pipeline.index_built: | |
| return "Index already built. Reload the page to rebuild." | |
| try: | |
| pipeline.build_index(progress=progress) | |
| count = pipeline.vector_store.count() | |
| return f"Index built successfully — {count} chunks indexed." | |
| except Exception as e: | |
| return f"Error: {e}" | |
| _stream_state = {} | |
| def retrieve_fn(question, alpha, top_k, use_llm, use_reformulation): | |
| if not question.strip(): | |
| return "Please enter a question.", pd.DataFrame() | |
| if not pipeline.index_built: | |
| return "Index not built yet. Click 'Build Index' first.", pd.DataFrame() | |
| cache_key = (question, alpha, top_k, use_llm, use_reformulation) | |
| cached = pipeline.answer_cache.get(cache_key) | |
| if cached: | |
| return cached["answer"], cached["sources"] | |
| search_query = question | |
| if use_reformulation and pipeline.rewriter: | |
| search_query = pipeline.rewriter.rewrite(question) | |
| retrieve_start = time.perf_counter() | |
| reranked = pipeline._retrieve(search_query, alpha, top_k) | |
| retrieval_time_ms = (time.perf_counter() - retrieve_start) * 1000 | |
| df = pipeline._build_sources_df(reranked) | |
| _stream_state["query"] = question | |
| _stream_state["search_query"] = search_query | |
| _stream_state["reranked"] = reranked | |
| _stream_state["alpha"] = alpha | |
| _stream_state["top_k"] = top_k | |
| _stream_state["use_llm"] = use_llm | |
| _stream_state["use_reformulation"] = use_reformulation | |
| _stream_state["retrieval_time_ms"] = retrieval_time_ms | |
| _stream_state["gen_start"] = time.perf_counter() | |
| if use_llm and pipeline.llm: | |
| return "**Retrieval complete. Generating answer...**", df | |
| return "(LLM disabled — retrieval only)", df | |
| def stream_fn(): | |
| if "query" not in _stream_state or "gen_start" not in _stream_state: | |
| return | |
| use_llm = _stream_state.get("use_llm", False) | |
| if not use_llm or not pipeline.llm: | |
| return | |
| question = _stream_state.pop("query") | |
| reranked = _stream_state.pop("reranked") | |
| gen_start = _stream_state.pop("gen_start") | |
| retrieval_time_ms = _stream_state.pop("retrieval_time_ms") | |
| prompt = pipeline.context_manager.assemble_prompt(question, reranked) | |
| full_answer = "" | |
| for chunk in pipeline.llm.generate_stream(prompt): | |
| if chunk: | |
| full_answer += chunk | |
| yield f"{full_answer}\n\n---\n*Generating...*" | |
| gen_time = (time.perf_counter() - gen_start) * 1000 | |
| total_time = retrieval_time_ms + gen_time | |
| answer_text = full_answer or "(LLM returned no response)" | |
| final = f"{answer_text}\n\n---\n*Retrieval: {retrieval_time_ms:.1f}ms | Generation: {gen_time:.1f}ms | Total: {total_time:.1f}ms*" | |
| cache_key = (question, _stream_state["alpha"], _stream_state["top_k"], | |
| use_llm, _stream_state["use_reformulation"]) | |
| pipeline.answer_cache.put(cache_key, {"answer": final, "sources": pipeline._build_sources_df(reranked)}) | |
| yield final | |
| def eval_fn(num_q): | |
| return pipeline.run_evaluation(int(num_q)) | |
| def browse_fn(max_rows): | |
| return pipeline.get_chunk_preview(int(max_rows)) | |
| with gr.Blocks( | |
| title="Production RAG Pipeline", | |
| ) as demo: | |
| gr.Markdown( | |
| "# Production RAG Pipeline\n" | |
| "Hybrid retrieval (BM25 + Dense) → Reranking → LLM generation via Groq. " | |
| f"Embedding: `{cfg.EMBEDDING_MODEL}` | Reranker: `{cfg.RERANKER_MODEL}`" | |
| ) | |
| with gr.Tab("Ask"): | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| query_input = gr.Textbox( | |
| label="Your Question", | |
| placeholder="e.g., What is the DeepSeek-V4 architecture?", | |
| lines=3, | |
| ) | |
| with gr.Row(): | |
| submit_btn = gr.Button("Ask", variant="primary") | |
| clear_btn = gr.Button("Clear") | |
| with gr.Column(scale=1): | |
| alpha_slider = gr.Slider( | |
| minimum=0.0, | |
| maximum=1.0, | |
| value=cfg.HYBRID_ALPHA, | |
| step=0.1, | |
| label="Hybrid α (0=BM25, 1=Dense)", | |
| ) | |
| top_k_dropdown = gr.Dropdown( | |
| choices=[3, 5, 10, 15, 20], | |
| value=5, | |
| label="Final Top-K Results", | |
| ) | |
| use_llm_checkbox = gr.Checkbox( | |
| label="Use LLM (Groq)", | |
| value=bool(cfg.GROQ_API_KEY), | |
| interactive=bool(cfg.GROQ_API_KEY), | |
| ) | |
| reformulate_checkbox = gr.Checkbox( | |
| label="Reformulate query (better retrieval)", | |
| value=False, | |
| interactive=True, | |
| ) | |
| answer_output = gr.Markdown(label="Answer", value="*Ask a question to get started.*") | |
| sources_output = gr.Dataframe(label="Retrieved Sources", wrap=True) | |
| click_event = submit_btn.click( | |
| fn=retrieve_fn, | |
| inputs=[query_input, alpha_slider, top_k_dropdown, use_llm_checkbox, reformulate_checkbox], | |
| outputs=[answer_output, sources_output], | |
| ) | |
| click_event.then( | |
| fn=stream_fn, | |
| outputs=[answer_output], | |
| ) | |
| clear_btn.click( | |
| fn=lambda: ("", pd.DataFrame()), | |
| outputs=[query_input, sources_output], | |
| ) | |
| with gr.Tab("Evaluate"): | |
| gr.Markdown("Run evaluation to measure retrieval accuracy and end-to-end latency.") | |
| with gr.Row(): | |
| num_q_slider = gr.Slider( | |
| minimum=5, maximum=50, value=cfg.EVAL_NUM_QUESTIONS, step=5, label="Number of Questions" | |
| ) | |
| eval_btn = gr.Button("Run Evaluation", variant="primary") | |
| eval_details = gr.Markdown(label="Sample Answers") | |
| eval_metrics = gr.Dataframe(label="Evaluation Metrics") | |
| eval_btn.click( | |
| fn=eval_fn, | |
| inputs=[num_q_slider], | |
| outputs=[eval_details, eval_metrics], | |
| ) | |
| with gr.Tab("Browse"): | |
| gr.Markdown("Browse the indexed document chunks.") | |
| with gr.Row(): | |
| max_rows_slider = gr.Slider( | |
| minimum=10, maximum=200, value=50, step=10, label="Max Rows" | |
| ) | |
| browse_btn = gr.Button("Refresh", variant="secondary") | |
| browse_output = gr.Dataframe(label="Chunks", wrap=True) | |
| browse_btn.click(fn=browse_fn, inputs=[max_rows_slider], outputs=[browse_output]) | |
| with gr.Tab("Configuration"): | |
| config_rows = [ | |
| ["Embedding Model", cfg.EMBEDDING_MODEL], | |
| ["Reranker Model", cfg.RERANKER_MODEL], | |
| ["Chunk Size", str(cfg.CHUNK_SIZE)], | |
| ["Chunk Overlap", str(cfg.CHUNK_OVERLAP)], | |
| ["HNSW ef_construction", str(cfg.HNSW_EF_CONSTRUCTION)], | |
| ["HNSW M", str(cfg.HNSW_M)], | |
| ["HNSW ef_search", str(cfg.HNSW_EF_SEARCH)], | |
| ["BM25 Top-K", str(cfg.BM25_TOP_K)], | |
| ["Dense Top-K", str(cfg.DENSE_TOP_K)], | |
| ["Hybrid Top-K", str(cfg.HYBRID_TOP_K)], | |
| ["Rerank Top-K", str(cfg.RERANK_TOP_K)], | |
| ["Final Top-K", str(cfg.FINAL_TOP_K)], | |
| ["Max Context Tokens", str(cfg.MAX_CONTEXT_TOKENS)], | |
| ["Dup Threshold", str(cfg.DUP_SIM_THRESHOLD)], | |
| ["Groq Model", cfg.GROQ_MODEL], | |
| ["Groq API Key Set", "Yes" if cfg.GROQ_API_KEY else "No"], | |
| ["Data File", cfg.DATA_PATH], | |
| ["Chroma DB Path", cfg.CHROMA_DB_PATH], | |
| ] | |
| config_df = pd.DataFrame(config_rows, columns=["Parameter", "Value"]) | |
| gr.Dataframe(value=config_df, label="Current Configuration") | |
| with gr.Tab("Build"): | |
| gr.Markdown( | |
| "Build or rebuild the index. This loads the document, chunks, embeds, " | |
| "and builds both the vector and BM25 indexes." | |
| ) | |
| build_btn = gr.Button("Build Index", variant="primary") | |
| build_output = gr.Markdown("*Index not built yet.*") | |
| build_btn.click(fn=build_index_fn, outputs=[build_output]) | |
| demo.queue() | |
| demo.launch( | |
| server_port=cfg.GRADIO_PORT, | |
| share=cfg.GRADIO_SHARE, | |
| theme=gr.themes.Soft(), | |
| css="footer {display:none !important}", | |
| ) | |