Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import os
 import tempfile
 import gradio as gr
@@ -16,7 +15,7 @@ import numpy as np
 from transformers import pipeline
 
 # CONFIG
-HF_GENERATION_MODEL = os.environ.get("HF_GENERATION_MODEL","google/flan-t5-large")  # change to DeepSeek if ready
+HF_GENERATION_MODEL = os.environ.get("HF_GENERATION_MODEL", "google/flan-t5-large")  # change to DeepSeek if ready
 EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 INDEX_PATH = "faiss_index.index"
 METADATA_PATH = "metadata.json"
@@ -48,7 +47,7 @@ def extract_text_from_excel(file):
 def extract_text_from_url(url):
     r = requests.get(url, timeout=10)
     soup = BeautifulSoup(r.text, "lxml")
-    for s in soup(["script","style","aside","nav","footer"]):
+    for s in soup(["script", "style", "aside", "nav", "footer"]):
         s.decompose()
     text = soup.get_text(separator="\n")
     return text
@@ -59,51 +58,76 @@ splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
 def ingest_sources(files, urls):
     docs = []
     metadata = []
+
     for f in files:
-        name = f.name
-tmp = tempfile.NamedTemporaryFile(delete=False)
-# handle both NamedString and normal file
-if hasattr(f, "read"):
-    tmp.write(f.read())
-else:
-    tmp.write(f.encode("utf-8"))
-tmp.flush()
-tmp.close()
+        # make sure we have a temp file
+        tmp = tempfile.NamedTemporaryFile(delete=False)
+
+        # handle different types of file objects
+        if hasattr(f, "read"):  # normal file
+            tmp.write(f.read())
+            name = getattr(f, "name", "uploaded_file")
+        elif isinstance(f, str):  # NamedString or text
+            tmp.write(f.encode("utf-8"))
+            name = "uploaded_text.txt"
+        elif isinstance(f, dict) and "data" in f:  # HF file dict
+            tmp.write(f["data"])
+            name = f.get("name", "uploaded_file")
+        else:
+            raise ValueError(f"Unknown file type: {type(f)}")
+
+        tmp.flush()
+        tmp.close()
+
+        # extract text depending on file type
         if name.lower().endswith(".pdf"):
             text = extract_text_from_pdf(tmp.name)
         elif name.lower().endswith(".docx"):
             text = extract_text_from_docx(tmp.name)
-        elif name.lower().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()
+
         os.unlink(tmp.name)
+
         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"})
+
+    # handle URLs
     for u in urls:
+        u = u.strip()
+        if not u:
+            continue
         try:
             text = extract_text_from_url(u)
             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"})
         except Exception as e:
             print("url error", u, e)
+
+    if not docs:
+        return "No valid documents or URLs found."
+
     embeddings = embed_model.encode(docs, show_progress_bar=True, convert_to_numpy=True)
     dim = embeddings.shape[1]
+
     if os.path.exists(INDEX_PATH):
         index = faiss.read_index(INDEX_PATH)
-        old_meta = json.load(open(METADATA_PATH,"r"))
+        old_meta = json.load(open(METADATA_PATH, "r"))
         index.add(embeddings)
         old_meta.extend(metadata)
-        json.dump(old_meta, open(METADATA_PATH,"w"))
+        json.dump(old_meta, open(METADATA_PATH, "w"))
     else:
         index = faiss.IndexFlatL2(dim)
         index.add(embeddings)
-        json.dump(metadata, open(METADATA_PATH,"w"))
+        json.dump(metadata, open(METADATA_PATH, "w"))
+
     faiss.write_index(index, INDEX_PATH)
     return f"Ingested {len(docs)} chunks from {len(files)} files and {len(urls)} urls."
 
@@ -113,27 +137,32 @@ def retrieve_topk(query, k=5):
         return []
     index = faiss.read_index(INDEX_PATH)
     D, I = index.search(q_emb, k)
-    metadata = json.load(open(METADATA_PATH,"r"))
+    metadata = json.load(open(METADATA_PATH, "r"))
     results = []
     for idx in I[0]:
         if idx < len(metadata):
             results.append((metadata[idx], idx))
     return results
 
-gen_pipeline = pipeline("text2text-generation", model=HF_GENERATION_MODEL, device=0 if os.environ.get("HF_DEVICE","cpu")!="cpu" else -1)
+gen_pipeline = pipeline("text2text-generation", model=HF_GENERATION_MODEL, device=0 if os.environ.get("HF_DEVICE", "cpu") != "cpu" else -1)
 
 def ask_prompt(prompt, top_k=5):
     hits = retrieve_topk(prompt, k=top_k)
     if not hits:
         return "No documents ingested. Use Ingest first."
+
     context_parts = []
     sources = []
     for meta, idx in hits:
         sources.append(f"{meta['source']} (chunk {meta['chunk']})")
         context_parts.append(f"[{meta['source']} - chunk {meta['chunk']}]")
+
     context = "\n\n".join(context_parts)
-    system_instruction = ("You are an AI research assistant. Use the contextual chunks below to answer the user's question. "
-                          "Provide a concise answer, then list sources in order of relevance with short snippets or page numbers.")
+    system_instruction = (
+        "You are an AI research assistant. Use the contextual chunks below to answer the user's question. "
+        "Provide a concise answer, then list sources in order of relevance."
+    )
+
     prompt_text = f"{system_instruction}\n\nCONTEXT:\n{context}\n\nQUESTION:\n{prompt}\n\nAnswer:"
     out = gen_pipeline(prompt_text, max_length=512, do_sample=False)[0]["generated_text"]
     out = out + "\n\nSources:\n" + "\n".join(sources)