| """PDF ingestion pipeline: chunk, embed with E5, upsert to Pinecone.""" |
|
|
| import argparse |
| import os |
| import uuid |
|
|
| import nltk |
| import torch |
| from pinecone import Pinecone, ServerlessSpec |
| from PyPDF2 import PdfReader |
| from sentence_transformers import SentenceTransformer |
| from tqdm import tqdm |
|
|
| nltk.download("punkt", quiet=True) |
| nltk.download("punkt_tab", quiet=True) |
|
|
|
|
| def load_settings(): |
| from pathlib import Path |
|
|
| from dotenv import load_dotenv |
|
|
| root = Path(__file__).resolve().parents[2] |
| load_dotenv(root / ".env") |
| return { |
| "api_key": os.getenv("PINECONE_API_KEY", ""), |
| "index_name": os.getenv("PINECONE_INDEX_NAME", "e5"), |
| "embedding_model": os.getenv("EMBEDDING_MODEL", "intfloat/multilingual-e5-large-instruct"), |
| } |
|
|
|
|
| def extract_text_from_pdf(pdf_path: str) -> str: |
| text = "" |
| with open(pdf_path, "rb") as file: |
| reader = PdfReader(file) |
| for page in reader.pages: |
| page_text = page.extract_text() or "" |
| text += page_text + "\n" |
| return text.strip() |
|
|
|
|
| def chunk_fixed_length(text: str, chunk_size: int = 1000, chunk_overlap: int = 100) -> list[str]: |
| if len(text) <= chunk_size: |
| return [text] |
|
|
| chunks: list[str] = [] |
| start = 0 |
| while start < len(text): |
| end = start + chunk_size |
| if end >= len(text): |
| chunks.append(text[start:]) |
| break |
|
|
| while end > start + chunk_size - 100 and end < len(text) and text[end] != " ": |
| end -= 1 |
| if end == start + chunk_size - 100: |
| end = start + chunk_size |
|
|
| chunks.append(text[start:end]) |
| start = end - chunk_overlap |
|
|
| return chunks |
|
|
|
|
| def setup_pinecone_index(api_key: str, index_name: str, dimension: int): |
| pc = Pinecone(api_key=api_key) |
| if not pc.has_index(name=index_name): |
| pc.create_index( |
| name=index_name, |
| dimension=dimension, |
| metric="dotproduct", |
| spec=ServerlessSpec(cloud="aws", region="us-east-1"), |
| ) |
| print(f"Created Pinecone index: {index_name}") |
| else: |
| print(f"Using existing Pinecone index: {index_name}") |
|
|
| return pc.Index(name=index_name) |
|
|
|
|
| def ingest_pdfs(pdf_dir: str, index, model: SentenceTransformer, chunk_size: int = 1000) -> None: |
| pdf_files = [f for f in os.listdir(pdf_dir) if f.lower().endswith(".pdf")] |
| if not pdf_files: |
| print(f"No PDF files found in {pdf_dir}") |
| return |
|
|
| batch: list[dict] = [] |
| batch_size = 100 |
|
|
| for pdf_file in pdf_files: |
| pdf_path = os.path.join(pdf_dir, pdf_file) |
| print(f"Processing {pdf_file}...") |
| text = extract_text_from_pdf(pdf_path) |
| if not text: |
| print(f"Skipping {pdf_file} - no text extracted") |
| continue |
|
|
| chunks = chunk_fixed_length(text, chunk_size=chunk_size) |
| print(f"Split {pdf_file} into {len(chunks)} chunks") |
|
|
| for i, chunk in enumerate(tqdm(chunks, desc=f"Embedding {pdf_file}")): |
| embedding = model.encode(f"passage: {chunk}").tolist() |
| batch.append( |
| { |
| "id": f"{pdf_file.replace('.pdf', '')}_{i}_{uuid.uuid4()}", |
| "values": embedding, |
| "metadata": { |
| "filename": pdf_file, |
| "chunk_index": i, |
| "text": chunk[:1000], |
| "full_text": chunk, |
| }, |
| } |
| ) |
|
|
| if len(batch) >= batch_size: |
| index.upsert(vectors=batch) |
| batch = [] |
|
|
| if batch: |
| index.upsert(vectors=batch) |
|
|
| stats = index.describe_index_stats() |
| print(f"Total vectors in index: {stats.get('total_vector_count', 0)}") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Ingest PDFs into Pinecone for RAG") |
| parser.add_argument("--pdf_dir", type=str, required=True, help="Directory with PDF files") |
| parser.add_argument("--chunk_size", type=int, default=1000) |
| args = parser.parse_args() |
|
|
| cfg = load_settings() |
| if not cfg["api_key"]: |
| raise SystemExit("PINECONE_API_KEY is not set") |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"Using device: {device}") |
|
|
| model = SentenceTransformer(cfg["embedding_model"], device=device) |
| dimension = model.get_sentence_embedding_dimension() |
| index = setup_pinecone_index(cfg["api_key"], cfg["index_name"], dimension) |
|
|
| ingest_pdfs(args.pdf_dir, index, model, chunk_size=args.chunk_size) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
