ingest fix

src/ingest.py

@@ -1,76 +1,56 @@
-# src/ingest.py
-import pickle
-from pathlib import Path
-
+import numpy as np
 import faiss
-import pdfplumber
-from datasets import load_dataset
+import pickle
 from sentence_transformers import SentenceTransformer
-
 from config import (
-    DATASET_NAME,
     FAISS_INDEX_PATH,
     DOCS_PATH,
     CHUNK_SIZE,
     CHUNK_OVERLAP,
 )
 
-
-def chunk_text(text: str):
-    chunks = []
-    start = 0
-
-    while start < len(text):
-        end = start + CHUNK_SIZE
-        chunks.append(text[start:end])
-        start = end - CHUNK_OVERLAP
-
-    return chunks
-
-
 def build_index():
-    print("📥 Loading HF dataset...")
-    dataset = load_dataset(DATASET_NAME, split="train")
-
-    embedder = SentenceTransformer("all-MiniLM-L6-v2")
-
-    documents = []
-    embeddings = []
+    from datasets import load_dataset
+    from pypdf import PdfReader
+    import io
 
-    print(f"📄 Dataset rows: {len(dataset)}")
+    dataset = load_dataset("OnlyTheTruth03/ott", split="train")
 
-    for row_idx, row in enumerate(dataset):
-        pdf_obj = row["pdf"]
+    texts = []
 
-        # ✅ THIS IS THE KEY
-        pdf_path = pdf_obj.path
+    for row in dataset:
+        pdf_obj = row["pdf"]  # HF auto parquet object
+        pdf_bytes = pdf_obj["bytes"]
 
-        print(f"➡️ Processing PDF {row_idx + 1}: {pdf_path}")
+        reader = PdfReader(io.BytesIO(pdf_bytes))
+        for page in reader.pages:
+            text = page.extract_text()
+            if text:
+                texts.append(text)
 
-        with pdfplumber.open(pdf_path) as pdf:
-            for page in pdf.pages:
-                text = page.extract_text()
-                if not text:
-                    continue
-
-                for chunk in chunk_text(text):
-                    documents.append(chunk)
-                    embeddings.append(embedder.encode(chunk))
+    # -------- Chunking --------
+    chunks = []
+    for text in texts:
+        start = 0
+        while start < len(text):
+            end = start + CHUNK_SIZE
+            chunks.append(text[start:end])
+            start = end - CHUNK_OVERLAP
 
-    if not documents:
-        raise RuntimeError("❌ No text extracted from PDFs")
+    # -------- Embeddings --------
+    model = SentenceTransformer("all-MiniLM-L6-v2")
+    embeddings = model.encode(chunks, show_progress_bar=True)
 
-    print(f"🧠 Creating FAISS index with {len(documents)} chunks")
+    # 🔥 CRITICAL FIX
+    embeddings = np.array(embeddings).astype("float32")
 
-    dim = len(embeddings[0])
+    dim = embeddings.shape[1]
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
 
-    FAISS_INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)
-
     faiss.write_index(index, str(FAISS_INDEX_PATH))
 
     with open(DOCS_PATH, "wb") as f:
-        pickle.dump(documents, f)
+        pickle.dump(chunks, f)
 
-    print("✅ FAISS index built successfully")
+    return index, chunks