Skip to content

Latest commit

 

History

121 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

VLM-ArabicExtractor

License: MIT Python 3.11+ Architecture: Clean

VLM-ArabicExtractor is an end-to-end, production-grade AI system for processing complex Arabic PDF documents (legal contracts, administrative forms, official records). It converts raw documents into preprocessed images, distills knowledge from state-of-the-art cloud Vision-Language Models (VLMs) into structured JSON datasets, and fine-tunes compact, open-weights Vision-Language Models (Gemma 3 4B) using LlamaFactory LoRA.


Business Goal & Value Proposition

The Problem

Traditional OCR systems (such as Tesseract) struggle with Arabic legal and administrative documents due to:

  • Complex right-to-left (RTL) typography, ligatures, and diacritics.
  • Unstructured output that loses document layout, tables, stamps, and section hierarchy.
  • High cost, latency, and data privacy concerns associated with relying indefinitely on external cloud APIs (e.g. Gemini 3, GPT-4 Vision) for enterprise document parsing at scale.

The Solution

VLM-ArabicExtractor solves this by establishing a complete Knowledge Distillation & SFT Pipeline:

  1. Teacher Distillation: Uses high-capability Cloud VLMs (e.g., gemini-3-flash-preview via LiteLLM) to extract structured JSON metadata, structural elements, and full page markdown content.
  2. Dataset Formatting: Splits annotations into multi-task SFT conversations in LlamaFactory ShareGPT format.
  3. Student Model Fine-Tuning: Fine-tunes an open-weights model (google/gemma-3-4b-it) using LoRA parameter-efficient SFT.
  4. High-Throughput VLLM Inference: Self-hosts the fine-tuned model via a high-performance VLLM OpenAI-compatible server.
  5. On-Premise Deployment: Yields a fast, privacy-preserving, zero-API-cost, domain-adapted Arabic VLM parser that can be deployed entirely on-premise.

Dataset

Source

All 71 PDF documents were collected from the official Saudi Ministry of Justice legislative portal: laws.moj.gov.sa

These are real Saudi legal regulations, administrative bylaws, and official court procedures - some of the most challenging Arabic document types in existence.

Scale

Stage Count
Raw PDFs collected 71
Preprocessed page images 2,058
SFT training samples (×2 tasks) 4,116

Each image was annotated twice - once for Task 1 (full markdown content & structural elements) and once for Task 2 (structured JSON metadata extraction) - doubling the effective dataset size from a single pass of cloud teacher distillation.

Document Complexity

The dataset covers a wide spectrum of real-world document difficulty that traditional OCR systems fundamentally fail on:

Challenge Description
Dense Arabic Typography Multi-column RTL layouts with ligatures, diacritics, and varied font weights
Embedded Tables Structured data grids with merged cells and handwritten values inside cells
Overlapping Stamps Official government seals and watermarks overlaid directly on body text
Handwritten Annotations Inline Arabic handwriting alongside printed text, dates, and signatures
Metadata Fields Document numbers, dates, references, and classifications in header/footer regions
Scan Quality Variance Mix of high-resolution scans and low-quality photocopies with noise and skew
Official Signatures Handwritten executive signatures often covering critical content

Sample Pages

These 4 pages are taken directly from the validation set (PDFs 0011 and 0012) to illustrate the challenges:

pdf0011 - page 8 pdf0011 - page 10
Dense Arabic legal text with overlapping stamps and a signature block Structured table with embedded handwritten values and official seal
Dense body text + overlapping archive stamp + signature Embedded table with handwritten cell values + official seal
pdf0012 - page 9 pdf0012 - page 10
Two-column legal Arabic layout with footer metadata Multi-section Arabic regulation layout with numbered clauses
Two-column layout, footer metadata, section headers Numbered clause structure, section titles, dense paragraphs

Architecture & Design Patterns

The codebase is engineered adhering strictly to SOLID design principles and clean architectural patterns:

src/
├── assets/                  # Centralized YAML configuration files
│   └── config.yaml
├── controllers/             # Workflow Orchestrators
│   ├── data_controller.py               # PDF → Preprocessed Image pipeline
│   ├── evaluation_controller.py         # Zero-shot VLM evaluation
│   ├── distillation_controller.py       # Teacher distillation & cost tracking
│   ├── dataset_formatting_controller.py # Distillation JSONL → LlamaFactory SFT JSON
│   └── finetuning_controller.py         # LlamaFactory setup & execution
├── helpers/                 # Domain Services & Utilities
│   ├── pdf_converter.py                 # PDF page rendering via pdf2image
│   ├── preprocessors.py                 # Image processing strategies
│   ├── prompts.py                       # Centralized Arabic OCR prompt registry
│   ├── cost_tracker.py                  # Token usage & USD cost monitoring
│   ├── sft_writer.py                    # Append-mode JSONL dataset writer
│   ├── image_collector.py               # Dataset filesystem scanner
│   ├── sft_reader.py                    # Distillation JSONL parser
│   ├── record_splitter.py               # Multi-task SFT record transformation
│   ├── dataset_writer.py                # LlamaFactory JSON dataset writer
│   ├── dataset_info_manager.py          # LlamaFactory dataset_info.json manager
│   ├── finetune_config_builder.py       # LlamaFactory train_lora YAML generator
│   ├── llamafactory_runner.py          # Subprocess runner for training CLI
│   └── llamafactory_installer.py       # Reproducible git checkout & package setup
├── models/                  # Configuration, DTOs & VLM Model Adapters
│   ├── base_vlm.py                      # Abstract VLM Model interface
│   ├── gemma_vlm.py                     # Local Gemma-3 Hugging Face adapter
│   ├── cloud_vlm.py                     # LiteLLM Cloud API adapter
│   ├── vllm_client.py                   # High-throughput VLLM API adapter
│   ├── factory.py                       # VLM Factory pattern implementation
│   ├── config_schema.py                 # Pydantic BaseSettings (.env sync)
│   └── schemas.py                       # Data Transfer Objects (DTOs)
├── stores/                  # Storage Layer Abstractions
│   └── file_store.py                    # Filesystem operations
└── utils/                   # Logging & Custom Domain Exceptions
    ├── logger.py                        # Structured console & file logger
    └── exceptions.py                    # Custom exception hierarchy

Applied Patterns

  • Strategy Pattern (src/helpers/preprocessors.py): Pluggable image preprocessing strategies (ImagePreprocessorStrategy, StandardArabicOCRPreprocessor).
  • Adapter Pattern (src/models/base_vlm.py, gemma_vlm.py, cloud_vlm.py): Unified BaseVLM abstract interface adapting both local Hugging Face Transformers/Optimum and provider-agnostic Cloud APIs (LiteLLM).
  • Factory Pattern (src/models/factory.py): VLMFactory decouples model backend creation from controllers.
  • Single Source of Truth Configuration (src/models/config_schema.py): AppConfig handles all environment variables (OPENAI_API_KEY, HF_TOKEN, HF_REPO_ID, WANDB_TOKEN) and synchronizes them to os.environ. HF_REPO_ID serves as the single mandatory Hugging Face repository identifier for both uploading fine-tuned checkpoints and downloading LoRA adapter snapshots for merging.
  • Dependency Injection: Controllers accept custom helpers and service instances for testability.

End-to-End Pipeline Stages

flowchart TD
    A[Raw PDF Documents] -->|1. preprocess| B[Preprocessed JPEG Images]
    B -->|2. evaluate| C[Base VLM vs Cloud VLM Benchmark]
    B -->|3. distill| D[Teacher Model Annotations ocr-image-sft.jsonl]
    D -->|4. format-dataset| E[LlamaFactory SFT Datasets train-v1.json / val-v1.json]
    E -->|5a. finetune-install| F[Reproducible LlamaFactory Env]
    E -->|5b. finetune-setup| G[dataset_info.json & ocr_finetune.yaml]
    F --> H[5c. finetune-train]
    G --> H
    H -->|LoRA SFT| I[Fine-Tuned Gemma-3 4B VLM Model]
    I -->|6. merge| J[Merged Standalone VLM Model]
    B -->|7. evaluate --model-id ./saves| J
    J -->|8. serve| K[VLLM Inference Server]
    B -->|9. evaluate --model-type vllm| K
Loading

Stage 1: PDF to Image Preprocessing (main.py preprocess)

  • Converts raw PDF pages into aspect-ratio-preserved, grayscale JPEG images.
  • Applies contrast adjustment factor ($1.5\times$) to optimize image clarity for OCR.

Stage 2: Model Evaluation (main.py evaluate)

  • Evaluates zero-shot or fine-tuned OCR & detail extraction capabilities on document page images using structured prompts.
  • Accepts any --model-id: Base HF models (e.g. google/gemma-3-4b-it), Merged model directories (e.g. ./saves/arabic-legal-documents-ocr-1.0), or Cloud APIs.
  • Supports preset task prompts: task1 (Markdown Content & Structure), task2 (Metadata & Classification), or default (Full Document Analysis).

Stage 3: Teacher Knowledge Distillation (main.py distill)

  • Annotates all document images across PDFs using the teacher cloud VLM (gemini-3-flash-preview).
  • Implements real-time token tracking (TokenUsageTracker) and outputs cost estimates ($/1M tokens).
  • Appends annotations to a crash-safe JSONL file (ocr-image-sft.jsonl).

Stage 4: Dataset Formatting (main.py format-dataset)

  • Reads distillation output and transforms each record into multi-task SFT conversations:
    • Task 1: Full page markdown content & structural elements extraction.
    • Task 2: Detailed structured JSON extraction.
  • Segregates hold-out validation PDFs (0012.pdf, 0005.pdf, 0011.pdf) to prevent data leakage.
  • Generates LlamaFactory ShareGPT-formatted JSON arrays (train-v1.json, val-v1.json).

Stage 5: LlamaFactory Fine-Tuning Pipeline (main.py finetune-*)

  • Install: Clones and installs LlamaFactory at commit 762b480131908d37736ad9aa3f12e87f8f7e6313 for 100% reproducibility.
  • Setup: Registers datasets in LlamaFactory/data/dataset_info.json and generates ocr_finetune.yaml containing model, LoRA rank (96), batch sizes, LR schedules, W&B tracking, and HF Hub export targets (HF_REPO_ID).
  • Train: Triggers llamafactory-cli train via controlled subprocess with real-time log streaming.

Stage 6: LoRA Adapter Merging (main.py merge-*)

  • Download: Downloads the specific LoRA checkpoint snapshot from HF_REPO_ID.
  • Setup: Writes merge-arabic-ocr.yaml config pointing base model to downloaded LoRA adapter.
  • Run: Executes llamafactory-cli export to produce a standalone merged model ready for self-hosted deployment.

Quick Start & Installation

Prerequisites

  • Python 3.11+
  • System dependency: poppler-utils (for PDF rendering)
# Ubuntu / Debian
sudo apt-get update && sudo apt-get install -y poppler-utils

Installation

# Clone the repository
git clone <repo-url>
cd VLM-ArabicExtractor

# Create and activate virtual environment
python -m venv .venv
source .venv/bin/activate

# Install requirements
pip install -r requirements.txt

Environment Configuration

Copy .env.example to .env and fill in your credentials:

cp .env.example .env
# App Storage Directories
DATA_DIR=./assets
OUTPUT_DIR=./images

# Cloud LLM / OpenRouter API Key (For Knowledge Distillation & Cloud Evaluation)
OPENAI_API_KEY=sk-or-v1-your-openrouter-key
OPENAI_BASE_URL=https://openrouter.ai/api/v1
VLLM_BASE_URL=http://localhost:8000/v1

# Hugging Face Authentication & Repository (Mandatory for model push & checkpoint download)
HF_TOKEN=hf_your_huggingface_token
HF_REPO_ID=your_username/arabic-legal-documents-ocr-1.0

# Weights & Biases Logging Token (For Fine-Tuning Experiment Tracking)
WANDB_TOKEN=wandb_v1_your_wandb_token

Usage Guide

1. Preprocess PDF Documents

python main.py preprocess \
    --input-dir ./data/pdfs \
    --output-dir ./images \
    --max-width 600 \
    --contrast 1.5

2. Evaluate Base Model, Cloud Teacher, Merged Model, or VLLM Server

The evaluate command runs OCR and detail extraction on a single document page image using --image-path.

The --model-id parameter is flexible and accepts:

  • Base Hugging Face Model: google/gemma-3-4b-it (downloads & runs the base VLM).
  • Fine-Tuned Merged Model: Path to local export directory (e.g. ./saves/arabic-legal-documents-ocr-1.0 or ./downloaded_checkpoint).
  • Cloud Teacher Model: openrouter/google/gemini-3-flash-preview or gemini-2.5-flash.
  • VLLM Server Model: arabic-legal-documents-ocr-1.0 (as deployed in the VLLM instance).

You can also specify the task prompt preset via --task:

  • default: Full document extraction schema (all metadata + content).
  • task1: Markdown content & structural elements extraction.
  • task2: Document classification, official marks, dates & physical properties.
# 1. Base Model Evaluation (Gemma-3 4B)
python main.py evaluate \
    --image-path ./images/0001/page_001.jpg \
    --model-id google/gemma-3-4b-it

# 2. Cloud Gemini-3 Flash Teacher Model Evaluation
python main.py evaluate \
    --image-path ./images/0001/page_001.jpg \
    --model-id openrouter/google/gemini-3-flash-preview

# 3. Merged Fine-Tuned Model Evaluation (Task 1: Markdown Content & Structure)
python main.py evaluate \
    --image-path ./images/0012/page_010.jpg \
    --model-id ./saves/arabic-legal-documents-ocr-1.0 \
    --task task1

# 4. Merged Fine-Tuned Model Evaluation (Task 2: Detailed Document Classification)
python main.py evaluate \
    --image-path ./images/0012/page_010.jpg \
    --model-id ./saves/arabic-legal-documents-ocr-1.0 \
    --task task2

# 5. High-Throughput VLLM Inference (Requires VLLM Server running)
python main.py evaluate \
    --image-path ./images/0001/page_001.jpg \
    --model-id /workspace/saves/arabic-legal-documents-ocr-1.0 \
    --model-type vllm

3. Run Knowledge Distillation

python main.py distill \
    --images-dir ./images \
    --pdf-dir ./data/pdfs \
    --output-file ./ocr-image-sft.jsonl \
    --model-id openrouter/google/gemini-3-flash-preview \
    --price-input 0.5 \
    --price-output 3.0

4. Format Dataset for SFT

python main.py format-dataset \
    --sft-file ./ocr-image-sft.jsonl \
    --output-dir ./datasets/llamafactory-ocr-finetune-data \
    --val-pdfs "0012.pdf,0005.pdf,0011.pdf" \
    --version "v1"

5. Fine-Tuning Pipeline Commands

Step 5a: Install LlamaFactory (Reproducible Commit)

python main.py finetune-install \
    --target-dir ./LlamaFactory \
    --commit 762b480131908d37736ad9aa3f12e87f8f7e6313

Step 5b: Configure & Generate YAML Setup

python main.py finetune-setup \
    --train-json ./datasets/llamafactory-ocr-finetune-data/train-v1.json \
    --val-json ./datasets/llamafactory-ocr-finetune-data/val-v1.json \
    --yaml-output ./LlamaFactory/examples/train_lora/ocr_finetune.yaml \
    --model-name google/gemma-3-4b-it \
    --epochs 20.0 \
    --learning-rate 1e-4 \
    --lora-rank 96

Step 5c: Run LlamaFactory SFT Training

python main.py finetune-train \
    --yaml-config ./LlamaFactory/examples/train_lora/ocr_finetune.yaml \
    --llama-factory-dir ./LlamaFactory

Full Step 5 End-to-End Command (Setup + Train)

python main.py finetune \
    --train-json ./datasets/llamafactory-ocr-finetune-data/train-v1.json \
    --val-json ./datasets/llamafactory-ocr-finetune-data/val-v1.json \
    --model-name google/gemma-3-4b-it \
    --epochs 20.0

Testing

The repository includes a comprehensive unit-test suite using pytest. All model backends and subprocess runners are mocked, allowing tests to run quickly without requiring API keys or GPU hardware.

# Run all unit tests
pytest tests/ -v

License

This project is released under the MIT License.

About

A production-grade AI pipeline for Arabic document processing. It uses knowledge distillation from cloud VLMs to generate SFT datasets, fine-tunes local Gemma-3 4B models via LlamaFactory, merges the adapters, and deploys a high-throughput VLLM server for cost-free, on-premise, highly accurate structured JSON extraction.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages