Update app.py

app.py

@@ -1,7 +1,7 @@
+# app.py
 import os
 import tempfile
 import gradio as gr
-from typing import List
 import json
 import pandas as pd
 import requests
@@ -14,33 +14,36 @@ import faiss
 import numpy as np
 from transformers import pipeline
 
-# -----------------------------
+# ==============================
 # CONFIG
-# -----------------------------
-HF_GENERATION_MODEL = os.environ.get("HF_GENERATION_MODEL", "google/flan-t5-large")  # You can switch later to DeepSeek
-EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-MiniLM-L3-v2"  # Faster, smaller
+# ==============================
+HF_GENERATION_MODEL = os.environ.get("HF_GENERATION_MODEL", "google/flan-t5-large")
+EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 INDEX_PATH = "faiss_index.index"
 METADATA_PATH = "metadata.json"
 
+# ==============================
 # Load embedding model
+# ==============================
 embed_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
 
-# -----------------------------
-# FILE HELPERS
-# -----------------------------
+# ==============================
+# Helper text extractors
+# ==============================
 def extract_text_from_pdf(file_path):
     reader = PdfReader(file_path)
-    return "\n\n".join(page.extract_text() or "" for page in reader.pages)
+    pages = [p.extract_text() or "" for p in reader.pages]
+    return "\n\n".join(pages)
 
 def extract_text_from_docx(file_path):
     doc = docx.Document(file_path)
     return "\n\n".join(p.text for p in doc.paragraphs)
 
 def extract_text_from_excel(file_path):
-    dfs = pd.read_excel(file_path, sheet_name=None)
+    df_dict = pd.read_excel(file_path, sheet_name=None)
     out = []
-    for name, df in dfs.items():
-        out.append(f"Sheet: {name}")
+    for sheet, df in df_dict.items():
+        out.append(f"Sheet: {sheet}")
         out.append(df.fillna("").to_csv(index=False))
     return "\n\n".join(out)
 
@@ -49,158 +52,169 @@ def extract_text_from_url(url):
     soup = BeautifulSoup(r.text, "lxml")
     for s in soup(["script", "style", "aside", "nav", "footer"]):
         s.decompose()
-    return soup.get_text(separator="\n")
+    text = soup.get_text(separator="\n")
+    return text
 
-# -----------------------------
-# CHUNKER (larger = faster)
-# -----------------------------
-splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=100)
+# ==============================
+# Text chunking setup
+# ==============================
+splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
 
-# -----------------------------
-# INGESTION
-# -----------------------------
+# ==============================
+# Ingestion function
+# ==============================
 def ingest_sources(files, urls):
-    docs, metadata = [], []
-
-    if os.path.exists(INDEX_PATH) and os.path.exists(METADATA_PATH):
-        return "Index already exists. Delete the files to re-ingest."
+    docs = []
+    metadata = []
 
+    # Handle uploaded files
     for f in files:
+        name = f.name
         tmp = tempfile.NamedTemporaryFile(delete=False)
         try:
             if hasattr(f, "read"):
-                data = f.read()
-                if isinstance(data, str):
-                    data = data.encode("utf-8")
-                tmp.write(data)
-                name = getattr(f, "name", "uploaded_file")
-            elif isinstance(f, dict) and "data" in f:
-                data = f["data"]
-                if isinstance(data, str):
-                    data = data.encode("utf-8")
-                tmp.write(data)
-                name = f.get("name", "uploaded_file")
-            elif isinstance(f, str):
-                tmp.write(f.encode("utf-8"))
-                name = "uploaded_text.txt"
+                tmp.write(f.read())
             else:
-                tmp.close()
-                os.unlink(tmp.name)
-                return f"Unknown upload type: {type(f)}"
-        finally:
+                tmp.write(f.encode("utf-8"))
             tmp.flush()
             tmp.close()
 
-        try:
-            low = name.lower()
-            if low.endswith(".pdf"):
+            if name.lower().endswith(".pdf"):
                 text = extract_text_from_pdf(tmp.name)
-            elif low.endswith(".docx"):
+            elif name.lower().endswith(".docx"):
                 text = extract_text_from_docx(tmp.name)
-            elif low.endswith((".xls", ".xlsx")):
+            elif name.lower().endswith((".xls", ".xlsx")):
                 text = extract_text_from_excel(tmp.name)
             else:
                 with open(tmp.name, "r", encoding="utf-8", errors="ignore") as fh:
                     text = fh.read()
-        except Exception as e:
-            print(f"Extraction error for {name}: {e}")
+        finally:
             os.unlink(tmp.name)
-            continue
-
-        os.unlink(tmp.name)
 
-        for i, c in enumerate(splitter.split_text(text)):
+        chunks = splitter.split_text(text)
+        for i, c in enumerate(chunks):
             docs.append(c)
-            metadata.append({"source": name, "chunk": i, "type": "file"})
+            metadata.append({"source": name, "chunk": i, "type": "file", "text": c})
 
-    for u in urls or []:
-        u = (u or "").strip()
-        if not u:
+    # Handle URLs
+    for u in urls:
+        if not u.strip():
             continue
         try:
             text = extract_text_from_url(u)
-            for i, c in enumerate(splitter.split_text(text)):
+            chunks = splitter.split_text(text)
+            for i, c in enumerate(chunks):
                 docs.append(c)
-                metadata.append({"source": u, "chunk": i, "type": "url"})
+                metadata.append({"source": u, "chunk": i, "type": "url", "text": c})
         except Exception as e:
-            print(f"URL fetch error for {u}: {e}")
+            print("URL error:", u, e)
 
     if not docs:
-        return "No content ingested (empty or failed files)."
+        return "No text extracted from files or URLs."
 
-    try:
-        embeddings = embed_model.encode(docs, show_progress_bar=True, convert_to_numpy=True)
-    except Exception as e:
-        return f"Embedding error: {e}"
+    embeddings = embed_model.encode(docs, show_progress_bar=True, convert_to_numpy=True)
+    dim = embeddings.shape[1]
 
-    try:
-        dim = embeddings.shape[1]
+    if os.path.exists(INDEX_PATH):
+        index = faiss.read_index(INDEX_PATH)
+        old_meta = json.load(open(METADATA_PATH, "r", encoding="utf-8"))
+        index.add(embeddings)
+        old_meta.extend(metadata)
+        json.dump(old_meta, open(METADATA_PATH, "w", encoding="utf-8"))
+    else:
         index = faiss.IndexFlatL2(dim)
         index.add(embeddings)
-        faiss.write_index(index, INDEX_PATH)
-        with open(METADATA_PATH, "w", encoding="utf-8") as fh:
-            json.dump(metadata, fh)
-    except Exception as e:
-        return f"Indexing error: {e}"
+        json.dump(metadata, open(METADATA_PATH, "w", encoding="utf-8"))
 
-    return f"Ingested {len(docs)} chunks from {len(files)} files and {len(urls)} URLs."
+    faiss.write_index(index, INDEX_PATH)
+    return f"Ingested {len(docs)} text chunks from {len(files)} files and {len(urls)} URLs."
 
-# -----------------------------
-# RETRIEVAL
-# -----------------------------
+# ==============================
+# Retrieve top matching chunks
+# ==============================
 def retrieve_topk(query, k=5):
     if not os.path.exists(INDEX_PATH):
         return []
     q_emb = embed_model.encode([query], convert_to_numpy=True)
     index = faiss.read_index(INDEX_PATH)
     D, I = index.search(q_emb, k)
-    metadata = json.load(open(METADATA_PATH))
+    metadata = json.load(open(METADATA_PATH, "r", encoding="utf-8"))
     results = []
     for idx in I[0]:
         if idx < len(metadata):
             results.append(metadata[idx])
     return results
 
-# -----------------------------
-# GENERATION PIPELINE
-# -----------------------------
-gen_pipeline = pipeline("text2text-generation", model=HF_GENERATION_MODEL, device=-1)
-
+# ==============================
+# Generation pipeline
+# ==============================
+gen_pipeline = pipeline(
+    "text2text-generation",
+    model=HF_GENERATION_MODEL,
+    device=0 if os.environ.get("HF_DEVICE", "cpu") != "cpu" else -1,
+)
+
+# ==============================
+# Ask prompt
+# ==============================
 def ask_prompt(prompt, top_k=5):
+    if not os.path.exists(INDEX_PATH) or not os.path.exists(METADATA_PATH):
+        return "No documents ingested yet."
+
     hits = retrieve_topk(prompt, k=top_k)
     if not hits:
-        return "No documents ingested yet."
+        return "No relevant context found. Try ingesting more content."
 
+    # Collect context text
+    context_parts = [h["text"] for h in hits if "text" in h]
     sources = [f"{h['source']} (chunk {h['chunk']})" for h in hits]
-    context = "\n\n".join(sources)
+
+    context = "\n\n".join(context_parts)
+    if not context.strip():
+        return "No readable text found in the ingested files."
 
     system_instruction = (
-        "You are a research assistant. Use the context below to answer the question clearly and briefly.\n"
+        "You are a helpful research assistant. Read the provided context carefully "
+        "and answer the question accurately and concisely."
     )
-    full_prompt = f"{system_instruction}\nCONTEXT:\n{context}\n\nQUESTION:\n{prompt}\n\nAnswer:"
 
-    out = gen_pipeline(full_prompt, max_length=400, do_sample=False)[0]["generated_text"]
+    full_prompt = f"{system_instruction}\n\nCONTEXT:\n{context}\n\nQUESTION:\n{prompt}\n\nAnswer:"
+
+    try:
+        out = gen_pipeline(full_prompt, max_length=400, do_sample=False)[0]["generated_text"]
+    except Exception as e:
+        return f"Model generation failed: {e}"
+
     return out + "\n\nSources:\n" + "\n".join(sources)
 
-# -----------------------------
-# GRADIO UI
-# -----------------------------
+# ==============================
+# Gradio UI
+# ==============================
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Research Assistant (light version)\nUpload PDFs, Word, Excel, or URLs. Click **Ingest**, then ask your question.")
-
+    gr.Markdown(
+        "# 🧠 Research Assistant (Prototype)\nUpload files or paste URLs, click **Ingest**, then ask your question."
+    )
     with gr.Row():
         with gr.Column():
-            file_in = gr.File(label="Upload files", file_count="multiple")
-            urls_in = gr.Textbox(label="URLs (one per line)", placeholder="https://example.com")
+            file_in = gr.File(
+                label="Upload files (pdf/docx/xlsx/txt)", file_count="multiple"
+            )
+            urls_in = gr.Textbox(
+                label="URLs (one per line)",
+                placeholder="https://example.com/article",
+            )
             ingest_btn = gr.Button("Ingest")
             ingest_output = gr.Textbox(label="Ingest status")
-
         with gr.Column():
-            prompt_in = gr.Textbox(label="Your question", lines=3)
+            prompt_in = gr.Textbox(label="Your question", lines=4)
             ask_btn = gr.Button("Ask")
-            answer_out = gr.Textbox(label="Answer", lines=10)
+            answer_out = gr.Textbox(label="Answer", lines=12)
 
-    ingest_btn.click(lambda f, u: ingest_sources(f or [], (u or "").splitlines()), inputs=[file_in, urls_in], outputs=ingest_output)
+    ingest_btn.click(
+        lambda files, urls: ingest_sources(files or [], (urls or "").splitlines()),
+        inputs=[file_in, urls_in],
+        outputs=ingest_output,
+    )
     ask_btn.click(lambda p: ask_prompt(p, top_k=5), inputs=prompt_in, outputs=answer_out)
 
 if __name__ == "__main__":