Create ingestion.py

ingestion.py

@@ -0,0 +1,41 @@
+import os
+import glob
+from datasets import Dataset
+from unstructured.partition.pdf import partition_pdf
+from transformers import RagTokenizer
+
+def ingest_and_push(dataset_name="username/mealplan-chunks"):
+    # Initialize tokenizer for token-aware splitting
+    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
+    texts, sources, pages = [], [], []
+
+    for pdf_path in glob.glob("pdfs/*.pdf"):
+        book = os.path.basename(pdf_path)
+        pages_data = partition_pdf(filename=pdf_path)
+        for page_num, page in enumerate(pages_data, start=1):
+            # Encode page text into token windows
+            enc = tokenizer(
+                page.text,
+                max_length=800,
+                truncation=True,
+                return_overflowing_tokens=True,
+                stride=50,
+                return_tensors="pt"
+            )
+            # Decode each token window back to text chunk
+            for token_ids in enc["input_ids"]:
+                chunk = tokenizer.decode(token_ids, skip_special_tokens=True)
+                texts.append(chunk)
+                sources.append(book)
+                pages.append(page_num)
+
+    # Build HF Dataset
+    ds = Dataset.from_dict({
+        "text": texts,
+        "source": sources,
+        "page": pages
+    })
+    ds.push_to_hub(dataset_name, token=True)
+
+if __name__ == "__main__":
+    ingest_and_push()