app.py

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("""
    <header>
        <h1>📄 RAG Document Q&A</h1>
        <p>Upload a document, ask questions — answers grounded in your content</p>
    </header>
    """)

    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("""
    <div style="text-align:center; color:gray; font-size:0.85rem; padding:1rem;">
    Embeddings: all-MiniLM-L6-v2 · Vector Store: FAISS · Generator: Qwen2.5-1.5B-Instruct · Chunk size: 512
    </div>
    """)

    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()