ai-embedder-engine - GitHub backup
Go to file
xiaoxue 96f9d9dc81 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
LICENSE 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
README.md 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
build_embeddings.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
chunker.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
config.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
embedder.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
faiss_index_builder.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
pdf_text.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
prepare_pdf.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
requirements.txt 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00
writer.py 同步完整源码 - 2026-05-25 2026-05-25 00:32:30 +08:00

README.md

🧠 AI-Embedder-Engine

Table of Contents

  1. Introduction
  2. Why Use This Project
  3. Key Features
  4. Quick Start
  5. Configuration
  6. How the Pipeline Works
  7. Data Formats
  8. Project Structure
  9. Logging & Monitoring
  10. Use Cases
  11. Performance Tips
  12. Vector Indexing with FAISS
  13. Disclaimer
  14. 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)

    python -m venv .venv && source .venv/bin/activate
    # Windows:
    .venv\Scripts\activate
    
  2. Install dependencies

    pip install -r requirements.txt
    
  3. Set your OpenAI API key
    Default environment variable name is: EMBEDDING_API_KEY

    export EMBEDDING_API_KEY=YOUR_OPENAI_KEY
    # Windows:
    set EMBEDDING_API_KEY=YOUR_OPENAI_KEY
    

    You can change this name in config.pyOPENAI_API_KEY_ENV.

📦 Run Instructions

  1. Run on a single PDF

    python build_embeddings.py --pdf "/path/to/document.pdf" --doc-id "DOC-ID-Example"
    

    Outputs

    • data/embeddings/DOC-ID-Example.parquetcanonical archive
    • data/manifest/manifest.jsonl → updated manifest entry (one per document)
  2. Batch Processing Multiple PDFs
    Edit prepare_pdf.py and fill the CHAPTERS list with (doc_id, pdf_path) pairs:

    python prepare_pdf.py
    

    Outputs

    • all embeddings saved under data/embeddings/
    • data/manifest/manifest.jsonl → updated manifest entry (one per document)
  3. 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)

    python faiss_index_builder.py
    

    Run on a single Parquet file

    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:

{
  "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):

{
  "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

[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.
  • 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., 150200 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 and the FAISS MIT 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 file for details.

© 2025 Ofir Ben Shabat • Contact: ofirbshofir@gmail.com