backends/__init__.py

"""
Backend interface for LightOnOCR-1B inference.
Supports both PyTorch and GGUF backends.
"""

from abc import ABC, abstractmethod
from typing import List, Tuple
from PIL import Image


class OCRBackend(ABC):
    """Abstract base class for OCR backends."""
    
    @abstractmethod
    def load_model(self):
        """Load the OCR model."""
        pass
    
    @abstractmethod
    def process_image(self, image: Image.Image, temperature: float = 0.1) -> str:
        """
        Process a single image and return extracted text.
        
        Args:
            image: PIL Image to process
            temperature: Sampling temperature (0 = greedy)
            
        Returns:
            Extracted text as string
        """
        pass
    
    @abstractmethod
    def get_backend_info(self) -> dict:
        """Return backend information (name, device, memory usage, etc.)."""
        pass


def get_available_backends() -> List[str]:
    """Return list of available backend names."""
    backends = ["pytorch"]
    
    # Check for GGUF support (binary or python package)
    from pathlib import Path
    project_root = Path(__file__).parent.parent
    cli_path = project_root / "llama.cpp" / "build" / "bin" / "llama-mtmd-cli"
    
    if cli_path.exists():
        backends.append("gguf")
    else:
        # Fallback check for python package (though we prefer CLI now)
        try:
            import llama_cpp
            backends.append("gguf")
        except ImportError:
            pass
    
    return backends


def create_backend(backend_name: str) -> OCRBackend:
    """
    Factory function to create backend instance.
    
    Args:
        backend_name: "pytorch" or "gguf"
        
    Returns:
        OCRBackend instance
    """
    if backend_name == "pytorch":
        from .pytorch_backend import PyTorchBackend
        return PyTorchBackend()
    elif backend_name == "gguf":
        from .gguf_backend import GGUFBackend
        return GGUFBackend()
    else:
        raise ValueError(f"Unknown backend: {backend_name}. Available: {get_available_backends()}")