Read dataset as temp file

src/ingest.py

@@ -1,47 +1,112 @@
-from datasets import load_dataset
-from pypdf import PdfReader
+# src/ingest.py
+
 import os
-import faiss
 import pickle
+import tempfile
+
+import faiss
+import numpy as np
+from datasets import load_dataset
+from pypdf import PdfReader
 from sentence_transformers import SentenceTransformer
 
+
+# ================= CONFIG =================
+DATASET_NAME = "OnlyTheTruth03/ott"
+SPLIT = "train"
+
 INDEX_DIR = "src/index"
-os.makedirs(INDEX_DIR, exist_ok=True)
+FAISS_PATH = os.path.join(INDEX_DIR, "faiss.index")
+DOCS_PATH = os.path.join(INDEX_DIR, "documents.pkl")
+
+EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+CHUNK_SIZE = 500
+CHUNK_OVERLAP = 50
+
+
+# ================= HELPERS =================
+def chunk_text(text, size=500, overlap=50):
+    words = text.split()
+    chunks = []
+
+    start = 0
+    while start < len(words):
+        end = start + size
+        chunk = " ".join(words[start:end])
+        chunks.append(chunk)
+        start += size - overlap
+
+    return chunks
+
+
+# ================= INGEST =================
+def ingest():
+    # Avoid re-indexing on every Streamlit rerun
+    if os.path.exists(FAISS_PATH) and os.path.exists(DOCS_PATH):
+        print("✅ FAISS index already exists. Skipping ingestion.")
+        return
+
+    print("📥 Loading HF dataset...")
+    dataset = load_dataset(DATASET_NAME, split=SPLIT)
+
+    print("🔎 Loading embedding model...")
+    model = SentenceTransformer(EMBED_MODEL_NAME)
+
+    all_chunks = []
+    documents = []
+
+    for row in dataset:
+        # HF datasets provide file objects, not paths
+        file_obj = list(row.values())[0]
+
+        # Write PDF bytes to temp file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+            tmp.write(file_obj["bytes"])
+            tmp_path = tmp.name
+
+        reader = PdfReader(tmp_path)
+
+        for page_num, page in enumerate(reader.pages):
+            text = page.extract_text()
+            if not text:
+                continue
 
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
+            chunks = chunk_text(text, CHUNK_SIZE, CHUNK_OVERLAP)
+            for chunk in chunks:
+                all_chunks.append(chunk)
+                documents.append({
+                    "text": chunk,
+                    "source": "ott",
+                    "page": page_num + 1
+                })
 
-dataset = load_dataset("OnlyTheTruth03/ott", split="train")
+        os.remove(tmp_path)
 
-print("Dataset columns:", dataset.column_names)
-print("First row:", dataset[0])
+    if not all_chunks:
+        raise RuntimeError("❌ No text extracted from PDFs.")
 
-documents = []
-texts = []
+    print(f"🧠 Creating embeddings for {len(all_chunks)} chunks...")
+    embeddings = model.encode(
+        all_chunks,
+        show_progress_bar=True,
+        convert_to_numpy=True
+    )
 
-for row in dataset:
-    #pdf_path = row["file"]
-    pdf_path = list(row.values())[0]
-    reader = PdfReader(pdf_path)
+    embeddings = embeddings.astype("float32")
 
-    for page_num, page in enumerate(reader.pages):
-        text = page.extract_text()
-        if text:
-            texts.append(text)
-            documents.append({
-                "text": text,
-                "source": "ott",
-                "page": page_num + 1
-            })
+    print("📦 Building FAISS index...")
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
 
-embeddings = embedder.encode(texts, show_progress_bar=True)
-embeddings = embeddings.astype("float32")
+    os.makedirs(INDEX_DIR, exist_ok=True)
 
-index = faiss.IndexFlatL2(embeddings.shape[1])
-index.add(embeddings)
+    faiss.write_index(index, FAISS_PATH)
+    with open(DOCS_PATH, "wb") as f:
+        pickle.dump(documents, f)
 
-faiss.write_index(index, f"{INDEX_DIR}/faiss.index")
+    print("✅ Ingestion completed successfully.")
 
-with open(f"{INDEX_DIR}/documents.pkl", "wb") as f:
-    pickle.dump(documents, f)
 
-print("✅ Ingestion complete")
+# ================= AUTO-RUN =================
+if __name__ == "__main__":
+    ingest()