Update app.py

app.py

@@ -1,7 +1,9 @@
-# app.py - Simple hybrid RAG (semantic + BM25, fused with RRF). PDF, DOCX, XLSX/CSV, TXT. Arabic + English.
+# 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.
 import io
 import os
 import re
+import unicodedata
 from collections import defaultdict
 
 import faiss
@@ -11,136 +13,352 @@ 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
 
-EMBED_MODEL = "paraphrase-multilingual-MiniLM-L12-v2"
-LLM_MODEL = "llama-3.3-70b-versatile"
-CHUNK_WORDS, OVERLAP, RRF_K = 80, 15, 60
+# ========================= 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
 
-st.set_page_config(page_title="RAG Assistant", layout="wide")
-st.title("RAG Assistant")
-st.caption("Semantic search + BM25, fused with Reciprocal Rank Fusion")
+# ========================= TEXT NORMALIZATION =========================
+DIACRITICS_RE = re.compile(r"[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06ED]")
+DIGIT_MAP = str.maketrans("٠١٢٣٤٥٦٧٨٩۰۱۲۳۴۵۶۷۸۹", "01234567890123456789")
 
+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):
-    text = re.sub(r"[\u064B-\u065F\u0640]", "", text)           # diacritics, tatweel
-    text = re.sub("[إأآ]", "ا", text).replace("ى", "ي").replace("ة", "ه")
-    return re.findall(r"\w+", text.lower())
+    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."]}
 
+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 read_file(name, data):
-    """Return a list of (label, text) pieces."""
+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):
     lower = name.lower()
     if lower.endswith(".pdf"):
-        doc = fitz.open(stream=data, filetype="pdf")
-        return [(f"{name}, p.{i}", p.get_text()) for i, p in enumerate(doc, 1)]
+        return extract_pdf(name, data)
     if lower.endswith(".docx"):
-        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)]
+        return extract_docx(name, data)
+    if lower.endswith(".pptx"):
+        return extract_pptx(name, data)
     if lower.endswith((".xlsx", ".csv")):
-        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 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))
     return chunks
 
+def cite_label(c):
+    return f"{c['source']}, {c['loc']}" if c["loc"] else c["source"]
 
-# ---------- index + search ----------
-@st.cache_resource(show_spinner="Loading embedding model (first run takes a minute)...")
+# ========================= INDEX + HYBRID SEARCH =========================
+@st.cache_resource(show_spinner="Loading embedding model...")
 def load_embedder():
-    return SentenceTransformer(EMBED_MODEL, device="cpu")
+    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 embed(embedder, texts):
-    return embedder.encode(texts, normalize_embeddings=True, show_progress_bar=False).astype("float32")
+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])
 
-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.")
+    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)
+
+    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.")
     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,
+)
 
-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()
+# 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..."):
             st.session_state.index = build_index(chunks, embedder)
-            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.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.text(c["text"])
 else:
-    st.info("Upload documents to begin.")
+    st.info("Upload documents, type a question, and press Ask.")