commit 96f9d9dc81eefc10fc3719d5d1014f4368cd6e22 Author: xiaoxue Date: Mon May 25 00:32:30 2026 +0800 同步完整源码 - 2026-05-25 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..fed2541 --- /dev/null +++ b/LICENSE @@ -0,0 +1,23 @@ +MIT License + +Copyright (c) 2025 Ofir Ben Shabat. +Author: Ofir Ben Shabat +Contact: ofirbshofir@gmail.com + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..ace6dd8 --- /dev/null +++ b/README.md @@ -0,0 +1,448 @@ +# 🧠 AI-Embedder-Engine + +## Table of Contents + +1. [Introduction](#introduction) +2. [Why Use This Project](#why-use-this-project) +3. [Key Features](#key-features) +4. [Quick Start](#quick-start) + - [Prerequisites](#-prerequisites) + - [Run Instructions](#-run-instructions) + - [Run on a Single PDF](#run-on-a-single-pdf) + - [Batch Processing Multiple PDFs](#batch-processing-multiple-pdfs) + - [FAISS Index Builder](#faiss-index-builder) +5. [Configuration](#configuration) +6. [How the Pipeline Works](#how-the-pipeline-works) +7. [Data Formats](#data-formats) + - [Manifest Schema](#-manifest-schema-jsonl) + - [Parquet Schema](#-parquet-schema) +8. [Project Structure](#-project-structure) +9. [Logging & Monitoring](#-logging--monitoring) +10. [Use Cases](#-use-cases) +11. [Performance Tips](#-performance-tips) +12. [Vector Indexing with FAISS](#vector-indexing-with-faiss) +13. [Disclaimer](#disclaimer) +14. [License](#-license) + + +--- + +## Introduction + +This project provides a **standalone embedding pipeline** that transforms large PDF documents into high-quality vector embeddings, stored in **Parquet** format with a **JSONL manifest** for indexing and inspection. +Originally developed for processing **official medical material**, the pipeline is fully domain-agnostic and can be applied to **medical, legal, technical, or academic** content. +By combining **efficient text extraction**, **adaptive chunking**, and **state-of-the-art embeddings**, it enables seamless integration into RAG workflows and advanced vector search systems. +With built-in support for **FAISS indexing (open source)**, the pipeline offers both scalability and speed, making it suitable for enterprise-grade information retrieval. + +--- + +## Why Use This Project + +- **Domain-Agnostic** – adaptable to medical, legal, academic, and enterprise documents. +- **End-to-End Workflow** – from PDF extraction → chunking → embeddings → Parquet storage → FAISS indexing. +- **Production-Ready Formats** – embeddings stored in **Parquet** + **JSONL manifest**. +- **FAISS-Ready** – supports **FAISS** (open source, developed by Meta AI). +- **Optimized for RAG** – plug directly into LLM-based applications. +- **Configurable by Design** – swap embedding models, tune chunking, or replace FAISS. +- **Open Source & Extensible** – modular Python components, easy to integrate. + +--- + +## Key Features + +- **PDF-to-Vector Pipeline** – complete workflow from raw PDFs to clean text, chunking, embeddings, Parquet storage, and FAISS indexing. +- **Scalable Storage** – embeddings are saved in **Parquet** format for efficient compression, schema enforcement, and compatibility with big-data tools. +- **Manifest Tracking** – every processed document is logged in a **JSONL manifest**, enabling version control and lightweight inspection. +- **Adaptive Chunking** – configurable `CHUNK_SIZE` and `OVERLAP` ensure optimal balance between context preservation and embedding efficiency. +- **Batch Embedding with Retries** – supports batching with automatic retries for stability and cost-efficient API calls. +- **FAISS Integration** – built-in indexing with **FAISS** for high-performance semantic search at scale. +- **Configurable & Extensible** – embedding models, chunking, batching, and indexing methods are all configurable to fit different project needs. +- **Logging & Cost Estimation** – detailed run-time logs with token and cost estimates before embedding execution. + +--- + +## Quick Start + +### 📦 Prerequisites + +1. **Create virtual environment & activate** + (Python 3.10+ recommended) + ```bash + python -m venv .venv && source .venv/bin/activate + # Windows: + .venv\Scripts\activate + ``` + +2. **Install dependencies** + ```bash + pip install -r requirements.txt + ``` + +3. **Set your OpenAI API key** + Default environment variable name is: EMBEDDING_API_KEY + ```bash + export EMBEDDING_API_KEY=YOUR_OPENAI_KEY + # Windows: + set EMBEDDING_API_KEY=YOUR_OPENAI_KEY + ``` + > You can change this name in `config.py` → `OPENAI_API_KEY_ENV`. + +### 📦 Run Instructions + +4. **Run on a single PDF** + ```bash + python build_embeddings.py --pdf "/path/to/document.pdf" --doc-id "DOC-ID-Example" + ``` + **Outputs** + - `data/embeddings/DOC-ID-Example.parquet` → **canonical archive** + - `data/manifest/manifest.jsonl` → updated manifest entry (one per document) + + +5. **Batch Processing Multiple PDFs** + Edit `prepare_pdf.py` and fill the `CHAPTERS` list with `(doc_id, pdf_path)` pairs: + ```bash + python prepare_pdf.py + ``` + **Outputs** + - all embeddings saved under `data/embeddings/` + - `data/manifest/manifest.jsonl` → updated manifest entry (one per document) + + +6. **FAISS Index Builder** + After embeddings are generated and stored in **Parquet** format, you can build a **FAISS index** for efficient semantic search. + **Run on all Parquet files (directory mode)** + ```bash + python faiss_index_builder.py + ``` + **Run on a single Parquet file** + ```bash + python faiss_index_builder.py --file-name "DOC-ID-Example" + ``` + **Outputs** + - `data/faiss/DOC-ID-Example.index` → FAISS index file + - `data/faiss/DOC-ID-Example_meta.json` → metadata (document IDs, schema, and parameters) + ⚡ Tip: You can configure index type (`flat` / `hnsw`) and metric (`ip` / `l2`) inside `faiss_index_builder.py` at the top section. + + +--- +## Configuration + +The pipeline is highly configurable. +Most settings are defined in **`config.py`** at the project root. + +### Embedding & Chunking +- `MODEL_NAME` – which embedding model to use (default: `text-embedding-3-large`) +- `EMBEDDING_DIM` – vector dimension expected by the chosen model +- `CHUNK_SIZE` – number of characters per chunk +- `OVERLAP` – number of overlapping characters between chunks +- `BATCH_SIZE` – number of API requests per batch +- `SPLIT_LENGTH` – optional parameter for future advanced splitting strategies + +### Storage +- `DATA_DIR` – base data folder (`data/`) +- `EMB_DIR` – where Parquet embeddings are saved +- `MANIFEST_DIR` – where the JSONL manifest is saved +- `MANIFEST_PATH` – path to the manifest file + +### API Key +- `OPENAI_API_KEY_ENV` – environment variable name for your embedding API key + (default: `EMBEDDING_API_KEY`) + +### Cost Estimation +- `AVG_CHARS_PER_TOKEN` – heuristic used when tokenization is not available +- `PRICE_PER_1K_TOKENS` – used to estimate cost before running embeddings + +--- + +> ⚙️ To customize behavior, edit `config.py` directly or override values in your own wrapper script. + +--- + +## How the Pipeline Works + +The pipeline is organized as a clear sequence of steps, transforming raw PDF documents into vector embeddings with supporting metadata: + + ┌───────────────┐ + │ PDF(s) │ + └───────┬───────┘ + │ Extract text + ▼ + ┌───────────────────┐ + │ Raw Text (per │ + │ document) │ + └────────┬──────────┘ + │ Clean + Chunk (CHUNK_SIZE, OVERLAP) + ▼ + ┌───────────────────┐ + │ Chunks (list) │ + └────────┬──────────┘ + │ Embed (MODEL_NAME, BATCH_SIZE, retries) + ▼ + ┌───────────────────┐ + │ Embeddings (N×D) │ + └────────┬──────────┘ + │ Write vectors → Parquet (per doc) + │ Append record → manifest.jsonl + ▼ + ┌──────────────────────────────┐ + │ data/embeddings/*.parquet │ + │ data/manifest/manifest.jsonl │ + └──────────────┬───────────────┘ + │ (optional) + │ Build FAISS index (flat / hnsw, ip / l2) + ▼ + ┌────────────────────────┐ + │ data/faiss/*.index │ + │ data/faiss/*_meta.json│ + └────────────┬───────────┘ + │ (optional) + ▼ + RAG / Semantic Search + + +1. **PDF → Raw Text** + - Extract text from input PDF files. + - Handles multi-page documents with cleaning to remove artifacts. + +2. **Text → Chunks** + - Apply configurable **chunking** (`CHUNK_SIZE`, `OVERLAP`). + - Each chunk preserves enough context while keeping within embedding model limits. + +3. **Chunks → Embeddings** + - Send each chunk to the embedding API (`MODEL_NAME`). + - Supports batching (`BATCH_SIZE`) and automatic retries for stability. + - Estimated token usage and cost are logged before execution. + +4. **Embeddings → Parquet Archive** + - Results are written into **compressed Parquet files** (`data/embeddings/`). + - Each row includes: `doc_id`, chunk text, embedding vector, metadata (hash, model, created_at). + +5. **Manifest Tracking** + - For every document processed, a lightweight record is appended to **`data/manifest/manifest.jsonl`**. + - The manifest allows quick inspection and version control without storing vectors. + +6. **Optional: FAISS Indexing** + - Parquet embeddings can be transformed into a **FAISS index** (`data/faiss/`). + - Supports both *exact* (`flat`) and *approximate* (`hnsw`) indexes. + - Metadata JSON accompanies each index for reproducibility and inspection. + +--- + +> 🧩 This modular flow allows you to plug in custom extractors, chunkers, embedding models, or indexing backends with minimal changes. + +--- +## Data Formats + +### 📑 Manifest Schema (JSONL) + +Each processed document is logged as a **single line** inside `data/manifest/manifest.jsonl`. +The format is **JSON Lines (JSONL)** – one JSON object per line, Git-friendly and easy to inspect. + +**Main fields**: +- `doc_id` *(string)* – unique document ID (e.g., `DOC-001` or `H21-CH01`). +- `source_path` *(string)* – absolute path to the source PDF. +- `num_chunks` *(int)* – number of chunks created after splitting and embedding. +- `model_name` *(string)* – embedding model used (e.g., `text-embedding-3-large`). +- `dim` *(int)* – vector dimension (e.g., `3072`). +- `chunk_size` *(int)* – configured chunk size. +- `overlap` *(int)* – overlap between consecutive chunks. +- `created_at` *(string, ISO8601)* – UTC timestamp when embeddings were created. +- `parquet_path` *(string)* – path to the generated Parquet file. + +📌 **Example**: +```json +{ + "doc_id": "DOC-001", + "source_path": "/abs/path/to/document.pdf", + "num_chunks": 145, + "model_name": "text-embedding-3-large", + "dim": 3072, + "chunk_size": 700, + "overlap": 200, + "created_at": "2025-08-18T10:22:33Z", + "parquet_path": "/abs/path/to/data/embeddings/DOC-001.parquet" +} +``` + +### 📑 Parquet Schema + +The **Parquet file** is the canonical archive containing embeddings. +Each row represents a **single text chunk** with its associated vector. + +**Main fields**: +- `doc_id` *(string)* – document ID. +- `source_path` *(string)* – original PDF path. +- `page_start` *(int | null)* – starting page (optional). +- `page_end` *(int | null)* – ending page (optional). +- `chunk_id` *(int)* – sequential index of the chunk. +- `start_char` *(int)* – start character offset in the raw text. +- `end_char` *(int)* – end character offset. +- `text` *(string)* – cleaned text of the chunk. +- `embedding` *(array[float32])* – embedding vector of size `dim`. +- `model_name` *(string)* – embedding model used. +- `dim` *(int)* – embedding dimension. +- `chunk_size` *(int)* – chunk size used. +- `overlap` *(int)* – overlap used. +- `hash` *(string, sha1)* – SHA-1 checksum of the chunk text. +- `created_at` *(string, ISO8601)* – creation timestamp. + +📌 **Example (JSON view of a Parquet row)**: +```json +{ + "doc_id": "DOC-001", + "source_path": "/abs/path/to/document.pdf", + "page_start": null, + "page_end": null, + "chunk_id": 42, + "start_char": 12000, + "end_char": 12700, + "text": "Clinical manifestations of the disease include ...", + "embedding": [0.0134, -0.0245, ..., 0.0098], + "model_name": "text-embedding-3-large", + "dim": 3072, + "chunk_size": 700, + "overlap": 200, + "hash": "a94a8fe5ccb19b991c4c0873d153e987a82fbbd3", + "created_at": "2025-08-18T10:22:33Z" +} +``` + +--- + +## 📂 Project Structure + +``` +AI-Embedder-Engine/ +│── data/ +│ ├── input_pdfs/ # Raw PDFs (source) +│ ├── embeddings/ # Output Parquet files +│ ├── manifest/ # JSONL manifest logs +│ ├── faiss/ # FAISS index + metadata +│ +│── prepare_pdf.py # PDF → Parquet pipeline +│── build_embeddings.py # Chunking + Embedding +│── faiss_index_builder.py # FAISS index builder +│── config.py # Config (chunk size, model, paths...) +│── requirements.txt # Dependencies +│── README.md # This file +``` + +--- + +## 📜 Logging & Monitoring + +The pipeline provides detailed logs during execution to help track progress, debug issues, and estimate costs. + +### What gets logged? +- **Processing steps** – PDF extraction, chunking, embedding, Parquet writing, FAISS indexing. +- **Token usage & cost estimates** – based on configured pricing (`PRICE_PER_1K_TOKENS`). +- **Retries & errors** – failed API calls are automatically retried with backoff, all attempts are logged. +- **Runtime stats** – number of chunks processed, time taken, embedding rate (chunks/sec). + +### Example Output +```text +[2025-08-18 14:32:11] INFO Starting embedding for DOC-001 (145 chunks) +[2025-08-18 14:32:15] INFO Estimated tokens: 22,800 | Estimated cost: $0.06 +[2025-08-18 14:32:25] INFO Batch 1/15 processed (10 chunks) in 3.2s +[2025-08-18 14:35:40] INFO Embedding complete: DOC-001.parquet written +[2025-08-18 14:35:40] INFO Manifest updated → data/manifest/manifest.jsonl +``` + +--- + +## 💡 Use Cases + +The **AI-Embedder-Engine** can be applied across many domains where **semantic search** and **RAG (Retrieval-Augmented Generation)** are required. + +### 📚 Academic & Research +- Index large collections of **scientific papers** or **theses**. +- Enable semantic queries across literature for faster discovery. + +### ⚖️ Legal Documents +- Process contracts, case law, and compliance reports. +- Build a **vector search engine** for clause retrieval and precedent matching. + +### 🏥 Medical Knowledge +- Originally built for **medical textbooks** and guidelines. +- Adaptable to clinical decision support and biomedical research. + +### 🏢 Enterprise Knowledge Base +- Convert internal **knowledge docs, manuals, and reports** into embeddings. +- Provide employees with an **AI-powered assistant** for internal search. + +### 📰 Media & Content +- Index large archives of **news articles, blogs, or reports**. +- Enhance recommendation engines with **semantic similarity search**. + +### 🔬 Technical Documentation +- Parse **developer docs, API references, and manuals**. +- Power **chatbots and assistants** that can answer technical questions directly from indexed content. + +--- + +## ⚡ Performance Tips + +Optimizing the pipeline for **speed, cost, and scalability** is critical for large datasets. + +### 🔹 Chunking Strategy +- Use **larger chunk sizes** for narrative text (e.g., books, articles). +- Use **smaller chunks** for dense content (e.g., legal, medical, or technical references). +- Always keep some **overlap** (e.g., 150–200 characters) to avoid losing context between chunks. + +### 🔹 Embedding Batches +- Increase the **batch size** to reduce API roundtrips (default is 10). +- If hitting API limits, reduce batch size to avoid errors. + +### 🔹 Cost Estimation +- Use the built-in **token estimator** to preview approximate costs before embedding. +- Consider filtering out **irrelevant pages** (like empty or index pages) to save tokens. + +### 🔹 FAISS Indexing +- For small datasets (<100K vectors), use **Flat (exact)** indexes. +- For large datasets (>1M vectors), switch to **HNSW** for faster approximate search. +- Normalize embeddings for **cosine similarity** (done automatically with `metric="ip"`). + +--- + +## Vector Indexing with FAISS + +The project includes built-in support for **FAISS** (Facebook AI Similarity Search), a high-performance vector index library. +This enables **semantic search** over millions of embeddings with either **exact (Flat)** or **approximate (HNSW)** methods. + +### 🔹 Why FAISS? +- **Scalability** – Handles datasets from thousands to billions of vectors. +- **Speed** – Provides millisecond query times even on large indexes. +- **Flexibility** – Supports multiple distance metrics (cosine/IP, L2, etc.). +- **Compatibility** – Works directly with the Parquet embeddings produced by this project. + +### 🔹 Index Types +- **Flat (Exact Search)** + - Stores all vectors directly. + - Best for smaller datasets (<100K vectors). + - Guarantees exact nearest neighbors. + +- **HNSW (Approximate Search)** + - Graph-based index structure for fast lookups. + - Best for large datasets (>1M vectors). + - Sacrifices a small amount of accuracy for massive speed gains. + +--- + +## Disclaimer +This project is not an official product of OpenAI or Meta. +It uses **OpenAI models** and the **FAISS library** under their respective licenses. +No ownership of these models or libraries is claimed. +The system is provided *"as is"*, without warranties or guarantees. +Usage is subject to the [OpenAI Terms of Use](https://openai.com/policies/terms-of-use) and the [FAISS MIT License](https://github.com/facebookresearch/faiss/blob/main/LICENSE). +For health-related applications, additional caution and compliance with medical AI guidelines is required. + +--- + +## 📄 License + +This project is licensed under the **MIT License**. +See the [LICENSE](./LICENSE) file for details. + +© 2025 Ofir Ben Shabat • Contact: ofirbshofir@gmail.com + +--- \ No newline at end of file diff --git a/build_embeddings.py b/build_embeddings.py new file mode 100644 index 0000000..443ef66 --- /dev/null +++ b/build_embeddings.py @@ -0,0 +1,123 @@ +import os +import argparse +import hashlib +from datetime import datetime, timezone +from typing import List, Dict, Any + +from config import ( + MODEL_NAME, EMBEDDING_DIM, CHUNK_SIZE, OVERLAP, BATCH_SIZE, + EMB_DIR, MANIFEST_PATH, AVG_CHARS_PER_TOKEN, PRICE_PER_1K_TOKENS +) +from pdf_text import extract_text_from_pdf +from chunker import basic_clean, chunk_text +from embedder import get_client, embed_texts +from writer import write_parquet, append_manifest + + +def _estimate_tokens(texts: List[str]) -> tuple[int, str]: + """Estimate total tokens for a list of texts. Prefer tiktoken; fallback to char heuristic.""" + try: + import tiktoken # type: ignore + enc = tiktoken.get_encoding("cl100k_base") + total = sum(len(enc.encode(t)) for t in texts) + return total, "tiktoken" + except Exception: + total_chars = sum(len(t) for t in texts) + approx = int(total_chars / AVG_CHARS_PER_TOKEN) + return approx, "chars-heuristic" + + +def sha1(text: str) -> str: + return hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest() + + +def build_for_pdf(pdf_path: str, doc_id: str) -> str: + if not os.path.isfile(pdf_path): + raise FileNotFoundError(f"PDF not found: {pdf_path}") + + print(f"[INFO] Reading PDF: {pdf_path}") + raw_text, total_pages = extract_text_from_pdf(pdf_path) + cleaned = basic_clean(raw_text) + + print(f"[INFO] Chunking (size={CHUNK_SIZE}, overlap={OVERLAP})…") + chunks = chunk_text(cleaned, CHUNK_SIZE, OVERLAP) + if not chunks: + raise RuntimeError("No chunks produced; check PDF text extraction or chunking settings.") + + texts = [c[2] for c in chunks] + positions = [(c[0], c[1]) for c in chunks] + + # --- FACT LOGS --- + print(f"[INFO] PAGES: {total_pages}") + print(f"[INFO] CHUNKS: {len(texts)}") + + # --- ESTIMATION LOGS --- + est_tokens, method = _estimate_tokens(texts) + est_cost = (est_tokens / 1000.0) * PRICE_PER_1K_TOKENS + print(f"[INFO] ESTIMATE: tokens≈{est_tokens:,} (via {method}) | cost≈${est_cost:.3f} | " + f"avg_chars_per_token={AVG_CHARS_PER_TOKEN}") + + print(f"[INFO] Embedding {len(texts)} chunks with model={MODEL_NAME} (batch={BATCH_SIZE})…") + client = get_client() + vectors = embed_texts(texts, client=client, model=MODEL_NAME, batch_size=BATCH_SIZE) + + if len(vectors) != len(texts): + raise RuntimeError("Mismatch between vectors and texts length.") + + # ISO 8601 UTC without microseconds + created_at = datetime.now(timezone.utc).replace(microsecond=0).isoformat() + source_path = os.path.abspath(pdf_path) + + # Prepare Parquet rows + rows: List[Dict[str, Any]] = [] + for idx, ((start, end), txt, vec) in enumerate(zip(positions, texts, vectors)): + rows.append({ + "doc_id": doc_id, + "source_path": source_path, + "page_start": None, # Optional: could be inferred with a more advanced pager + "page_end": None, + "chunk_id": idx, + "start_char": start, + "end_char": end, + "text": txt, + "embedding": vec, + "model_name": MODEL_NAME, + "dim": EMBEDDING_DIM, + "chunk_size": CHUNK_SIZE, + "overlap": OVERLAP, + "hash": sha1(txt), + "created_at": created_at + }) + + out_path = os.path.join(EMB_DIR, f"{doc_id}.parquet") + write_parquet(rows, out_path) + + manifest_entry = { + "doc_id": doc_id, + "source_path": source_path, + "num_chunks": len(rows), + "model_name": MODEL_NAME, + "dim": EMBEDDING_DIM, + "chunk_size": CHUNK_SIZE, + "overlap": OVERLAP, + "created_at": created_at, + "parquet_path": out_path + } + append_manifest([manifest_entry]) + + print(f"[INFO] Manifest appended at: {MANIFEST_PATH}") + return out_path + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Build embeddings Parquet for a single PDF.") + parser.add_argument("--pdf", required=True, help="Path to the PDF file.") + parser.add_argument("--doc-id", required=True, help="Document ID, e.g., DOC-001-Employee-Handbook.") + args = parser.parse_args() + + try: + path = build_for_pdf(args.pdf, args.doc_id) + print(f"[DONE] Embeddings built: {path}") + except Exception as e: + print(f"[ERROR] {e}") + raise \ No newline at end of file diff --git a/chunker.py b/chunker.py new file mode 100644 index 0000000..a9c9904 --- /dev/null +++ b/chunker.py @@ -0,0 +1,28 @@ +from typing import List, Tuple + +def basic_clean(text: str) -> str: + """Minimal cleaning only: normalize broken lines and spaces, do NOT remove headers/footers.""" + # Join obvious hyphenated line-breaks and normalize spaces + text = text.replace('\r', '') + # Collapse multiple spaces but keep newlines + lines = [l.strip() for l in text.split('\n')] + return "\n".join(lines) + +def chunk_text(text: str, chunk_size: int, overlap: int) -> List[Tuple[int, int, str]]: + """Split text into overlapping chunks by character count. + Returns list of tuples: (start_char, end_char, chunk_text).""" + if chunk_size <= 0: + raise ValueError("chunk_size must be > 0") + if overlap < 0 or overlap >= chunk_size: + raise ValueError("overlap must be >= 0 and < chunk_size") + + chunks: List[Tuple[int, int, str]] = [] + n = len(text) + start = 0 + while start < n: + end = min(start + chunk_size, n) + chunks.append((start, end, text[start:end])) + if end == n: + break + start = end - overlap # overlap ensures context continuity + return chunks diff --git a/config.py b/config.py new file mode 100644 index 0000000..09ad1f3 --- /dev/null +++ b/config.py @@ -0,0 +1,36 @@ +import os +from dotenv import load_dotenv + +load_dotenv() + +# ==== Embedding & Chunking params ==== +MODEL_NAME = "text-embedding-3-large" # default embedding model +EMBEDDING_DIM = 3072 # dimension for the model above +CHUNK_SIZE = 700 +OVERLAP = 200 +BATCH_SIZE = 10 # items per embedding request +MAX_RETRIES = 3 # automatic retries for transient API errors +RETRY_BACKOFF_SEC = 2.0 # backoff between retries (exponential/backoff logic in client) +SPLIT_LENGTH = 128_000 # reserved for future splitting strategies + +# ==== Storage ==== +BASE_DIR = os.path.dirname(__file__) +DATA_DIR = os.path.join(BASE_DIR, "data") +EMB_DIR = os.path.join(DATA_DIR, "embeddings") +MANIFEST_DIR = os.path.join(DATA_DIR, "manifest") +os.makedirs(EMB_DIR, exist_ok=True) +os.makedirs(MANIFEST_DIR, exist_ok=True) + +MANIFEST_PATH = os.path.join(MANIFEST_DIR, "manifest.jsonl") + +# ==== Environment ==== +# Use a neutral, public-friendly env var name. Keep legacy key for backward compatibility if needed. +OPENAI_API_KEY_ENV = "EMBEDDING_API_KEY" + +# ==== Cost estimation (informative only) ==== +AVG_CHARS_PER_TOKEN = 4.0 # rough average; depends on language/script +PRICE_PER_1K_TOKENS = 0.00013 # $ per 1K tokens for text-embedding-3-large (update if provider changes) + +# ==== Canonical archive note ==== +# Vectors are stored in Parquet (ZSTD compressed) as the canonical archive. +# The JSONL manifest holds metadata only (no vectors) for readability and version control. diff --git a/embedder.py b/embedder.py new file mode 100644 index 0000000..19771a6 --- /dev/null +++ b/embedder.py @@ -0,0 +1,34 @@ +import os +from typing import List, Dict, Any +import time +import numpy as np +from openai import OpenAI +from config import MODEL_NAME, OPENAI_API_KEY_ENV, BATCH_SIZE + +def get_client() -> OpenAI: + api_key = os.getenv(OPENAI_API_KEY_ENV) + if not api_key: + raise RuntimeError(f"Missing OpenAI API key env var: {OPENAI_API_KEY_ENV}") + return OpenAI(api_key=api_key) + +def embed_texts(texts: List[str], client: OpenAI = None, model: str = MODEL_NAME, batch_size: int = BATCH_SIZE) -> List[List[float]]: + """Embed a list of texts in batches. Returns list of vectors (list[float]).""" + if client is None: + client = get_client() + + vectors: List[List[float]] = [] + total = len(texts) + for i in range(0, total, batch_size): + batch = texts[i:i+batch_size] + # Retry loop (simple backoff); you can replace with tenacity if desired. + for attempt in range(5): + try: + resp = client.embeddings.create(model=model, input=batch) + for item in resp.data: + vectors.append(item.embedding) + break + except Exception as e: + wait = 2 ** attempt + print(f"[WARN] Embedding batch {i//batch_size+1}/{(total+batch_size-1)//batch_size} failed: {e}. Retrying in {wait}s...") + time.sleep(wait) + return vectors \ No newline at end of file diff --git a/faiss_index_builder.py b/faiss_index_builder.py new file mode 100644 index 0000000..f88d5a4 --- /dev/null +++ b/faiss_index_builder.py @@ -0,0 +1,325 @@ +from __future__ import annotations +import json +import time +import logging +from pathlib import Path +from typing import List, Tuple + +import faiss +import numpy as np +import pandas as pd + +# ----------------------------------------------------------- +# User CONFIG – edit these values only +# ----------------------------------------------------------- + +# Base project path (relative to repo root) +BASE_DIR = Path(__file__).resolve().parent + +# Input folder: Parquet embeddings +PATH_TO_EMBEDDED = BASE_DIR / "data" / "embeddings" + +# Output folder: FAISS index + metadata +PATH_TO_FAISS = BASE_DIR / "data" / "faiss" + +# Dataset version tag (example) +DATASET_VERSION = "v1.0.0" + +# Index settings +INDEX_TYPE = "flat" # "flat" (exact) or "hnsw" (ANN) +METRIC = "ip" # "ip" (cosine-like, with normalization) or "l2" +HNSW_M = 32 # used only if INDEX_TYPE="hnsw" +HNSW_EFC = 200 # used only if INDEX_TYPE="hnsw" + +# Meta columns to include if present in Parquet +META_COLS_PREFERRED = [ + "doc_id", "text", "chapter", "section", "page", + "source_path", "hash", "model_name", "dim", "created_at", +] + +# --- Single-file mode (used when RUN_MODE="single") --- +FILE_NAME = "" # without .parquet (leave empty for dir mode) + +# --- Directory mode (used when RUN_MODE="dir") --- +GLOB_PATTERN = "*.parquet" # which files to index +MERGE_INTO_SINGLE = False # False=one index per file; True=one merged index +MERGED_NAME = "merged" # used only if MERGE_INTO_SINGLE=True + +# Which action to run when executing this file directly: +RUN_MODE = "dir" # "single" / "dir" / "none" +# ----------------------------------------------------------- + +# ============ Logging setup ============ +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%H:%M:%S" +) + +# ============ Helpers ============ +def _ensure_dir(path: Path) -> None: + path.mkdir(parents=True, exist_ok=True) + +def _stack_embeddings(df: pd.DataFrame) -> Tuple[np.ndarray, int]: + """ + Ensure 'embedding' column exists and stack to (N, D) float32 matrix. + Returns: (embeddings, dim) + """ + if "embedding" not in df.columns: + raise ValueError("Missing 'embedding' column in the Parquet file.") + + dims = {len(v) for v in df["embedding"].head(100) if isinstance(v, (list, tuple, np.ndarray))} + if len(dims) != 1: + raise ValueError(f"Inconsistent embedding dims in sample: {dims}") + dim = list(dims)[0] + + try: + emb = np.vstack(df["embedding"].values).astype("float32") + except Exception as e: + raise ValueError(f"Failed to stack embeddings to 2D float32 array: {e}") + if emb.ndim != 2 or emb.shape[1] != dim: + raise ValueError(f"Unexpected embeddings shape {emb.shape}, expected (N, {dim})") + return emb, dim + +def _build_faiss_index( + emb: np.ndarray, + *, + metric: str = "ip", # "ip" (cosine-like via L2 normalization) or "l2" + index_type: str = "flat", # "flat" or "hnsw" + hnsw_m: int = 32, + hnsw_efc: int = 200, +) -> faiss.Index: + """ + Build a FAISS index for the given embeddings matrix (N, D). + - metric="ip": we L2-normalize embeddings -> inner-product ~ cosine + - index_type="flat": exact search; "hnsw": fast ANN for large N + """ + d = emb.shape[1] + if metric not in ("ip", "l2"): + raise ValueError("metric must be 'ip' or 'l2'") + + if metric == "ip": + faiss.normalize_L2(emb) # cosine-like scoring + + if index_type == "flat": + index = faiss.IndexFlatIP(d) if metric == "ip" else faiss.IndexFlatL2(d) + elif index_type == "hnsw": + index = faiss.IndexHNSWFlat(d, hnsw_m) + index.hnsw.efConstruction = hnsw_efc + else: + raise ValueError("index_type must be 'flat' or 'hnsw'") + + index.add(emb) + return index + +def _write_artifacts( + index: faiss.Index, + meta_records: List[dict], + out_index_path: Path, + out_meta_path: Path, + meta_header: dict, +) -> Tuple[Path, Path]: + faiss.write_index(index, str(out_index_path)) + with out_meta_path.open("w", encoding="utf-8") as f: + json.dump({**meta_header, "records": meta_records}, f, ensure_ascii=False) + return out_index_path, out_meta_path + + +# ============ Public API ============ + +def build_index_from_parquet_for_single_file( + path_to_embedded: str, + path_to_faiss: str, + file_name: str, + *, + dataset_version: str = "unknown", + index_type: str = "flat", # "flat" or "hnsw" + metric: str = "ip", # "ip" (cosine-like) or "l2" + meta_cols_preferred: List[str] | None = None, +) -> Tuple[str, str]: + """ + Build a FAISS index from a single Parquet file that already contains an 'embedding' column. + """ + t0 = time.time() + meta_cols_preferred = meta_cols_preferred or META_COLS_PREFERRED + + embedded_dir = Path(path_to_embedded) + out_dir = Path(path_to_faiss) + _ensure_dir(out_dir) + + parquet_path = embedded_dir / f"{file_name}.parquet" + if not parquet_path.exists(): + raise FileNotFoundError(f"Parquet not found: {parquet_path}") + + logging.info(f"Loading Parquet: {parquet_path}") + df = pd.read_parquet(parquet_path) + logging.info(f"Rows: {len(df):,}") + + emb, dim = _stack_embeddings(df) + logging.info(f"Embedding dim: {dim}") + + logging.info(f"Building FAISS index (type={index_type}, metric={metric}) ...") + index = _build_faiss_index( + emb, metric=metric, index_type=index_type, hnsw_m=HNSW_M, hnsw_efc=HNSW_EFC + ) + + index_path = out_dir / f"{file_name}.index" + meta_path = out_dir / f"{file_name}_meta.json" + + meta_cols = [c for c in meta_cols_preferred if c in df.columns] + meta_records = df[meta_cols].to_dict(orient="records") + + header = { + "dataset_version": dataset_version, + "index_name": file_name, + "index_type": index_type, + "metric": metric, + "dim": dim, + "rows": len(df), + "meta_cols": meta_cols, + "source_parquet": str(parquet_path), + } + + idx_p, meta_p = _write_artifacts(index, meta_records, index_path, meta_path, header) + dt = time.time() - t0 + logging.info(f"✅ Wrote index: {idx_p}") + logging.info(f"✅ Wrote meta : {meta_p}") + logging.info(f"Done in {dt:.2f}s") + return str(idx_p), str(meta_p) + + +def build_index_from_parquet_for_files( + path_to_embedded: str, + path_to_faiss: str, + *, + dataset_version: str = "unknown", + index_type: str = "flat", + metric: str = "ip", + glob_pattern: str = "*.parquet", + merge_into_single: bool = False, + merged_name: str = "merged", + meta_cols_preferred: List[str] | None = None, +) -> List[Tuple[str, str]]: + """ + Build FAISS indices for all Parquet files under `path_to_embedded`. + """ + meta_cols_preferred = meta_cols_preferred or META_COLS_PREFERRED + + embedded_dir = Path(path_to_embedded) + out_dir = Path(path_to_faiss) + _ensure_dir(out_dir) + + parquet_paths = sorted(embedded_dir.glob(glob_pattern)) + total = len(parquet_paths) + if total == 0: + raise FileNotFoundError(f"No parquet files under: {embedded_dir} (pattern={glob_pattern})") + + results: List[Tuple[str, str]] = [] + + if merge_into_single: + logging.info(f"Merging {total} parquet files into one FAISS index: {merged_name}") + t0 = time.time() + frames = [] + for i, p in enumerate(parquet_paths, start=1): + logging.info(f"[{i}/{total}] Reading: {p.name}") + df = pd.read_parquet(p) + if "embedding" not in df.columns: + raise ValueError(f"Missing 'embedding' column in {p}") + frames.append(df) + + df_all = pd.concat(frames, ignore_index=True) + emb, dim = _stack_embeddings(df_all) + logging.info(f"Total rows: {len(df_all):,} | dim: {dim}") + + logging.info(f"Building FAISS index (type={index_type}, metric={metric}) ...") + index = _build_faiss_index( + emb, metric=metric, index_type=index_type, hnsw_m=HNSW_M, hnsw_efc=HNSW_EFC + ) + + index_path = out_dir / f"{merged_name}.index" + meta_path = out_dir / f"{merged_name}_meta.json" + + meta_cols = [c for c in meta_cols_preferred if c in df_all.columns] + meta_records = df_all[meta_cols].to_dict(orient="records") + header = { + "dataset_version": dataset_version, + "index_name": merged_name, + "index_type": index_type, + "metric": metric, + "dim": dim, + "rows": len(df_all), + "meta_cols": meta_cols, + "sources": [str(p) for p in parquet_paths], + } + + _write_artifacts(index, meta_records, index_path, meta_path, header) + dt = time.time() - t0 + logging.info(f"✅ Wrote merged index: {index_path}") + logging.info(f"✅ Wrote merged meta : {meta_path}") + logging.info(f"Done in {dt:.2f}s") + results.append((str(index_path), str(meta_path))) + return results + + # One FAISS index per Parquet file, with % progress and ETA + start = time.time() + for i, p in enumerate(parquet_paths, start=1): + t_loop = time.time() + pct = (i / total) * 100.0 + logging.info(f"[{i}/{total}] ({pct:5.1f}%) Processing: {p.name}") + + try: + idx_p, meta_p = build_index_from_parquet_for_single_file( + path_to_embedded=str(embedded_dir), + path_to_faiss=str(out_dir), + file_name=p.stem, + dataset_version=dataset_version, + index_type=index_type, + metric=metric, + meta_cols_preferred=meta_cols_preferred, + ) + results.append((idx_p, meta_p)) + except Exception as e: + logging.error(f"⚠️ Skipped {p.name} due to error: {e}") + finally: + # ETA + elapsed_total = time.time() - start + avg_per_file = elapsed_total / i + remaining = total - i + eta_sec = avg_per_file * remaining + logging.info( + f"Elapsed: {elapsed_total:.1f}s | Avg/file: {avg_per_file:.2f}s | ETA: {eta_sec:.1f}s" + ) + logging.info("-" * 60) + + logging.info(f"✅ Completed building FAISS indexes for {len(results)} files " + f"(out of {total}). Output dir: {out_dir}") + return results + + +# ============ Optional direct run using CONFIG ============ + +if __name__ == "__main__": + if RUN_MODE == "single": + build_index_from_parquet_for_single_file( + path_to_embedded=PATH_TO_EMBEDDED, + path_to_faiss=PATH_TO_FAISS, + file_name=FILE_NAME, + dataset_version=DATASET_VERSION, + index_type=INDEX_TYPE, + metric=METRIC, + meta_cols_preferred=META_COLS_PREFERRED, + ) + elif RUN_MODE == "dir": + build_index_from_parquet_for_files( + path_to_embedded=PATH_TO_EMBEDDED, + path_to_faiss=PATH_TO_FAISS, + dataset_version=DATASET_VERSION, + index_type=INDEX_TYPE, + metric=METRIC, + glob_pattern=GLOB_PATTERN, + merge_into_single=MERGE_INTO_SINGLE, + merged_name=MERGED_NAME, + meta_cols_preferred=META_COLS_PREFERRED, + ) + else: + logging.info("Configured RUN_MODE='none'. Edit the CONFIG at the top and set RUN_MODE to 'single' or 'dir'.") diff --git a/pdf_text.py b/pdf_text.py new file mode 100644 index 0000000..9ee3c3e --- /dev/null +++ b/pdf_text.py @@ -0,0 +1,17 @@ +from typing import List, Tuple +from pypdf import PdfReader + +def extract_text_from_pdf(pdf_path: str) -> Tuple[str, int]: + """Extract text and return (text, total_pages).""" + reader = PdfReader(pdf_path) + texts: List[str] = [] + total = len(reader.pages) + for i, page in enumerate(reader.pages, start=1): + print(f"[INFO] Extracting text from page {i}/{total}...") + try: + txt = page.extract_text() or "" + except Exception as e: + print(f"[WARN] Failed to extract page {i}: {e}") + txt = "" + texts.append(txt) + return "\n".join(texts), total \ No newline at end of file diff --git a/prepare_pdf.py b/prepare_pdf.py new file mode 100644 index 0000000..2c54b0a --- /dev/null +++ b/prepare_pdf.py @@ -0,0 +1,67 @@ +# ------------------------------------------------------------ +# prepare_pdf.py +# ------------------------------------------------------------ +# Utility script for running PDF→Parquet conversions +# either for a predefined list of documents or for a single file. +# +# Mode 1 (recommended): +# Fill the CHAPTERS/DOCUMENTS list and run: +# python prepare_pdf.py +# +# Mode 2 (single PDF from CLI): +# python prepare_pdf.py --pdf "/path/to/notebook.pdf" --doc-id "DOC-001-File-Notebook" +# +# Output: Each run produces Parquet embeddings and a JSONL manifest +# +# Requirements: +# Set the same API key environment variable as used in the project: +# export EMBEDDING_API_KEY=... +# ------------------------------------------------------------ + +import argparse +from typing import List, Tuple +from build_embeddings import build_for_pdf + +# ====== List of documents with (doc_id, path) ====== +CHAPTERS: List[Tuple[str, str]] = [ + # Example entries: + # ("DOC-001-Employee-Handbook", r"/path/to/Employee-Handbook.pdf"), + # ("DOC-002-Product-Manual", r"/path/to/Product-Manual.pdf"), + # ("DOC-003-Research-Paper", r"/path/to/Research-Paper.pdf"), + + # Leave empty if running with --pdf and --doc-id arguments only. +] + +def run_single(pdf: str, doc_id: str) -> None: + print(f"[RUN] {doc_id} ← {pdf}") + out = build_for_pdf(pdf, doc_id) + print(f"[OK ] {doc_id} → {out}") + +def run_batch(pairs: List[Tuple[str, str]]) -> None: + if not pairs: + print("[WARN] CHAPTERS list is empty. Fill it or use --pdf/--doc-id parameters.") + return + errors = 0 + for doc_id, pdf in pairs: + try: + run_single(pdf, doc_id) + except Exception as e: + errors += 1 + print(f"[FAIL] {doc_id}: {e}") + print(f"\n[DONE] Run complete. Total: {len(pairs)}, Success: {len(pairs) - errors}, Failures: {errors}") + +def main(): + parser = argparse.ArgumentParser(description="Prepare PDFs → Parquet (single or batch).") + parser.add_argument("--pdf", help="Path to a single PDF file.") + parser.add_argument("--doc-id", help="Document identifier, e.g. DOC-001-Employee-Handbook.") + args = parser.parse_args() + + if args.pdf and args.doc_id: + run_single(args.pdf, args.doc_id) + elif args.pdf or args.doc_id: + parser.error("Both --pdf and --doc-id must be provided together, or neither (use CHAPTERS list).") + else: + run_batch(CHAPTERS) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..9dfa6e4 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +pypdf==5.0.1 +pyarrow==16.1.0 +numpy>=1.26.0 +python-dotenv>=1.0.1 +openai>=1.43.0 diff --git a/writer.py b/writer.py new file mode 100644 index 0000000..44a64fb --- /dev/null +++ b/writer.py @@ -0,0 +1,33 @@ +import os +from typing import List, Dict, Any +from datetime import datetime + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from config import EMB_DIR, MANIFEST_PATH + +def _ensure_float32(vectors: List[List[float]]) -> List[List[float]]: + return [np.asarray(v, dtype=np.float32).tolist() for v in vectors] + +def write_parquet( + rows: List[Dict[str, Any]], + parquet_out_path: str +) -> None: + """Write rows with schema to Parquet (ZSTD).""" + table = pa.Table.from_pylist(rows) + pq.write_table(table, parquet_out_path, compression="zstd" ) + print(f"[INFO] Parquet saved: {parquet_out_path}") + +def append_manifest(entries: List[Dict[str, Any]]) -> None: + os.makedirs(os.path.dirname(MANIFEST_PATH), exist_ok=True) + with open(MANIFEST_PATH, "a", encoding="utf-8") as f: + for e in entries: + f.write(__to_jsonl_line(e) + "\n") + print(f"[INFO] Manifest appended: {MANIFEST_PATH}") + +def __to_jsonl_line(obj: Dict[str, Any]) -> str: + # Basic compact JSON; avoid pretty for smaller size + import json + return json.dumps(obj, ensure_ascii=False, separators=(",", ":")) \ No newline at end of file