xiaoxue_admin
/
ai-embedder-engine
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

同步完整源码 - 2026-05-25

This commit is contained in:
xiaoxue 2026-05-25 00:32:30 +08:00
commit 96f9d9dc81
11 changed files with 1139 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
LICENSE Normal file
View File

@ -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.

448
README.md Normal file
View File

@ -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
---

123
build_embeddings.py Normal file
View File

@ -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

28
chunker.py Normal file
View File

@ -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

36
config.py Normal file
View File

@ -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.

34
embedder.py Normal file
View File

@ -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

325
faiss_index_builder.py Normal file
View File

@ -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'.")

17
pdf_text.py Normal file
View File

@ -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

67
prepare_pdf.py Normal file
View File

@ -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()

5
requirements.txt Normal file
View File

@ -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

33
writer.py Normal file
View File

@ -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=(",", ":"))