import gradio as gr import torch import numpy as np from sentence_transformers import SentenceTransformer import faiss from transformers import AutoModelForCausalLM, AutoTokenizer from pathlib import Path # ── Config ── EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2" GEN_MODEL = "Qwen/Qwen2.5-1.5B-Instruct" CHUNK_SIZE = 512 CHUNK_OVERLAP = 0.15 TOP_K = 5 theme = ( gr.themes.Soft(primary_hue="indigo", neutral_hue="slate") .set(button_primary_background_fill_hover="#4f46e5") ) # ── Load models ── print("Loading embedding model...") embedder = SentenceTransformer(EMBED_MODEL) EMBED_DIM = embedder.get_sentence_embedding_dimension() print("Loading generation model (Qwen2.5-1.5B)...") gen_tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL) gen_model = AutoModelForCausalLM.from_pretrained( GEN_MODEL, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto", ) # ── Document store (in-memory) ── documents = [] # list of {"id", "title", "text"} all_chunks = [] # list of chunk strings chunk_to_doc = [] # chunk index → document index index = None # FAISS index def chunk_document(text, chunk_size, overlap=0.15): stride = int(chunk_size * (1 - overlap)) chunks = [] start = 0 while start < len(text): end = min(start + chunk_size, len(text)) chunk = text[start:end] if len(chunk.strip()) > 20: chunks.append(chunk) if end == len(text): break start += stride return chunks def rebuild_index(): global all_chunks, chunk_to_doc, index all_chunks = [] chunk_to_doc = [] for doc_idx, doc in enumerate(documents): chunks = chunk_document(doc["text"], CHUNK_SIZE, CHUNK_OVERLAP) for c in chunks: all_chunks.append(c) chunk_to_doc.append(doc_idx) if not all_chunks: index = None return 0 embeddings = embedder.encode(all_chunks, show_progress_bar=False) index = faiss.IndexFlatL2(EMBED_DIM) index.add(embeddings.astype(np.float32)) return len(all_chunks) def add_pdf_text(text, title): doc_id = f"doc_{len(documents)}" documents.append({"id": doc_id, "title": title, "text": text}) n = rebuild_index() return n def search(query, top_k=TOP_K): if index is None or index.ntotal == 0: return [] q_emb = embedder.encode([query])[0] D, I = index.search(q_emb.reshape(1, -1).astype(np.float32), min(top_k, index.ntotal)) results = [] for rank, (dist, idx) in enumerate(zip(D[0], I[0])): doc_idx = chunk_to_doc[idx] results.append({ "rank": rank + 1, "chunk": all_chunks[idx][:500], "distance": float(dist), "similarity": float(1 / (1 + dist)), "source": documents[doc_idx]["title"][:80], }) return results def generate_answer(question, history): if index is None or index.ntotal == 0: return "Please upload a document first." # Retrieve results = search(question) if not results: return "No relevant documents found. Try uploading a different document." context_text = "\n\n".join([ f"[Excerpt {r['rank']}]: {r['chunk']}" for r in results ]) prompt = ( "You are a document analysis assistant. Answer the question using ONLY the provided document excerpts.\n" "If the excerpts don't contain enough information, say so clearly.\n\n" f"Document Excerpts:\n{context_text}\n\n" f"Question: {question}\n\nAnswer:" ) inputs = gen_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048) inputs = {k: v.to(gen_model.device) for k, v in inputs.items()} with torch.no_grad(): outputs = gen_model.generate( **inputs, max_new_tokens=256, temperature=0.3, do_sample=True, pad_token_id=gen_tokenizer.eos_token_id, ) answer = gen_tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True) return answer.strip() def upload_and_ask(file, question, history): if file is None: return question, history, "Please upload a PDF or text file first.", "" try: path = Path(file.name) text = path.read_text(encoding="utf-8", errors="replace") if len(text) < 50: text = "" import subprocess result = subprocess.run(["pdftotext", str(path), "-"], capture_output=True, text=True) text = result.stdout if result.stdout else "Could not extract text from PDF." except Exception as e: # Try pdftotext try: import subprocess result = subprocess.run(["pdftotext", str(file.name), "-"], capture_output=True, text=True) text = result.stdout if result.stdout else f"Error extracting text: {e}" except Exception: text = f"Could not process file: {e}" title = Path(file.name).stem n_chunks = add_pdf_text(text, title) status = f"✅ Indexed {n_chunks:,} chunks from '{title}' ({len(text):,} chars)" if question.strip(): answer = generate_answer(question, history) history.append({"role": "user", "content": question}) history.append({"role": "assistant", "content": answer}) else: answer = "" history.append({"role": "assistant", "content": f"Document '{title}' loaded and indexed ({n_chunks:,} chunks). Ask me anything about it!"}) return question, history, status, "" def chat_fn(message, history): if index is None or index.ntotal == 0: return "Please upload a document first using the file uploader above." answer = generate_answer(message, history) return answer # ── UI ── with gr.Blocks(theme=theme, title="RAG Document Q&A", css=""" .gradio-container { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important; } header { text-align: center; padding: 1.5rem; background: linear-gradient(135deg, #6366f1, #4f46e5); color: white; border-radius: 12px; margin-bottom: 1rem; } header h1 { margin: 0; font-size: 1.75rem; } header p { margin: 0.25rem 0 0; opacity: 0.9; font-size: 0.9rem; } .status-ok { color: #16a34a; font-weight: 500; } """) as demo: gr.HTML("""

📄 RAG Document Q&A

Upload a document, ask questions — answers grounded in your content

""") with gr.Row(): with gr.Column(scale=2): file_input = gr.File(label="Upload Document (PDF or TXT)", file_types=[".pdf", ".txt"]) with gr.Column(scale=1): upload_btn = gr.Button("Upload & Index", variant="primary") status_display = gr.Markdown("📎 Upload a document to get started.") gr.Markdown("---") chatbot = gr.ChatInterface( fn=chat_fn, title="Ask Questions About Your Document", description="Your questions are answered using chunks from the uploaded document.", examples=[ "What is this document about?", "Summarize the key points", "What methodology was used?", "What are the main conclusions?", ], ) gr.Markdown("""
Embeddings: all-MiniLM-L6-v2 · Vector Store: FAISS · Generator: Qwen2.5-1.5B-Instruct · Chunk size: 512
""") def handle_upload(file): if file is None: return "📎 Upload a document to get started." try: path = Path(file.name) text = path.read_text(encoding="utf-8", errors="replace") if len(text) < 50: raise ValueError("Short text") except Exception: try: import subprocess result = subprocess.run(["pdftotext", str(file.name), "-"], capture_output=True, text=True) text = result.stdout if result.stdout else "Could not extract text from PDF." except Exception as e: return f"❌ Error: {e}" title = Path(file.name).stem n_chunks = add_pdf_text(text, title) return f"✅ Indexed **{n_chunks:,}** chunks from **'{title}'** ({len(text):,} chars). Start asking questions!" upload_btn.click(handle_upload, inputs=[file_input], outputs=[status_display]) file_input.change(handle_upload, inputs=[file_input], outputs=[status_display]) if __name__ == "__main__": demo.launch()