adapters/__init__.py

"""
Adapters module registry for 9 OCR & Document AI models.
"""

from typing import Dict, Type
from adapters.base import BaseOCRAdapter
from adapters.deepseek_adapter import DeepSeekOCRAdapter
from adapters.qwen_adapter import QwenVLAdapter
from adapters.pp_structure_adapter import PPStructureAdapter
from adapters.nuextract_adapter import NuExtractAdapter
from adapters.pp_ocr_adapter import PPOCRAdapter
from adapters.olmocr_adapter import OlmOCRAdapter
from adapters.mineru_adapter import MinerUAdapter
from adapters.nanonets_adapter import NanonetsOCRAdapter
from adapters.hunyuan_adapter import HunyuanOCRAdapter

# 9 Models Catalog
ADAPTER_REGISTRY: Dict[str, Type[BaseOCRAdapter]] = {
    "DeepSeek-OCR": DeepSeekOCRAdapter,
    "Qwen3-VL": QwenVLAdapter,
    "PP-StructureV3": PPStructureAdapter,
    "NuExtract3": NuExtractAdapter,
    "PP-OCRv5": PPOCRAdapter,
    "olmOCR": OlmOCRAdapter,
    "MinerU 3": MinerUAdapter,
    "Nanonets-OCR2-3B": NanonetsOCRAdapter,
    "HunyuanOCR": HunyuanOCRAdapter
}

MODEL_CATALOG = {
    "DeepSeek-OCR": {
        "id": "deepseek-ai/DeepSeek-OCR",
        "description": "High-compression Optical Character Recognition engine with visual patch tiling.",
        "default_type": "markdown",
        "supports_bounding_box": True
    },
    "Qwen3-VL": {
        "id": "Qwen/Qwen2.5-VL-3B-Instruct",
        "description": "Multilingual Vision-Language Model with visual grounding and coordinate localization.",
        "default_type": "markdown",
        "supports_bounding_box": True
    },
    "PP-StructureV3": {
        "id": "PaddleOCR/PP-StructureV3",
        "description": "Multi-column layout analysis and HTML table structure recognition engine.",
        "default_type": "markdown",
        "supports_bounding_box": True
    },
    "NuExtract3": {
        "id": "numind/NuExtract-v1.5",
        "description": "Multimodal structured information & schema extraction engine.",
        "default_type": "json",
        "supports_bounding_box": True
    },
    "PP-OCRv5": {
        "id": "PaddleOCR/PP-OCRv5",
        "description": "High-speed text detection & recognition engine with Vietnamese dictionary support.",
        "default_type": "text",
        "supports_bounding_box": True
    },
    "olmOCR": {
        "id": "allenai/olmOCR-7B-0225-preview",
        "description": "Document & academic paper Markdown extraction engine with reading order linearization.",
        "default_type": "markdown",
        "supports_bounding_box": True
    },
    "MinerU 3": {
        "id": "opendatalab/MinerU2.5-Pro-2605-1.2B",
        "description": "Scientific document extraction, LaTeX formula recognition, and complex layout parser.",
        "default_type": "markdown",
        "supports_bounding_box": True
    },
    "Nanonets-OCR2-3B": {
        "id": "nanonets/Nanonets-OCR2-3B",
        "description": "3B Multimodal OCR specialized in structured Markdown, flowcharts (Mermaid), and tables.",
        "default_type": "markdown",
        "supports_bounding_box": True
    },
    "HunyuanOCR": {
        "id": "tencent/HunyuanOCR",
        "description": "Lightweight 1B VLM with SigLIP-v2 visual encoder and XD-RoPE 2D/3D spatial alignment.",
        "default_type": "markdown",
        "supports_bounding_box": True
    }
}

AVAILABLE_MODELS = list(ADAPTER_REGISTRY.keys())


def get_adapter_by_name(model_name: str) -> BaseOCRAdapter:
    """Instantiate and return the OCR adapter class for the given model name."""
    if model_name not in ADAPTER_REGISTRY:
        raise ValueError(f"Unknown model '{model_name}'. Available: {AVAILABLE_MODELS}")
    adapter_cls = ADAPTER_REGISTRY[model_name]
    return adapter_cls()


def get_model_info(model_name: str) -> Dict[str, Any]:
    """Retrieve metadata information for a model."""
    if model_name not in MODEL_CATALOG:
        raise ValueError(f"Unknown model '{model_name}'. Available: {AVAILABLE_MODELS}")
    return MODEL_CATALOG[model_name]