Update app.py

app.py

@@ -1,9 +1,7 @@
-# app.py - Hybrid RAG (FAISS semantic + BM25, fused with RRF) for multi-format documents.
-# Handles: PDFs (text + tables), DOCX (with tables), PPTX, XLSX/CSV, TXT/MD. Arabic + English.
+# app.py - Simple hybrid RAG (semantic + BM25, fused with RRF). PDF, DOCX, XLSX/CSV, TXT. Arabic + English.
 import io
 import os
 import re
-import unicodedata
 from collections import defaultdict
 
 import faiss
@@ -13,352 +11,136 @@ import pandas as pd
 import streamlit as st
 from docx import Document
 from langchain_groq import ChatGroq
-from pptx import Presentation
 from rank_bm25 import BM25Okapi
 from sentence_transformers import SentenceTransformer
 
-# ========================= CONSTANTS =========================
-EMBED_MODEL_NAME = "paraphrase-multilingual-MiniLM-L12-v2"  # Arabic + English, 384-dim
-LLM_MODEL_NAME = "llama-3.3-70b-versatile"
-CHUNK_WORDS = 80          # stays under the embedding model's 128-token limit
-CHUNK_OVERLAP = 15
-TABLE_CHUNK_WORDS = 60    # table/spreadsheet rows are token-heavy, so smaller groups
-RRF_K = 60
+EMBED_MODEL = "paraphrase-multilingual-MiniLM-L12-v2"
+LLM_MODEL = "llama-3.3-70b-versatile"
+CHUNK_WORDS, OVERLAP, RRF_K = 80, 15, 60
 
-# ========================= TEXT NORMALIZATION =========================
-DIACRITICS_RE = re.compile(r"[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06ED]")
-DIGIT_MAP = str.maketrans("٠١٢٣٤٥٦٧٨٩۰۱۲۳۴۵۶۷۸۹", "01234567890123456789")
+st.set_page_config(page_title="RAG Assistant", layout="wide")
+st.title("RAG Assistant")
+st.caption("Semantic search + BM25, fused with Reciprocal Rank Fusion")
 
-def light_clean(text):
-    text = unicodedata.normalize("NFKC", text)
-    text = DIACRITICS_RE.sub("", text)
-    return text.replace("\u0640", "")
-
-def normalize(text):
-    text = light_clean(text).translate(DIGIT_MAP)
-    text = re.sub("[إأآٱ]", "ا", text)
-    text = text.replace("ى", "ي").replace("ة", "ه")
-    return text.lower()
 
+# ---------- text helpers ----------
 def tokenize(text):
-    tokens = re.findall(r"[\w\+#]+", normalize(text))
-    return [t[2:] if len(t) > 4 and t.startswith("ال") else t for t in tokens]
-
-# ========================= EXTRACTION =========================
-def make_block(text, source, loc, kind):
-    return {"text": text, "source": source, "loc": loc, "kind": kind}
-
-def pdf_tables(page):
-    out = []
-    try:
-        for t in page.find_tables().tables:
-            lines = []
-            for r in t.extract():
-                cells = [(c or "").replace("\n", " ").strip() for c in r]
-                if any(cells):
-                    lines.append(" | ".join(cells))
-            if len(lines) >= 2:
-                out.append("\n".join(lines))
-    except Exception:
-        pass
-    return out
-
-def extract_pdf(name, data):
-    blocks = []
-    stats = {"pages": 0, "tables": 0, "notes": []}
-    doc = fitz.open(stream=data, filetype="pdf")
-    empty_pages = []
-    for i, page in enumerate(doc, 1):
-        stats["pages"] += 1
-        loc = f"p.{i}"
-        text = unicodedata.normalize("NFKC", page.get_text("text") or "").strip()
-        for tbl in pdf_tables(page):
-            blocks.append(make_block(tbl, name, loc, "table"))
-            stats["tables"] += 1
-        if len(text) < 50:
-            empty_pages.append(i)
-        if text:
-            blocks.append(make_block(text, name, loc, "text"))
-    if empty_pages:
-        shown = ", ".join(map(str, empty_pages[:15])) + ("..." if len(empty_pages) > 15 else "")
-        stats["notes"].append(f"Little or no text on page(s) {shown}. Scanned pages cannot be read by this app.")
-    return blocks, stats
-
-def extract_docx(name, data):
-    doc = Document(io.BytesIO(data))
-    blocks, tables = [], 0
-    text = "\n".join(p.text for p in doc.paragraphs if p.text.strip())
-    if text.strip():
-        blocks.append(make_block(text, name, None, "text"))
-    for table in doc.tables:
-        lines = []
-        for row in table.rows:
-            cells = [c.text.replace("\n", " ").strip() for c in row.cells]
-            if any(cells):
-                lines.append(" | ".join(cells))
-        if len(lines) >= 2:
-            blocks.append(make_block("\n".join(lines), name, None, "table"))
-            tables += 1
-    return blocks, {"pages": 1, "tables": tables,
-                    "notes": ["Images embedded inside the DOCX are not read."]}
-
-def extract_pptx(name, data):
-    prs = Presentation(io.BytesIO(data))
-    blocks, tables = [], 0
-    for i, slide in enumerate(prs.slides, 1):
-        loc = f"slide {i}"
-        texts = []
-        for shape in slide.shapes:
-            if shape.has_text_frame and shape.text_frame.text.strip():
-                texts.append(shape.text_frame.text.strip())
-            if getattr(shape, "has_table", False) and shape.has_table:
-                lines = [" | ".join(c.text.strip() for c in row.cells) for row in shape.table.rows]
-                if len(lines) >= 2:
-                    blocks.append(make_block("\n".join(lines), name, loc, "table"))
-                    tables += 1
-        if slide.has_notes_slide and slide.notes_slide.notes_text_frame.text.strip():
-            texts.append("Notes: " + slide.notes_slide.notes_text_frame.text.strip())
-        if texts:
-            blocks.append(make_block("\n".join(texts), name, loc, "text"))
-    return blocks, {"pages": len(prs.slides), "tables": tables,
-                    "notes": ["Pictures and charts on slides are not read."]}
+    text = re.sub(r"[\u064B-\u065F\u0640]", "", text)           # diacritics, tatweel
+    text = re.sub("[إأآ]", "ا", text).replace("ى", "ي").replace("ة", "ه")
+    return re.findall(r"\w+", text.lower())
 
-def df_to_rows(df):
-    df = df.dropna(how="all").dropna(axis=1, how="all").fillna("")
-    headers = [str(c) for c in df.columns]
-    lines = []
-    for _, row in df.iterrows():
-        cells = []
-        for h, v in zip(headers, row):
-            v = str(v).strip()
-            if v:
-                cells.append(v if h.startswith("Unnamed") else f"{h}: {v}")
-        if cells:
-            lines.append(" | ".join(cells))
-    return "\n".join(lines)
 
-def extract_spreadsheet(name, data):
-    blocks = []
-    if name.lower().endswith(".csv"):
-        df = None
-        for enc in ("utf-8-sig", "cp1256", "latin-1"):
-            try:
-                df = pd.read_csv(io.BytesIO(data), encoding=enc)
-                break
-            except Exception:
-                continue
-        sheets = {"": df} if df is not None else {}
-    else:
-        sheets = pd.read_excel(io.BytesIO(data), sheet_name=None)
-    for sheet, df in sheets.items():
-        text = df_to_rows(df)
-        if text:
-            blocks.append(make_block(text, name, f"sheet {sheet}" if sheet else None, "rows"))
-    return blocks, {"pages": len(sheets), "tables": len(blocks), "notes": []}
-
-def extract_text_file(name, data):
-    text = data.decode("utf-8", errors="ignore")
-    blocks = [make_block(text, name, None, "text")] if text.strip() else []
-    return blocks, {"pages": 1, "tables": 0, "notes": []}
-
-@st.cache_data(show_spinner=False)
-def extract_file(name, data):
+def read_file(name, data):
+    """Return a list of (label, text) pieces."""
     lower = name.lower()
     if lower.endswith(".pdf"):
-        return extract_pdf(name, data)
+        doc = fitz.open(stream=data, filetype="pdf")
+        return [(f"{name}, p.{i}", p.get_text()) for i, p in enumerate(doc, 1)]
     if lower.endswith(".docx"):
-        return extract_docx(name, data)
-    if lower.endswith(".pptx"):
-        return extract_pptx(name, data)
+        d = Document(io.BytesIO(data))
+        text = "\n".join(p.text for p in d.paragraphs)
+        for t in d.tables:
+            text += "\n" + "\n".join(" | ".join(c.text for c in r.cells) for r in t.rows)
+        return [(name, text)]
     if lower.endswith((".xlsx", ".csv")):
-        return extract_spreadsheet(name, data)
-    if lower.endswith((".txt", ".md")):
-        return extract_text_file(name, data)
-    return [], {"pages": 0, "tables": 0, "notes": ["Unsupported file type."]}
-
-# ========================= CHUNKING =========================
-def make_chunk(block, text):
-    return {"text": text, "source": block["source"], "loc": block["loc"], "kind": block["kind"]}
-
-def chunk_blocks(blocks):
-    chunks = []
-    step = max(1, CHUNK_WORDS - CHUNK_OVERLAP)
-    for b in blocks:
-        if b["kind"] in ("table", "rows"):
-            lines = [l for l in b["text"].split("\n") if l.strip()]
-            header = lines[0] if b["kind"] == "table" else None
-            body = lines[1:] if b["kind"] == "table" else lines
-            group, wc = [], 0
-            for line in body:
-                lw = len(line.split())
-                if group and wc + lw > TABLE_CHUNK_WORDS:
-                    chunks.append(make_chunk(b, "\n".join(([header] if header else []) + group)))
-                    group, wc = [], 0
-                group.append(line)
-                wc += lw
-            if group:
-                chunks.append(make_chunk(b, "\n".join(([header] if header else []) + group)))
-        else:
-            words = b["text"].split()
-            for i in range(0, len(words), step):
-                piece = " ".join(words[i:i + CHUNK_WORDS])
-                if len(piece.strip()) > 20:
-                    chunks.append(make_chunk(b, piece))
+        sheets = ({"": pd.read_csv(io.BytesIO(data), encoding="utf-8-sig")} if lower.endswith(".csv")
+                  else pd.read_excel(io.BytesIO(data), sheet_name=None))
+        pieces = []
+        for sheet, df in sheets.items():
+            rows = [" | ".join(f"{c}: {v}" for c, v in r.items() if pd.notna(v))
+                    for _, r in df.iterrows()]
+            for i in range(0, len(rows), 10):                    # 10 rows per piece
+                pieces.append((f"{name}, {sheet} rows {i + 1}+", "\n".join(rows[i:i + 10])))
+        return pieces
+    return [(name, data.decode("utf-8", errors="ignore"))]       # txt / md
+
+
+def chunk(pieces):
+    chunks, step = [], CHUNK_WORDS - OVERLAP
+    for label, text in pieces:
+        words = text.split()
+        for i in range(0, len(words), step):
+            part = " ".join(words[i:i + CHUNK_WORDS])
+            if len(part) > 20:
+                chunks.append({"label": label, "text": part})
     return chunks
 
-def cite_label(c):
-    return f"{c['source']}, {c['loc']}" if c["loc"] else c["source"]
 
-# ========================= INDEX + HYBRID SEARCH =========================
-@st.cache_resource(show_spinner="Loading embedding model...")
+# ---------- index + search ----------
+@st.cache_resource(show_spinner="Loading embedding model (first run takes a minute)...")
 def load_embedder():
-    return SentenceTransformer(EMBED_MODEL_NAME, cache_folder="/tmp/model_cache", device="cpu")
-
-def build_index(chunks, embedder):
-    texts = [light_clean(c["text"]) for c in chunks]
-    emb = embedder.encode(texts, batch_size=32, convert_to_numpy=True,
-                          normalize_embeddings=True, show_progress_bar=False).astype("float32")
-    faiss_index = faiss.IndexFlatIP(emb.shape[1])  # inner product on normalized vectors = cosine
-    faiss_index.add(emb)
-    bm25 = BM25Okapi([tokenize(c["text"]) for c in chunks])
-    return {"chunks": chunks, "faiss": faiss_index, "bm25": bm25}
-
-def hybrid_search(query, index, embedder, top_k, w_sem):
-    """FAISS semantic + BM25, fused with Reciprocal Rank Fusion."""
-    q_vec = embedder.encode([light_clean(query)], convert_to_numpy=True,
-                            normalize_embeddings=True, show_progress_bar=False).astype("float32")
-    n = len(index["chunks"])
-    pool = min(n, top_k * 4)
-
-    scores, ids = index["faiss"].search(q_vec, pool)
-    sem_scores = {int(i): float(s) for i, s in zip(ids[0], scores[0]) if i >= 0}
-    best_sem = float(scores[0][0])
+    return SentenceTransformer(EMBED_MODEL, device="cpu")
 
-    bm = index["bm25"].get_scores(tokenize(query))
-    bm_rank = np.argsort(-bm)[:pool]
 
-    fused = defaultdict(float)
-    for r, i in enumerate(ids[0]):
-        if i >= 0:
-            fused[int(i)] += w_sem / (RRF_K + r + 1)
-    for r, i in enumerate(bm_rank):
-        if bm[i] > 0:
-            fused[int(i)] += (1 - w_sem) / (RRF_K + r + 1)
+def embed(embedder, texts):
+    return embedder.encode(texts, normalize_embeddings=True, show_progress_bar=False).astype("float32")
 
-    top = sorted(fused, key=fused.get, reverse=True)[:top_k]
-    return top, best_sem, sem_scores, bm
 
-# ========================= GENERATION =========================
-def generate_answer(llm, query, retrieved):
-    context = "\n\n---\n\n".join(f"[{cite_label(c)}]\n{c['text']}" for c in retrieved)
-    prompt = f"""You are a document assistant. Answer ONLY from the context below.
-- Answer in the same language as the question.
-- After each fact, cite its source exactly as shown in the context, e.g. [file, p.N].
-- Table rows use " | " between cells; the first line of a table chunk is its header.
-- If the context does not contain the answer, say it was not found in the documents and
-  suggest contacting the document owner. Do not guess.
-
-Context:
-{context}
-
-Question: {query}
-Answer:"""
-    return llm.invoke(prompt).content
-
-# ========================= STREAMLIT APP =========================
-st.set_page_config(page_title="Hybrid RAG Assistant", page_icon="📚", layout="wide")
-st.title("📚 Hybrid RAG Assistant")
-st.caption("FAISS semantic search + BM25, fused with Reciprocal Rank Fusion")
-
-with st.sidebar:
-    st.header("Settings")
-    top_k = st.slider("Chunks sent to the LLM (Top-K)", 3, 10, 5)
-    w_sem = st.slider("Semantic weight (BM25 gets the rest)", 0.0, 1.0, 0.6, 0.1)
-    min_rel = st.slider("Minimum relevance to answer", 0.0, 0.8, 0.25, 0.05,
-                        help="If the best semantic match is below this, the app says 'not found' instead of guessing.")
-
-embedder = load_embedder()
-
-groq_key = None
-try:
-    groq_key = st.secrets.get("GROQ_API_KEY")
-except Exception:
-    pass
-groq_key = groq_key or os.getenv("GROQ_API_KEY")
-if not groq_key:
-    st.error("GROQ_API_KEY not found. Add it to Streamlit secrets or environment variables.")
+def build_index(chunks, embedder):
+    vecs = embed(embedder, [c["text"] for c in chunks])
+    index = faiss.IndexFlatIP(vecs.shape[1])
+    index.add(vecs)
+    return {"chunks": chunks, "faiss": index, "bm25": BM25Okapi([tokenize(c["text"]) for c in chunks])}
+
+
+def search(query, idx, embedder, top_k):
+    pool = min(len(idx["chunks"]), top_k * 4)
+    _, ids = idx["faiss"].search(embed(embedder, [query]), pool)
+    bm = idx["bm25"].get_scores(tokenize(query))
+    bm_ids = [i for i in np.argsort(-bm)[:pool] if bm[i] > 0]
+    scores = defaultdict(float)
+    for ranking in (ids[0], bm_ids):
+        for r, i in enumerate(ranking):
+            if i >= 0:
+                scores[int(i)] += 1 / (RRF_K + r + 1)
+    return sorted(scores, key=scores.get, reverse=True)[:top_k]
+
+
+# ---------- app ----------
+key = os.getenv("GROQ_API_KEY")
+if not key:
+    try:
+        key = st.secrets.get("GROQ_API_KEY")
+    except Exception:
+        pass
+if not key:
+    st.error("GROQ_API_KEY not found. Add it as a secret in the Space settings.")
     st.stop()
-llm = ChatGroq(temperature=0.1, model_name=LLM_MODEL_NAME, groq_api_key=groq_key)
-
-uploaded_files = st.file_uploader(
-    "Upload documents",
-    type=["pdf", "docx", "pptx", "xlsx", "csv", "txt", "md"],
-    accept_multiple_files=True,
-)
 
-# Rebuild the index whenever the set of uploaded files changes
-if uploaded_files:
-    signature = tuple(sorted((f.name, f.size) for f in uploaded_files))
-    if st.session_state.get("signature") != signature:
-        all_blocks, report = [], {}
-        with st.spinner("Reading documents..."):
-            for f in uploaded_files:
-                try:
-                    blocks, stats = extract_file(f.name, f.getvalue())
-                except Exception as e:
-                    blocks, stats = [], {"pages": 0, "tables": 0, "notes": [f"Could not read file: {e}"]}
-                all_blocks.extend(blocks)
-                report[f.name] = stats
-        chunks = chunk_blocks(all_blocks)
-        if not chunks:
-            st.error("No readable text found in the uploaded files.")
-            st.session_state.pop("index", None)
-            st.session_state.pop("signature", None)
-            with st.expander("Extraction report", expanded=True):
-                for fname, s in report.items():
-                    for note in s["notes"]:
-                        st.write(f"**{fname}**: {note}")
-            st.stop()
-        with st.spinner("Building FAISS + BM25 indexes..."):
+top_k = st.sidebar.slider("Chunks sent to the LLM", 3, 10, 5)
+files = st.file_uploader("Upload documents", type=["pdf", "docx", "xlsx", "csv", "txt", "md"],
+                         accept_multiple_files=True)
+
+if files:
+    embedder = load_embedder()
+    sig = tuple((f.name, f.size) for f in files)
+    if st.session_state.get("sig") != sig:
+        with st.spinner("Indexing..."):
+            pieces = [p for f in files for p in read_file(f.name, f.getvalue())]
+            chunks = chunk(pieces)
+            if not chunks:
+                st.error("No readable text found (scanned PDFs are not supported).")
+                st.stop()
             st.session_state.index = build_index(chunks, embedder)
-        st.session_state.report = report
-        st.session_state.signature = signature
-    st.success(f"Indexed {len(st.session_state.index['chunks'])} chunks from {len(uploaded_files)} file(s)")
-
-    with st.expander("📋 Extraction report"):
-        for fname, s in st.session_state.report.items():
-            st.write(f"**{fname}**: {s['pages']} page(s)/sheet(s), {s['tables']} table(s)")
-            for note in s["notes"]:
-                st.caption(f"⚠️ {note}")
-else:
-    st.session_state.pop("index", None)
-    st.session_state.pop("signature", None)
-
-# Query form: the LLM is only called when the user presses the button
-with st.form("ask"):
-    query = st.text_input("Ask a question about your documents (Arabic or English):")
-    submitted = st.form_submit_button("Ask")
-
-if submitted and query.strip():
-    index = st.session_state.get("index")
-    if not index:
-        st.warning("Upload at least one document first.")
-    else:
-        top, best_sem, sem_scores, bm = hybrid_search(query, index, embedder, top_k, w_sem)
-        retrieved = [index["chunks"][i] for i in top]
-
-        st.subheader("Answer")
-        if not retrieved or best_sem < min_rel:
-            st.info("Not found in the uploaded documents.")
-        else:
-            with st.spinner("Generating answer..."):
-                st.write(generate_answer(llm, query, retrieved))
-
-        with st.expander("🔍 Retrieved chunks"):
-            for rank, i in enumerate(top, 1):
-                c = index["chunks"][i]
-                sem_txt = f"{sem_scores[i]:.2f}" if i in sem_scores else "n/a"
-                st.markdown(f"**{rank}. {cite_label(c)}** ({c['kind']})  ·  semantic {sem_txt}  ·  BM25 {bm[i]:.2f}")
+            st.session_state.sig = sig
+    st.success(f"Indexed {len(st.session_state.index['chunks'])} chunks")
+
+    query = st.text_input("Ask a question (Arabic or English):")
+    if st.button("Ask") and query.strip():
+        idx = st.session_state.index
+        hits = [idx["chunks"][i] for i in search(query, idx, embedder, top_k)]
+        context = "\n\n---\n\n".join(f"[{c['label']}]\n{c['text']}" for c in hits)
+        prompt = (
+            "Answer ONLY from the context below, in the same language as the question. "
+            "Cite sources like [file, p.N]. If the answer is not in the context, say it was "
+            f"not found in the documents.\n\nContext:\n{context}\n\nQuestion: {query}\nAnswer:"
+        )
+        llm = ChatGroq(temperature=0.1, model_name=LLM_MODEL, groq_api_key=key)
+        with st.spinner("Generating answer..."):
+            st.write(llm.invoke(prompt).content)
+        with st.expander("Retrieved chunks"):
+            for c in hits:
+                st.markdown(f"**{c['label']}**")
                 st.text(c["text"])
 else:
-    st.info("Upload documents, type a question, and press Ask.")
+    st.info("Upload documents to begin.")