update

ingest.py

@@ -1,5 +1,5 @@
 import os
-from datasets import load_dataset
+from datasets import load_dataset, Features, Value
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
@@ -10,33 +10,47 @@ def run_ingestion():
     os.makedirs(KB_DIR, exist_ok=True)
 
     print(f"⬇️ Loading dataset from {HF_DATASET_REPO}...")
-    # decode(False) prevents the library from turning bytes into pdfplumber objects
-    dataset = load_dataset(HF_DATASET_REPO, split="train").with_format("binary").decode(False)
     
+    # We cast the 'pdf' column to binary to prevent automatic decoding into a 'PDF' object
+    features = Features({"pdf": Value("binary"), "file_name": Value("string")})
+    
+    # Use decode=False to get raw bytes
+    dataset = load_dataset(HF_DATASET_REPO, split="train", decode=False)
+    
+    print(f"📊 Dataset columns: {dataset.column_names}")
+
     pdf_paths = []
     for i, row in enumerate(dataset):
-        # Determine filename and raw data column
-        fname = row.get("filename") or row.get("file_name") or f"doc_{i}.pdf"
-        # Access the raw 'bytes' from the 'pdf' column
-        pdf_feature = row.get("pdf")
+        # Your logs show the column is named 'pdf'
+        pdf_data = row.get("pdf")
         
-        if pdf_feature is None:
+        # Determine filename
+        fname = row.get("file_name") or row.get("filename") or f"document_{i}.pdf"
+        path = os.path.join(KB_DIR, fname)
+
+        if pdf_data is None:
+            print(f"⚠️ Row {i} has no data, skipping.")
             continue
 
-        path = os.path.join(KB_DIR, fname)
         with open(path, "wb") as f:
-            if isinstance(pdf_feature, dict) and "bytes" in pdf_feature:
-                f.write(pdf_feature["bytes"])
+            # When decode=False, pdf_data is usually a dict like {'bytes': b'...', 'path': None}
+            if isinstance(pdf_data, dict) and "bytes" in pdf_data:
+                f.write(pdf_data["bytes"])
             else:
-                f.write(pdf_feature)
+                f.write(pdf_data)
+        
         pdf_paths.append(path)
 
     print(f"📄 Processing {len(pdf_paths)} documents...")
     docs = []
     for p in pdf_paths:
-        loader = PyPDFLoader(p)
-        docs.extend(loader.load())
+        try:
+            loader = PyPDFLoader(p)
+            docs.extend(loader.load())
+        except Exception as e:
+            print(f"❌ Could not load {p}: {e}")
 
+    # Split text into chunks
     splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
     splits = splitter.split_documents(docs)
 
@@ -48,7 +62,7 @@ def run_ingestion():
         embedding=embeddings,
         persist_directory=CHROMA_DIR
     )
-    print(f"✅ Ingestion complete. Data saved to {CHROMA_DIR}")
+    print(f"✅ Ingestion complete. DB saved to {CHROMA_DIR}")
 
 if __name__ == "__main__":
     run_ingestion()