""" Data ingestion script — loads jamescalam/ai-arxiv2-chunks into PostgreSQL. Usage: python -m scripts.ingest_data [--limit N] [--batch-size 256] [--role researcher] The script: 1. Streams the HuggingFace dataset (no full download required for large corpora). 2. Upserts documents (deduplication by arxiv_id). 3. Embeds chunk content locally with sentence-transformers. 4. Bulk-inserts chunks with their 384-dim vectors via COPY (fast). 5. Grants access to the specified role for all ingested documents. NOTE: 241k chunks × 384 floats ≈ 370 MB of vectors. Embedding on CPU takes ~2-4 hours for the full dataset; use --limit for quick smoke tests. """ import argparse import logging import time import uuid from typing import Iterator import numpy as np import psycopg2 import psycopg2.extras from datasets import load_dataset from sentence_transformers import SentenceTransformer from tqdm import tqdm from app.config import get_settings logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") logger = logging.getLogger(__name__) DATASET_NAME = "jamescalam/ai-arxiv2-chunks" EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2" EMBED_BATCH = 64 # chunks per embedding batch DB_BATCH = 256 # rows per DB insert # ── Dataset helpers ─────────────────────────────────────────────────────────── def _iter_dataset(limit: int | None) -> Iterator[dict]: """Stream dataset rows; respect optional limit.""" ds = load_dataset(DATASET_NAME, split="train", streaming=True) for i, row in enumerate(ds): if limit and i >= limit: break yield row def _arxiv_url(arxiv_id: str) -> str: return f"https://arxiv.org/abs/{arxiv_id}" def _full_citation(row: dict) -> str: title = row.get("title", "") arxiv_id = row.get("doi", row.get("arxiv_id", "")) return f"{title}. arXiv:{arxiv_id}" if arxiv_id else title # ── Database helpers ────────────────────────────────────────────────────────── def _get_or_create_document(cur, row: dict) -> str: """Upsert a document row and return its UUID.""" arxiv_id = row.get("doi", row.get("arxiv_id", str(uuid.uuid4()))) cur.execute( """ INSERT INTO documents (arxiv_id, title, source_url, full_citation, metadata) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (arxiv_id) DO UPDATE SET title = EXCLUDED.title RETURNING id """, ( arxiv_id, row.get("title"), _arxiv_url(arxiv_id), _full_citation(row), psycopg2.extras.Json({"source": row.get("source", "")}), ), ) return cur.fetchone()[0] def _get_role_id(cur, role_name: str) -> str: cur.execute("SELECT id FROM roles WHERE name = %s", (role_name,)) result = cur.fetchone() if result is None: raise ValueError(f"Role '{role_name}' not found. Run init_db.py first.") return result[0] def _grant_access(cur, role_id: str, doc_ids: list[str]) -> None: psycopg2.extras.execute_values( cur, "INSERT INTO role_document_access (role_id, document_id) VALUES %s ON CONFLICT DO NOTHING", [(role_id, doc_id) for doc_id in doc_ids], ) def _get_existing_chunk_indices(cur, doc_id: str) -> set: """Return the set of chunk_index values already stored for this document.""" cur.execute("SELECT chunk_index FROM chunks WHERE document_id = %s", (doc_id,)) return {row[0] for row in cur.fetchall()} def _bulk_insert_chunks(cur, rows: list[tuple]) -> None: """rows: (doc_id, content, embedding_str, chunk_index, metadata_json)""" psycopg2.extras.execute_values( cur, """ INSERT INTO chunks (document_id, content, embedding, chunk_index, metadata) VALUES %s ON CONFLICT (document_id, chunk_index) DO NOTHING """, rows, template="(%s, %s, %s::vector, %s, %s)", ) # ── Main ingestion loop ──────────────────────────────────────────────────────── def ingest(limit: int | None, batch_size: int, role_name: str) -> None: settings = get_settings() embedder = SentenceTransformer(EMBED_MODEL) conn = psycopg2.connect(settings.database_sync_url) conn.autocommit = False with conn.cursor() as cur: role_id = _get_role_id(cur, role_name) logger.info( "Starting ingestion — dataset=%s limit=%s role=%s", DATASET_NAME, limit or "ALL", role_name, ) pending_chunks: list[tuple] = [] doc_ids_seen: set[str] = set() doc_chunk_cache: dict[str, set] = {} # doc_id -> set of already-stored chunk_index total_chunks = 0 skipped_chunks = 0 failed_chunks = 0 texts_buf: list[str] = [] meta_buf: list[dict] = [] def _flush(force: bool = False): nonlocal total_chunks, failed_chunks if not texts_buf: return if not force and len(texts_buf) < EMBED_BATCH: return embeddings = embedder.encode(texts_buf, batch_size=EMBED_BATCH, show_progress_bar=False) for emb, meta in zip(embeddings, meta_buf): pending_chunks.append(( meta["doc_id"], meta["content"], str(emb.tolist()), meta["chunk_index"], psycopg2.extras.Json(meta.get("metadata", {})), )) texts_buf.clear() meta_buf.clear() if force or len(pending_chunks) >= batch_size: try: with conn.cursor() as cur: _bulk_insert_chunks(cur, pending_chunks) if doc_ids_seen: _grant_access(cur, role_id, list(doc_ids_seen)) conn.commit() total_chunks += len(pending_chunks) logger.info("Inserted %d chunks (total: %d)", len(pending_chunks), total_chunks) except Exception as batch_exc: conn.rollback() logger.warning("Batch insert failed (%s) — retrying row-by-row", batch_exc) inserted = 0 for chunk_row in pending_chunks: try: with conn.cursor() as cur: _bulk_insert_chunks(cur, [chunk_row]) conn.commit() inserted += 1 except Exception as row_exc: conn.rollback() logger.error( "Skipping chunk (doc=%s idx=%s): %s", chunk_row[0], chunk_row[3], row_exc, ) failed_chunks += 1 total_chunks += inserted logger.info("Row-by-row retry: %d inserted, %d failed", inserted, failed_chunks) finally: pending_chunks.clear() doc_ids_seen.clear() try: for row in tqdm(_iter_dataset(limit), desc="Ingesting", unit="chunk"): with conn.cursor() as cur: doc_id = _get_or_create_document(cur, row) if doc_id not in doc_chunk_cache: doc_chunk_cache[doc_id] = _get_existing_chunk_indices(cur, doc_id) conn.commit() doc_ids_seen.add(doc_id) chunk_index = row.get("chunk-id", row.get("chunk_index", 0)) content = row.get("chunk", row.get("text", "")).replace("\x00", "") if not content.strip(): continue if chunk_index in doc_chunk_cache[doc_id]: skipped_chunks += 1 continue texts_buf.append(content) meta_buf.append({ "doc_id": doc_id, "content": content, "chunk_index": chunk_index, "metadata": {"source": row.get("source", "")}, }) _flush() _flush(force=True) # flush remainder finally: conn.close() logger.info( "Ingestion complete. Inserted: %d Skipped (already exist): %d Failed: %d", total_chunks, skipped_chunks, failed_chunks, ) # ── CLI ─────────────────────────────────────────────────────────────────────── def _parse_args(): parser = argparse.ArgumentParser(description="Ingest ArXiv chunks into pgvector") parser.add_argument("--limit", type=int, default=None, help="Max chunks to ingest") parser.add_argument("--batch-size", type=int, default=DB_BATCH, help="DB insert batch size") parser.add_argument("--role", default="researcher", help="Role to grant document access") return parser.parse_args() if __name__ == "__main__": args = _parse_args() ingest(limit=args.limit, batch_size=args.batch_size, role_name=args.role)