document-classifier/
├── src/
│   ├── __init__.py
│   ├── config.py
│   ├── preprocessing/
│   │   ├── __init__.py
│   │   ├── ocr.py              # OCR processing
│   │   ├── layout.py           # Layout analysis
│   │   └── normalizer.py       # Text normalization
│   ├── models/
│   │   ├── __init__.py
│   │   ├── classifier.py       # Document classifier
│   │   ├── extractor.py        # Field extraction
│   │   └── validator.py        # Output validation
│   ├── training/
│   │   ├── __init__.py
│   │   ├── dataset.py          # Dataset handling
│   │   ├── trainer.py          # Training loop
│   │   └── evaluation.py       # Model evaluation
│   ├── inference/
│   │   ├── __init__.py
│   │   ├── pipeline.py         # Inference pipeline
│   │   └── optimization.py     # Model optimization
│   └── api/
│       ├── __init__.py
│       └── main.py             # FastAPI endpoints
├── notebooks/
│   ├── data_exploration.ipynb
│   └── model_analysis.ipynb
├── tests/
└── requirements.txt

# src/config.py
from pydantic_settings import BaseSettings
from typing import List, Dict
from pathlib import Path

class Settings(BaseSettings):
    # Model Settings
    base_model: str = "microsoft/layoutlmv3-base"
    num_labels: int = 10
    max_seq_length: int = 512

    # Training Settings
    learning_rate: float = 5e-5
    batch_size: int = 8
    num_epochs: int = 10
    warmup_ratio: float = 0.1

    # LoRA Settings
    lora_r: int = 16
    lora_alpha: int = 32
    lora_dropout: float = 0.1

    # Paths
    data_dir: Path = Path("./data")
    model_dir: Path = Path("./models")
    output_dir: Path = Path("./outputs")

    # Document Types
    document_types: List[str] = [
        "invoice",
        "contract",
        "receipt",
        "report",
        "letter",
        "form",
        "resume",
        "statement",
        "memo",
        "other"
    ]

    # Extraction Fields by Document Type
    extraction_fields: Dict[str, List[str]] = {
        "invoice": ["invoice_number", "date", "total", "vendor", "line_items"],
        "contract": ["parties", "effective_date", "terms", "signatures"],
        "receipt": ["merchant", "date", "total", "items"],
        "resume": ["name", "email", "phone", "experience", "education"]
    }

    class Config:
        env_file = ".env"

settings = Settings()

# src/preprocessing/layout.py
from typing import Dict, List, Tuple
import numpy as np
from PIL import Image
from dataclasses import dataclass
import pytesseract
from pdf2image import convert_from_path

@dataclass
class TextBox:
    text: str
    bbox: Tuple[int, int, int, int]  # x1, y1, x2, y2
    confidence: float

@dataclass
class DocumentLayout:
    text_boxes: List[TextBox]
    full_text: str
    image_size: Tuple[int, int]
    page_number: int

class LayoutAnalyzer:
    """Analyzes document layout and extracts text with positions."""

    def __init__(self, ocr_engine: str = "tesseract"):
        self.ocr_engine = ocr_engine

    def process_pdf(self, pdf_path: str) -> List[DocumentLayout]:
        """Process PDF and extract layout for each page."""
        images = convert_from_path(pdf_path, dpi=300)
        layouts = []

        for i, image in enumerate(images):
            layout = self.process_image(image)
            layout.page_number = i + 1
            layouts.append(layout)

        return layouts

    def process_image(self, image: Image.Image) -> DocumentLayout:
        """Process a single image and extract layout."""
        # Get OCR data with bounding boxes
        ocr_data = pytesseract.image_to_data(
            image,
            output_type=pytesseract.Output.DICT
        )

        text_boxes = []
        for i, text in enumerate(ocr_data["text"]):
            if text.strip():
                bbox = (
                    ocr_data["left"][i],
                    ocr_data["top"][i],
                    ocr_data["left"][i] + ocr_data["width"][i],
                    ocr_data["top"][i] + ocr_data["height"][i]
                )
                confidence = float(ocr_data["conf"][i]) / 100

                text_boxes.append(TextBox(
                    text=text,
                    bbox=bbox,
                    confidence=confidence
                ))

        full_text = " ".join([tb.text for tb in text_boxes])

        return DocumentLayout(
            text_boxes=text_boxes,
            full_text=full_text,
            image_size=image.size,
            page_number=0
        )

    def normalize_bboxes(
        self,
        layout: DocumentLayout,
        target_size: int = 1000
    ) -> List[List[int]]:
        """Normalize bounding boxes to 0-1000 range for LayoutLM."""
        width, height = layout.image_size
        normalized = []

        for box in layout.text_boxes:
            x1, y1, x2, y2 = box.bbox
            normalized.append([
                int(x1 / width * target_size),
                int(y1 / height * target_size),
                int(x2 / width * target_size),
                int(y2 / height * target_size)
            ])

        return normalized

# src/models/classifier.py
import torch
import torch.nn as nn
from transformers import (
    LayoutLMv3ForSequenceClassification,
    LayoutLMv3Processor,
    AutoTokenizer
)
from peft import get_peft_model, LoraConfig, TaskType
from typing import Dict, List, Optional
from ..config import settings

class DocumentClassifier:
    """LayoutLMv3-based document classifier with LoRA fine-tuning."""

    def __init__(self, model_path: Optional[str] = None):
        self.processor = LayoutLMv3Processor.from_pretrained(
            settings.base_model,
            apply_ocr=False
        )

        if model_path:
            self.model = LayoutLMv3ForSequenceClassification.from_pretrained(
                model_path,
                num_labels=settings.num_labels
            )
        else:
            self.model = LayoutLMv3ForSequenceClassification.from_pretrained(
                settings.base_model,
                num_labels=settings.num_labels
            )

        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self.model.to(self.device)

        self.id2label = {
            i: label for i, label in enumerate(settings.document_types)
        }
        self.label2id = {
            label: i for i, label in enumerate(settings.document_types)
        }

    def apply_lora(self):
        """Apply LoRA for efficient fine-tuning."""
        lora_config = LoraConfig(
            task_type=TaskType.SEQ_CLS,
            r=settings.lora_r,
            lora_alpha=settings.lora_alpha,
            lora_dropout=settings.lora_dropout,
            target_modules=["query", "value"]
        )
        self.model = get_peft_model(self.model, lora_config)
        self.model.print_trainable_parameters()

    def preprocess(
        self,
        image,
        words: List[str],
        boxes: List[List[int]]
    ) -> Dict[str, torch.Tensor]:
        """Preprocess document for model input."""
        encoding = self.processor(
            image,
            words,
            boxes=boxes,
            max_length=settings.max_seq_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        return {k: v.to(self.device) for k, v in encoding.items()}

    def predict(
        self,
        image,
        words: List[str],
        boxes: List[List[int]]
    ) -> Dict[str, float]:
        """Predict document type with confidence scores."""
        self.model.eval()

        inputs = self.preprocess(image, words, boxes)

        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)

        predictions = {}
        for i, prob in enumerate(probs[0].cpu().numpy()):
            predictions[self.id2label[i]] = float(prob)

        return predictions

    def get_top_prediction(
        self,
        image,
        words: List[str],
        boxes: List[List[int]]
    ) -> tuple:
        """Get top prediction with confidence."""
        predictions = self.predict(image, words, boxes)
        top_label = max(predictions, key=predictions.get)
        return top_label, predictions[top_label]

# src/models/extractor.py
import torch
from transformers import (
    LayoutLMv3ForTokenClassification,
    LayoutLMv3Processor
)
from typing import Dict, List, Tuple
from ..config import settings

class FieldExtractor:
    """Extracts structured fields from documents using token classification."""

    def __init__(self, model_path: str = None):
        self.processor = LayoutLMv3Processor.from_pretrained(
            settings.base_model,
            apply_ocr=False
        )

        # Build label vocabulary for all field types
        self.build_label_vocab()

        if model_path:
            self.model = LayoutLMv3ForTokenClassification.from_pretrained(
                model_path,
                num_labels=len(self.label2id)
            )
        else:
            self.model = LayoutLMv3ForTokenClassification.from_pretrained(
                settings.base_model,
                num_labels=len(self.label2id)
            )

        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self.model.to(self.device)

    def build_label_vocab(self):
        """Build BIO tagging vocabulary."""
        labels = ["O"]  # Outside
        for doc_type, fields in settings.extraction_fields.items():
            for field in fields:
                labels.append(f"B-{field}")
                labels.append(f"I-{field}")

        self.id2label = {i: l for i, l in enumerate(labels)}
        self.label2id = {l: i for i, l in enumerate(labels)}

    def extract(
        self,
        image,
        words: List[str],
        boxes: List[List[int]],
        document_type: str
    ) -> Dict[str, List[str]]:
        """Extract fields from document."""
        self.model.eval()

        encoding = self.processor(
            image,
            words,
            boxes=boxes,
            max_length=settings.max_seq_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )

        inputs = {k: v.to(self.device) for k, v in encoding.items()}

        with torch.no_grad():
            outputs = self.model(**inputs)
            predictions = torch.argmax(outputs.logits, dim=-1)

        # Convert predictions to field values
        expected_fields = settings.extraction_fields.get(document_type, [])
        extracted = {field: [] for field in expected_fields}

        current_field = None
        current_value = []

        for i, (word, pred_id) in enumerate(zip(words, predictions[0])):
            if i >= len(predictions[0]):
                break

            label = self.id2label[pred_id.item()]

            if label.startswith("B-"):
                # Save previous field
                if current_field and current_value:
                    extracted[current_field].append(" ".join(current_value))

                # Start new field
                current_field = label[2:]
                current_value = [word]

            elif label.startswith("I-") and current_field == label[2:]:
                current_value.append(word)

            else:
                # Save and reset
                if current_field and current_value:
                    extracted[current_field].append(" ".join(current_value))
                current_field = None
                current_value = []

        # Save last field
        if current_field and current_value:
            extracted[current_field].append(" ".join(current_value))

        return extracted

# src/training/trainer.py
import torch
from torch.utils.data import DataLoader
from transformers import (
    get_linear_schedule_with_warmup,
    AdamW
)
from tqdm import tqdm
from typing import Dict, Optional
import wandb
from ..config import settings
from .dataset import DocumentDataset
from .evaluation import Evaluator

class DocumentTrainer:
    """Training pipeline for document understanding models."""

    def __init__(
        self,
        model,
        train_dataset: DocumentDataset,
        val_dataset: Optional[DocumentDataset] = None,
        use_wandb: bool = True
    ):
        self.model = model
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.use_wandb = use_wandb

        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self.model.to(self.device)

        if use_wandb:
            wandb.init(project="document-classifier")

    def train(self) -> Dict[str, float]:
        """Run training loop."""
        train_loader = DataLoader(
            self.train_dataset,
            batch_size=settings.batch_size,
            shuffle=True,
            num_workers=4
        )

        # Optimizer
        optimizer = AdamW(
            self.model.parameters(),
            lr=settings.learning_rate,
            weight_decay=0.01
        )

        # Scheduler
        total_steps = len(train_loader) * settings.num_epochs
        warmup_steps = int(total_steps * settings.warmup_ratio)
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=total_steps
        )

        best_val_acc = 0
        training_history = []

        for epoch in range(settings.num_epochs):
            # Training phase
            self.model.train()
            total_loss = 0
            correct = 0
            total = 0

            progress = tqdm(train_loader, desc=f"Epoch {epoch + 1}")
            for batch in progress:
                # Move to device
                batch = {k: v.to(self.device) for k, v in batch.items()}

                # Forward pass
                outputs = self.model(**batch)
                loss = outputs.loss

                # Backward pass
                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()
                scheduler.step()

                # Track metrics
                total_loss += loss.item()
                predictions = torch.argmax(outputs.logits, dim=-1)
                correct += (predictions == batch["labels"]).sum().item()
                total += batch["labels"].size(0)

                progress.set_postfix({
                    "loss": loss.item(),
                    "acc": correct / total
                })

            avg_loss = total_loss / len(train_loader)
            train_acc = correct / total

            # Validation phase
            val_metrics = {}
            if self.val_dataset:
                val_metrics = self.evaluate()

                # Save best model
                if val_metrics.get("accuracy", 0) > best_val_acc:
                    best_val_acc = val_metrics["accuracy"]
                    self.save_model(settings.model_dir / "best_model")

            # Log metrics
            metrics = {
                "epoch": epoch + 1,
                "train_loss": avg_loss,
                "train_accuracy": train_acc,
                **{f"val_{k}": v for k, v in val_metrics.items()}
            }
            training_history.append(metrics)

            if self.use_wandb:
                wandb.log(metrics)

            print(f"Epoch {epoch + 1}: loss={avg_loss:.4f}, "
                  f"train_acc={train_acc:.4f}, "
                  f"val_acc={val_metrics.get('accuracy', 'N/A')}")

        return training_history

    def evaluate(self) -> Dict[str, float]:
        """Evaluate on validation set."""
        self.model.eval()
        evaluator = Evaluator(self.model, self.val_dataset, self.device)
        return evaluator.evaluate()

    def save_model(self, path: str):
        """Save model checkpoint."""
        self.model.save_pretrained(path)

# src/inference/optimization.py
import torch
import onnx
import onnxruntime as ort
from transformers import LayoutLMv3ForSequenceClassification
from typing import Dict
import numpy as np

class ModelOptimizer:
    """Optimizes models for production inference."""

    def __init__(self, model_path: str):
        self.model_path = model_path
        self.model = LayoutLMv3ForSequenceClassification.from_pretrained(
            model_path
        )

    def quantize_dynamic(self, output_path: str):
        """Apply dynamic quantization for CPU inference."""
        quantized_model = torch.quantization.quantize_dynamic(
            self.model,
            {torch.nn.Linear},
            dtype=torch.qint8
        )
        torch.save(quantized_model.state_dict(), output_path)
        return quantized_model

    def export_onnx(
        self,
        output_path: str,
        opset_version: int = 14
    ):
        """Export model to ONNX format."""
        self.model.eval()

        # Create dummy inputs
        dummy_inputs = {
            "input_ids": torch.randint(0, 1000, (1, 512)),
            "attention_mask": torch.ones(1, 512, dtype=torch.long),
            "bbox": torch.randint(0, 1000, (1, 512, 4)),
            "pixel_values": torch.randn(1, 3, 224, 224)
        }

        # Export
        torch.onnx.export(
            self.model,
            tuple(dummy_inputs.values()),
            output_path,
            input_names=list(dummy_inputs.keys()),
            output_names=["logits"],
            dynamic_axes={
                "input_ids": {0: "batch_size", 1: "sequence"},
                "attention_mask": {0: "batch_size", 1: "sequence"},
                "bbox": {0: "batch_size", 1: "sequence"},
                "pixel_values": {0: "batch_size"}
            },
            opset_version=opset_version
        )

        # Verify
        onnx_model = onnx.load(output_path)
        onnx.checker.check_model(onnx_model)

        return output_path


class ONNXInference:
    """ONNX Runtime inference for optimized models."""

    def __init__(self, model_path: str, use_gpu: bool = False):
        providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
        self.session = ort.InferenceSession(model_path, providers=providers)

    def predict(
        self,
        input_ids: np.ndarray,
        attention_mask: np.ndarray,
        bbox: np.ndarray,
        pixel_values: np.ndarray
    ) -> np.ndarray:
        """Run inference with ONNX Runtime."""
        inputs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "bbox": bbox,
            "pixel_values": pixel_values
        }

        outputs = self.session.run(None, inputs)
        return outputs[0]

    def get_latency_stats(self, num_runs: int = 100) -> Dict[str, float]:
        """Benchmark inference latency."""
        import time

        # Dummy input
        inputs = {
            "input_ids": np.random.randint(0, 1000, (1, 512)).astype(np.int64),
            "attention_mask": np.ones((1, 512), dtype=np.int64),
            "bbox": np.random.randint(0, 1000, (1, 512, 4)).astype(np.int64),
            "pixel_values": np.random.randn(1, 3, 224, 224).astype(np.float32)
        }

        # Warmup
        for _ in range(10):
            self.session.run(None, inputs)

        # Benchmark
        latencies = []
        for _ in range(num_runs):
            start = time.time()
            self.session.run(None, inputs)
            latencies.append((time.time() - start) * 1000)

        return {
            "mean_ms": np.mean(latencies),
            "p50_ms": np.percentile(latencies, 50),
            "p95_ms": np.percentile(latencies, 95),
            "p99_ms": np.percentile(latencies, 99)
        }

# src/api/main.py
from fastapi import FastAPI, UploadFile, File, HTTPException
from pydantic import BaseModel
from typing import Dict, List, Optional
import tempfile
from PIL import Image
import io

from ..preprocessing.layout import LayoutAnalyzer
from ..models.classifier import DocumentClassifier
from ..models.extractor import FieldExtractor
from ..config import settings

app = FastAPI(
    title="Document Understanding API",
    description="Classify and extract data from documents"
)

# Initialize models
layout_analyzer = LayoutAnalyzer()
classifier = DocumentClassifier(settings.model_dir / "classifier")
extractor = FieldExtractor(settings.model_dir / "extractor")

class ClassificationResult(BaseModel):
    document_type: str
    confidence: float
    all_scores: Dict[str, float]

class ExtractionResult(BaseModel):
    document_type: str
    fields: Dict[str, List[str]]
    confidence: float

class ProcessingResult(BaseModel):
    classification: ClassificationResult
    extraction: Optional[ExtractionResult]
    page_count: int

@app.post("/classify", response_model=ClassificationResult)
async def classify_document(file: UploadFile = File(...)):
    """Classify a document type."""
    # Save uploaded file
    content = await file.read()

    if file.filename.lower().endswith(".pdf"):
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
            f.write(content)
            layouts = layout_analyzer.process_pdf(f.name)
            layout = layouts[0]  # Use first page for classification
    else:
        image = Image.open(io.BytesIO(content))
        layout = layout_analyzer.process_image(image)

    # Get words and boxes
    words = [tb.text for tb in layout.text_boxes]
    boxes = layout_analyzer.normalize_bboxes(layout)

    # Classify
    predictions = classifier.predict(
        Image.open(io.BytesIO(content)) if not file.filename.lower().endswith(".pdf") else None,
        words,
        boxes
    )

    top_label = max(predictions, key=predictions.get)

    return ClassificationResult(
        document_type=top_label,
        confidence=predictions[top_label],
        all_scores=predictions
    )

@app.post("/extract", response_model=ExtractionResult)
async def extract_fields(
    file: UploadFile = File(...),
    document_type: Optional[str] = None
):
    """Extract structured fields from a document."""
    content = await file.read()

    if file.filename.lower().endswith(".pdf"):
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
            f.write(content)
            layouts = layout_analyzer.process_pdf(f.name)
            layout = layouts[0]
            image = None
    else:
        image = Image.open(io.BytesIO(content))
        layout = layout_analyzer.process_image(image)

    words = [tb.text for tb in layout.text_boxes]
    boxes = layout_analyzer.normalize_bboxes(layout)

    # Classify if not provided
    if not document_type:
        top_label, confidence = classifier.get_top_prediction(
            image, words, boxes
        )
        document_type = top_label
    else:
        confidence = 1.0

    # Extract fields
    fields = extractor.extract(image, words, boxes, document_type)

    return ExtractionResult(
        document_type=document_type,
        fields=fields,
        confidence=confidence
    )

@app.post("/process", response_model=ProcessingResult)
async def process_document(file: UploadFile = File(...)):
    """Full document processing pipeline."""
    content = await file.read()

    # Process all pages
    if file.filename.lower().endswith(".pdf"):
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
            f.write(content)
            layouts = layout_analyzer.process_pdf(f.name)
    else:
        image = Image.open(io.BytesIO(content))
        layouts = [layout_analyzer.process_image(image)]

    # Use first page for classification
    layout = layouts[0]
    words = [tb.text for tb in layout.text_boxes]
    boxes = layout_analyzer.normalize_bboxes(layout)

    # Classify
    predictions = classifier.predict(None, words, boxes)
    top_label = max(predictions, key=predictions.get)
    confidence = predictions[top_label]

    classification = ClassificationResult(
        document_type=top_label,
        confidence=confidence,
        all_scores=predictions
    )

    # Extract if we have fields for this document type
    extraction = None
    if top_label in settings.extraction_fields:
        fields = extractor.extract(None, words, boxes, top_label)
        extraction = ExtractionResult(
            document_type=top_label,
            fields=fields,
            confidence=confidence
        )

    return ProcessingResult(
        classification=classification,
        extraction=extraction,
        page_count=len(layouts)
    )

@app.get("/health")
async def health():
    return {"status": "healthy"}