"""Model specifications and presets for stage models.

This module defines:
1. VlmModelSpec - Model configuration with engine-specific overrides
2. StageModelPreset - Preset combining model, engine, and stage config
3. StagePresetMixin - Mixin for stage options to manage presets
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Set

from pydantic import BaseModel, Field

from docling.datamodel.pipeline_options_vlm_model import (
    ResponseFormat,
    TransformersModelType,
    TransformersPromptStyle,
)
from docling.datamodel.vlm_engine_options import BaseVlmEngineOptions
from docling.datamodel.vlm_prompts import (
    CHANDRA_OCR_LAYOUT_PROMPT,
    DOTS_LAYOUT_PROMPT,
)
from docling.models.inference_engines.image_classification.base import (
    ImageClassificationEngineType,
)
from docling.models.inference_engines.object_detection.base import (
    ObjectDetectionEngineType,
)
from docling.models.inference_engines.vlm.base import VlmEngineType

if TYPE_CHECKING:
    from docling.datamodel.image_classification_engine_options import (
        BaseImageClassificationEngineOptions,
    )
    from docling.datamodel.object_detection_engine_options import (
        BaseObjectDetectionEngineOptions,
    )

_log = logging.getLogger(__name__)


# =============================================================================
# ENGINE-SPECIFIC MODEL CONFIGURATION
# =============================================================================


class EngineModelConfig(BaseModel):
    """Engine-specific model configuration.

    Allows overriding model settings for specific engines.
    For example, MLX might use a different repo_id than Transformers.
    """

    repo_id: str | None = Field(
        default=None, description="Override model repository ID for this engine"
    )

    revision: str | None = Field(
        default=None, description="Override model revision for this engine"
    )

    torch_dtype: str | None = Field(
        default=None,
        description="Override torch dtype for this engine (e.g., 'bfloat16')",
    )

    extra_config: Dict[str, Any] = Field(
        default_factory=dict, description="Additional engine-specific configuration"
    )

    def merge_with(
        self, base_repo_id: str, base_revision: str = "main"
    ) -> EngineModelConfig:
        """Merge with base configuration.

        Args:
            base_repo_id: Base repository ID
            base_revision: Base revision

        Returns:
            Merged configuration with overrides applied
        """
        return EngineModelConfig(
            repo_id=self.repo_id or base_repo_id,
            revision=self.revision or base_revision,
            torch_dtype=self.torch_dtype,
            extra_config=self.extra_config,
        )


class ApiModelConfig(BaseModel):
    """API-specific model configuration.

    For API engines, configuration is simpler - just params to send.
    """

    params: Dict[str, Any] = Field(
        default_factory=dict,
        description="API parameters (model name, max_tokens, etc.)",
    )

    def merge_with(self, base_params: Dict[str, Any]) -> ApiModelConfig:
        """Merge with base parameters.

        Args:
            base_params: Base API parameters

        Returns:
            Merged configuration with overrides applied
        """
        merged_params = {**base_params, **self.params}
        return ApiModelConfig(params=merged_params)


# =============================================================================
# VLM MODEL SPECIFICATION
# =============================================================================


class VlmModelSpec(BaseModel):
    """Specification for a VLM model.

    This defines the model configuration that is independent of the engine.
    It includes:
    - Default model repository ID
    - Prompt template
    - Response format
    - Engine-specific overrides
    """

    name: str = Field(description="Human-readable model name")

    default_repo_id: str = Field(description="Default HuggingFace repository ID")

    revision: str = Field(default="main", description="Default model revision")

    prompt: str = Field(description="Prompt template for this model")

    response_format: ResponseFormat = Field(
        description="Expected response format from the model"
    )

    supported_engines: Set[VlmEngineType] | None = Field(
        default=None, description="Set of supported engines (None = all supported)"
    )

    engine_overrides: Dict[VlmEngineType, EngineModelConfig] = Field(
        default_factory=dict, description="Engine-specific configuration overrides"
    )

    api_overrides: Dict[VlmEngineType, ApiModelConfig] = Field(
        default_factory=dict, description="API-specific configuration overrides"
    )

    trust_remote_code: bool = Field(
        default=False, description="Whether to trust remote code for this model"
    )

    stop_strings: List[str] = Field(
        default_factory=list, description="Stop strings for generation"
    )

    max_new_tokens: int = Field(
        default=4096, description="Maximum number of new tokens to generate"
    )

    temperature: float = Field(
        default=0.0,
        description="Sampling temperature for generation. 0.0 uses greedy decoding.",
    )

    def get_repo_id(self, engine_type: VlmEngineType) -> str:
        """Get the repository ID for a specific engine.

        Args:
            engine_type: The engine type

        Returns:
            Repository ID (with engine override if applicable)
        """
        if engine_type in self.engine_overrides:
            override = self.engine_overrides[engine_type]
            return override.repo_id or self.default_repo_id
        return self.default_repo_id

    def get_revision(self, engine_type: VlmEngineType) -> str:
        """Get the model revision for a specific engine.

        Args:
            engine_type: The engine type

        Returns:
            Model revision (with engine override if applicable)
        """
        if engine_type in self.engine_overrides:
            override = self.engine_overrides[engine_type]
            return override.revision or self.revision
        return self.revision

    def get_api_params(self, engine_type: VlmEngineType) -> Dict[str, Any]:
        """Get API parameters for a specific engine.

        Args:
            engine_type: The engine type

        Returns:
            API parameters (with engine override if applicable)
        """
        base_params = {"model": self.default_repo_id}

        if engine_type in self.api_overrides:
            override = self.api_overrides[engine_type]
            return override.merge_with(base_params).params

        return base_params

    def is_engine_supported(self, engine_type: VlmEngineType) -> bool:
        """Check if an engine is supported by this model.

        Args:
            engine_type: The engine type to check

        Returns:
            True if supported, False otherwise
        """
        if self.supported_engines is None:
            return True
        return engine_type in self.supported_engines

    def get_engine_config(self, engine_type: VlmEngineType) -> EngineModelConfig:
        """Get EngineModelConfig for a specific engine type.

        This is the single source of truth for generating engine-specific
        configuration from the model spec.

        Args:
            engine_type: The engine type to get config for

        Returns:
            EngineModelConfig with repo_id, revision, and engine-specific extra_config
        """
        # Get repo_id and revision (with engine-specific overrides if present)
        repo_id = self.get_repo_id(engine_type)
        revision = self.get_revision(engine_type)

        # Get engine-specific extra_config and torch_dtype
        extra_config = {}
        torch_dtype = None
        if engine_type in self.engine_overrides:
            extra_config = self.engine_overrides[engine_type].extra_config.copy()
            torch_dtype = self.engine_overrides[engine_type].torch_dtype

        return EngineModelConfig(
            repo_id=repo_id,
            revision=revision,
            torch_dtype=torch_dtype,
            extra_config=extra_config,
        )

    def has_explicit_engine_export(self, engine_type: VlmEngineType) -> bool:
        """Check if this model has an explicit export for the given engine.

        An explicit export means either:
        1. The engine has an entry in engine_overrides, OR
        2. The engine is explicitly listed in supported_engines (not None)

        This is used by auto_inline to determine if it should attempt to use
        a specific engine. For example, MLX should only be used if the model
        explicitly declares MLX support, either via engine_overrides or
        supported_engines.

        Args:
            engine_type: The engine type to check

        Returns:
            True if there's an explicit export, False otherwise

        Examples:
            >>> # Model with MLX export (different repo_id)
            >>> spec = VlmModelSpec(
            ...     name="Test",
            ...     default_repo_id="org/model",
            ...     engine_overrides={
            ...         VlmEngineType.MLX: EngineModelConfig(repo_id="org/model-mlx")
            ...     }
            ... )
            >>> spec.has_explicit_engine_export(VlmEngineType.MLX)
            True

            >>> # Model with same-repo MLX support (native handler, no separate weights)
            >>> spec = VlmModelSpec(
            ...     name="Test",
            ...     default_repo_id="org/model",
            ...     engine_overrides={VlmEngineType.MLX: EngineModelConfig()}
            ... )
            >>> spec.has_explicit_engine_export(VlmEngineType.MLX)
            True

            >>> # Model without MLX export
            >>> spec = VlmModelSpec(name="Test", default_repo_id="org/model")
            >>> spec.has_explicit_engine_export(VlmEngineType.MLX)
            False

            >>> # Model with explicit supported_engines
            >>> spec = VlmModelSpec(
            ...     name="Test",
            ...     default_repo_id="org/model",
            ...     supported_engines={VlmEngineType.MLX}
            ... )
            >>> spec.has_explicit_engine_export(VlmEngineType.MLX)
            True
        """
        # If supported_engines is explicitly set and includes this engine
        if self.supported_engines is not None:
            return engine_type in self.supported_engines

        return engine_type in self.engine_overrides


# =============================================================================
# OBJECT DETECTION MODEL SPECIFICATION
# =============================================================================


class ObjectDetectionModelSpec(BaseModel):
    """Specification for an object detection model.

    Simpler than VlmModelSpec - no prompts, no preprocessing params.
    Preprocessing comes from HuggingFace preprocessor configs.
    Model files are assumed to be at the root of the HuggingFace repo.
    """

    name: str = Field(description="Human-readable model name")

    repo_id: str = Field(description="Default HuggingFace repository ID")

    revision: str = Field(default="main", description="Default model revision")

    engine_overrides: Dict[ObjectDetectionEngineType, EngineModelConfig] = Field(
        default_factory=dict,
        description="Engine-specific configuration overrides",
    )

    def get_engine_config(
        self, engine_type: ObjectDetectionEngineType
    ) -> EngineModelConfig:
        """Get EngineModelConfig for a specific object-detection engine.

        Args:
            engine_type: The engine type being requested

        Returns:
            EngineModelConfig populated with repo/revision and engine overrides
        """
        override = self.engine_overrides.get(engine_type)
        if override is not None:
            return override.merge_with(self.repo_id, self.revision)
        return EngineModelConfig(repo_id=self.repo_id, revision=self.revision)

    def get_repo_id(self, engine_type: ObjectDetectionEngineType) -> str:
        """Get repository ID for specific engine.

        Args:
            engine_type: The engine type

        Returns:
            Repository ID (with engine override if applicable)
        """
        override = self.engine_overrides.get(engine_type)
        if override and override.repo_id:
            return override.repo_id
        return self.repo_id

    def get_revision(self, engine_type: ObjectDetectionEngineType) -> str:
        """Get revision for specific engine.

        Args:
            engine_type: The engine type

        Returns:
            Model revision (with engine override if applicable)
        """
        override = self.engine_overrides.get(engine_type)
        if override and override.revision:
            return override.revision
        return self.revision


# =============================================================================
# IMAGE CLASSIFICATION MODEL SPECIFICATION
# =============================================================================


class ImageClassificationModelSpec(BaseModel):
    """Specification for an image-classification model."""

    name: str = Field(description="Human-readable model name")

    repo_id: str = Field(description="Default HuggingFace repository ID")

    revision: str = Field(default="main", description="Default model revision")

    engine_overrides: Dict[ImageClassificationEngineType, EngineModelConfig] = Field(
        default_factory=dict,
        description="Engine-specific configuration overrides",
    )

    def get_engine_config(
        self, engine_type: ImageClassificationEngineType
    ) -> EngineModelConfig:
        """Get EngineModelConfig for a specific image-classification engine."""
        override = self.engine_overrides.get(engine_type)
        if override is not None:
            return override.merge_with(self.repo_id, self.revision)
        return EngineModelConfig(repo_id=self.repo_id, revision=self.revision)

    def get_repo_id(self, engine_type: ImageClassificationEngineType) -> str:
        """Get repository ID for specific engine."""
        override = self.engine_overrides.get(engine_type)
        if override and override.repo_id:
            return override.repo_id
        return self.repo_id

    def get_revision(self, engine_type: ImageClassificationEngineType) -> str:
        """Get revision for specific engine."""
        override = self.engine_overrides.get(engine_type)
        if override and override.revision:
            return override.revision
        return self.revision


# =============================================================================
# STAGE PRESET SYSTEM
# =============================================================================


class StageModelPreset(BaseModel):
    """A preset configuration combining stage, model, and prompt.

    Presets provide convenient named configurations that users can
    reference by ID instead of manually configuring everything.
    """

    preset_id: str = Field(
        description="Simple preset identifier (e.g., 'smolvlm', 'granite')"
    )

    name: str = Field(description="Human-readable preset name")

    description: str = Field(description="Description of what this preset does")

    model_spec: VlmModelSpec = Field(description="Model specification for this preset")

    scale: float = Field(default=2.0, description="Image scaling factor")

    max_size: int | None = Field(default=None, description="Maximum image dimension")

    default_engine_type: VlmEngineType = Field(
        default=VlmEngineType.AUTO_INLINE,
        description="Default engine to use with this preset",
    )

    stage_options: Dict[str, Any] = Field(
        default_factory=dict, description="Additional stage-specific options"
    )

    @property
    def supported_engines(self) -> Set[VlmEngineType]:
        """Get supported engines from model spec."""
        if self.model_spec.supported_engines is None:
            return set(VlmEngineType)
        return self.model_spec.supported_engines


class StagePresetMixin:
    """Mixin for stage options classes that support presets.

    Each stage options class that uses this mixin manages its own presets.
    This is more decentralized than a global registry.

    Usage:
        class MyStageOptions(StagePresetMixin, BaseModel):
            ...

        # Register presets
        MyStageOptions.register_preset(preset1)
        MyStageOptions.register_preset(preset2)

        # Use presets
        options = MyStageOptions.from_preset("preset1")
    """

    # Class variable to store presets for this specific stage
    # Note: Each subclass gets its own _presets dict via __init_subclass__
    _presets: ClassVar[Dict[str, StageModelPreset]]

    def __init_subclass__(cls, **kwargs):
        """Initialize each subclass with its own preset registry.

        This ensures that each stage options class has an isolated preset
        registry, preventing namespace collisions across different stages.
        """
        super().__init_subclass__(**kwargs)
        # Each subclass gets its own _presets dictionary
        cls._presets = {}

    @classmethod
    def register_preset(cls, preset: StageModelPreset) -> None:
        """Register a preset for this stage options class.

        Args:
            preset: The preset to register

        Note:
            If preset ID already registered, it will be silently skipped.
            This allows for idempotent registration at module import time.
        """
        if preset.preset_id not in cls._presets:
            cls._presets[preset.preset_id] = preset
        else:
            _log.error(
                f"Preset '{preset.preset_id}' already registered for {cls.__name__}"
            )

    @classmethod
    def get_preset(cls, preset_id: str) -> StageModelPreset:
        """Get a specific preset.

        Args:
            preset_id: The preset identifier

        Returns:
            The requested preset

        Raises:
            KeyError: If preset not found
        """
        if preset_id not in cls._presets:
            raise KeyError(
                f"Preset '{preset_id}' not found for {cls.__name__}. "
                f"Available presets: {list(cls._presets.keys())}"
            )
        return cls._presets[preset_id]

    @classmethod
    def list_presets(cls) -> List[StageModelPreset]:
        """List all presets for this stage.

        Returns:
            List of presets
        """
        return list(cls._presets.values())

    @classmethod
    def list_preset_ids(cls) -> List[str]:
        """List all preset IDs for this stage.

        Returns:
            List of preset IDs
        """
        return list(cls._presets.keys())

    @classmethod
    def get_preset_info(cls) -> List[Dict[str, str]]:
        """Get summary info for all presets (useful for CLI).

        Returns:
            List of dicts with preset_id, name, description, model
        """
        return [
            {
                "preset_id": p.preset_id,
                "name": p.name,
                "description": p.description,
                "model": p.model_spec.name,
                "default_engine": p.default_engine_type.value,
            }
            for p in cls._presets.values()
        ]

    @classmethod
    def from_preset(
        cls,
        preset_id: str,
        engine_options: BaseVlmEngineOptions | None = None,
        **overrides,
    ):
        """Create options from a registered preset.

        Args:
            preset_id: The preset identifier
            engine_options: Optional engine override
            **overrides: Additional option overrides

        Returns:
            Instance of the stage options class
        """
        from docling.datamodel.vlm_engine_options import (
            ApiVlmEngineOptions,
            AutoInlineVlmEngineOptions,
            MlxVlmEngineOptions,
            TransformersVlmEngineOptions,
            VllmVlmEngineOptions,
        )

        preset = cls.get_preset(preset_id)

        # Create engine options if not provided
        trust_remote = preset.model_spec.trust_remote_code
        if engine_options is None:
            if preset.default_engine_type == VlmEngineType.AUTO_INLINE:
                engine_options = AutoInlineVlmEngineOptions()
            elif VlmEngineType.is_api_variant(preset.default_engine_type):
                engine_options = ApiVlmEngineOptions(
                    engine_type=preset.default_engine_type
                )
            elif preset.default_engine_type == VlmEngineType.TRANSFORMERS:
                engine_options = TransformersVlmEngineOptions(
                    trust_remote_code=trust_remote,
                )
            elif preset.default_engine_type == VlmEngineType.MLX:
                engine_options = MlxVlmEngineOptions(
                    trust_remote_code=trust_remote,
                )
            elif preset.default_engine_type == VlmEngineType.VLLM:
                engine_options = VllmVlmEngineOptions(
                    trust_remote_code=trust_remote,
                )
            else:
                engine_options = AutoInlineVlmEngineOptions()

        # Create instance with preset values
        # Type ignore because cls is the concrete options class, not the mixin
        instance = cls(  # type: ignore[call-arg]
            model_spec=preset.model_spec,
            engine_options=engine_options,
            scale=preset.scale,
            max_size=preset.max_size,
            **preset.stage_options,
        )

        # Apply overrides
        for key, value in overrides.items():
            setattr(instance, key, value)

        return instance


class ObjectDetectionStagePreset(BaseModel):
    """Preset definition for object detection-powered stages."""

    preset_id: str = Field(description="Preset identifier")
    name: str = Field(description="Human-readable preset name")
    description: str = Field(description="Description of this preset")
    model_spec: ObjectDetectionModelSpec = Field(
        description="Object detection model specification"
    )
    default_engine_type: ObjectDetectionEngineType = Field(
        default=ObjectDetectionEngineType.ONNXRUNTIME,
        description="Default inference engine to use",
    )
    stage_options: Dict[str, Any] = Field(
        default_factory=dict, description="Additional stage-specific defaults"
    )


class ObjectDetectionStagePresetMixin:
    """Mixin to enable preset loading for object detection stages."""

    _presets: ClassVar[Dict[str, ObjectDetectionStagePreset]]

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        cls._presets = {}

    @classmethod
    def register_preset(cls, preset: ObjectDetectionStagePreset) -> None:
        if preset.preset_id not in cls._presets:
            cls._presets[preset.preset_id] = preset
        else:
            _log.error(
                f"Preset '{preset.preset_id}' already registered for {cls.__name__}"
            )

    @classmethod
    def get_preset(cls, preset_id: str) -> ObjectDetectionStagePreset:
        if preset_id not in cls._presets:
            raise KeyError(
                f"Preset '{preset_id}' not found for {cls.__name__}. "
                f"Available presets: {list(cls._presets.keys())}"
            )
        return cls._presets[preset_id]

    @classmethod
    def list_presets(cls) -> List[ObjectDetectionStagePreset]:
        return list(cls._presets.values())

    @classmethod
    def list_preset_ids(cls) -> List[str]:
        return list(cls._presets.keys())

    @classmethod
    def get_preset_info(cls) -> List[Dict[str, str]]:
        return [
            {
                "preset_id": p.preset_id,
                "name": p.name,
                "description": p.description,
                "model": p.model_spec.name,
                "default_engine": p.default_engine_type.value,
            }
            for p in cls._presets.values()
        ]

    @classmethod
    def from_preset(
        cls,
        preset_id: str,
        engine_options: BaseObjectDetectionEngineOptions | None = None,
        **overrides: Any,
    ):
        from docling.datamodel.object_detection_engine_options import (
            ApiKserveV2ObjectDetectionEngineOptions,
            OnnxRuntimeObjectDetectionEngineOptions,
            TransformersObjectDetectionEngineOptions,
        )

        preset = cls.get_preset(preset_id)

        if engine_options is None:
            if preset.default_engine_type == ObjectDetectionEngineType.ONNXRUNTIME:
                engine_options = OnnxRuntimeObjectDetectionEngineOptions()
            elif preset.default_engine_type == ObjectDetectionEngineType.TRANSFORMERS:
                engine_options = TransformersObjectDetectionEngineOptions()
            elif preset.default_engine_type == ObjectDetectionEngineType.API_KSERVE_V2:
                raise ValueError(
                    f"Preset '{preset_id}' uses API_KSERVE_V2 engine which requires explicit "
                    "engine_options with a 'url' parameter. Please provide "
                    "engine_options=ApiKserveV2ObjectDetectionEngineOptions(url='...') "
                    "when calling from_preset()."
                )
            else:
                raise ValueError(
                    f"Unsupported engine type {preset.default_engine_type} for presets"
                )

        instance = cls(  # type: ignore[call-arg]
            model_spec=preset.model_spec,
            engine_options=engine_options,
            **preset.stage_options,
        )

        for key, value in overrides.items():
            setattr(instance, key, value)

        return instance


class ImageClassificationStagePreset(BaseModel):
    """Preset definition for image classification-powered stages."""

    preset_id: str = Field(description="Preset identifier")
    name: str = Field(description="Human-readable preset name")
    description: str = Field(description="Description of this preset")
    model_spec: ImageClassificationModelSpec = Field(
        description="Image classification model specification"
    )
    default_engine_type: ImageClassificationEngineType = Field(
        default=ImageClassificationEngineType.TRANSFORMERS,
        description="Default inference engine to use",
    )
    stage_options: Dict[str, Any] = Field(
        default_factory=dict, description="Additional stage-specific defaults"
    )


class ImageClassificationStagePresetMixin:
    """Mixin to enable preset loading for image-classification stages."""

    _presets: ClassVar[Dict[str, ImageClassificationStagePreset]]

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        cls._presets = {}

    @classmethod
    def register_preset(cls, preset: ImageClassificationStagePreset) -> None:
        if preset.preset_id not in cls._presets:
            cls._presets[preset.preset_id] = preset
        else:
            _log.error(
                f"Preset '{preset.preset_id}' already registered for {cls.__name__}"
            )

    @classmethod
    def get_preset(cls, preset_id: str) -> ImageClassificationStagePreset:
        if preset_id not in cls._presets:
            raise KeyError(
                f"Preset '{preset_id}' not found for {cls.__name__}. "
                f"Available presets: {list(cls._presets.keys())}"
            )
        return cls._presets[preset_id]

    @classmethod
    def list_presets(cls) -> List[ImageClassificationStagePreset]:
        return list(cls._presets.values())

    @classmethod
    def list_preset_ids(cls) -> List[str]:
        return list(cls._presets.keys())

    @classmethod
    def get_preset_info(cls) -> List[Dict[str, str]]:
        return [
            {
                "preset_id": p.preset_id,
                "name": p.name,
                "description": p.description,
                "model": p.model_spec.name,
                "default_engine": p.default_engine_type.value,
            }
            for p in cls._presets.values()
        ]

    @classmethod
    def from_preset(
        cls,
        preset_id: str,
        engine_options: BaseImageClassificationEngineOptions | None = None,
        **overrides: Any,
    ):
        from docling.datamodel.image_classification_engine_options import (
            ApiKserveV2ImageClassificationEngineOptions,
            OnnxRuntimeImageClassificationEngineOptions,
            TransformersImageClassificationEngineOptions,
        )

        preset = cls.get_preset(preset_id)

        if engine_options is None:
            if preset.default_engine_type == ImageClassificationEngineType.ONNXRUNTIME:
                engine_options = OnnxRuntimeImageClassificationEngineOptions()
            elif (
                preset.default_engine_type == ImageClassificationEngineType.TRANSFORMERS
            ):
                engine_options = TransformersImageClassificationEngineOptions()
            elif (
                preset.default_engine_type
                == ImageClassificationEngineType.API_KSERVE_V2
            ):
                raise ValueError(
                    f"Preset '{preset_id}' uses API_KSERVE_V2 engine which requires explicit "
                    "engine_options with a 'url' parameter. Please provide "
                    "engine_options=ApiKserveV2ImageClassificationEngineOptions(url='...') "
                    "when calling from_preset()."
                )
            else:
                raise ValueError(
                    f"Unsupported engine type {preset.default_engine_type} for presets"
                )

        instance = cls(  # type: ignore[call-arg]
            model_spec=preset.model_spec,
            engine_options=engine_options,
            **preset.stage_options,
        )

        for key, value in overrides.items():
            setattr(instance, key, value)

        return instance


# =============================================================================
# PRESET DEFINITIONS
# =============================================================================

# -----------------------------------------------------------------------------
# SHARED MODEL SPECS (for reuse across multiple stages)
# -----------------------------------------------------------------------------

# Shared Granite Docling model spec used across VLM_CONVERT and CODE_FORMULA stages
# Note: prompt and response_format are intentionally excluded here as they vary per stage
GRANITE_DOCLING_MODEL_SPEC_BASE = {
    "name": "Granite-Docling-258M",
    "default_repo_id": "ibm-granite/granite-docling-258M",
    "stop_strings": ["</doctag>", "<|end_of_text|>"],
    "max_new_tokens": 8192,
    "engine_overrides": {
        VlmEngineType.MLX: EngineModelConfig(
            repo_id="ibm-granite/granite-docling-258M-mlx"
        ),
        VlmEngineType.TRANSFORMERS: EngineModelConfig(
            extra_config={
                "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                "extra_generation_config": {"skip_special_tokens": False},
            }
        ),
    },
    "api_overrides": {
        VlmEngineType.API_OLLAMA: ApiModelConfig(
            params={"model": "ibm/granite-docling:258m"}
        ),
    },
}

# Shared Pixtral model spec used across VLM_CONVERT and PICTURE_DESCRIPTION stages
PIXTRAL_MODEL_SPEC_BASE = {
    "name": "Pixtral-12B",
    "default_repo_id": "mistral-community/pixtral-12b",
    "engine_overrides": {
        VlmEngineType.MLX: EngineModelConfig(repo_id="mlx-community/pixtral-12b-bf16"),
        VlmEngineType.TRANSFORMERS: EngineModelConfig(
            extra_config={
                "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
            }
        ),
    },
}

# Shared Granite Vision model spec used across VLM_CONVERT and PICTURE_DESCRIPTION stages
GRANITE_VISION_MODEL_SPEC_BASE = {
    "name": "Granite-Vision-3.3-2B",
    "default_repo_id": "ibm-granite/granite-vision-3.3-2b",
    "supported_engines": {
        VlmEngineType.TRANSFORMERS,
        VlmEngineType.VLLM,
        VlmEngineType.API_OLLAMA,
        VlmEngineType.API_LMSTUDIO,
    },
    "engine_overrides": {
        VlmEngineType.TRANSFORMERS: EngineModelConfig(
            extra_config={
                "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
            }
        ),
    },
    "api_overrides": {
        VlmEngineType.API_OLLAMA: ApiModelConfig(
            params={"model": "granite3.3-vision:2b"}
        ),
    },
}

# -----------------------------------------------------------------------------
# OBJECT DETECTION PRESETS
# -----------------------------------------------------------------------------

OBJECT_DETECTION_LAYOUT_HERON = ObjectDetectionStagePreset(
    preset_id="layout_heron_default",
    name="Layout Heron",
    description="RT-DETR layout-heron model (ResNet50)",
    model_spec=ObjectDetectionModelSpec(
        name="layout_heron",
        repo_id="docling-project/docling-layout-heron",
        revision="main",
        engine_overrides={
            ObjectDetectionEngineType.ONNXRUNTIME: EngineModelConfig(
                repo_id="docling-project/docling-layout-heron-onnx",
                extra_config={"model_filename": "model.onnx"},
            )
        },
    ),
    default_engine_type=ObjectDetectionEngineType.TRANSFORMERS,
)


# -----------------------------------------------------------------------------
# IMAGE CLASSIFICATION PRESETS
# -----------------------------------------------------------------------------

IMAGE_CLASSIFICATION_DOCUMENT_FIGURE = ImageClassificationStagePreset(
    preset_id="document_figure_classifier_v2",
    name="Document Figure Classifier v2",
    description="EfficientNet model for classifying document pictures",
    model_spec=ImageClassificationModelSpec(
        name="document_figure_classifier_v2",
        repo_id="docling-project/DocumentFigureClassifier-v2.5",
        revision="main",
    ),
    default_engine_type=ImageClassificationEngineType.TRANSFORMERS,
)


# -----------------------------------------------------------------------------
# VLM_CONVERT PRESETS (for full page conversion)
# -----------------------------------------------------------------------------

VLM_CONVERT_SMOLDOCLING = StageModelPreset(
    preset_id="smoldocling",
    name="SmolDocling",
    description="Lightweight DocTags model optimized for document conversion (256M parameters)",
    model_spec=VlmModelSpec(
        name="SmolDocling-256M",
        default_repo_id="docling-project/SmolDocling-256M-preview",
        prompt="Convert this page to docling.",
        response_format=ResponseFormat.DOCTAGS,
        stop_strings=["</doctag>", "<end_of_utterance>"],
        engine_overrides={
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="docling-project/SmolDocling-256M-preview-mlx-bf16"
            ),
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                },
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_GRANITE_DOCLING = StageModelPreset(
    preset_id="granite_docling",
    name="Granite-Docling",
    description="IBM Granite DocTags model for document conversion (258M parameters)",
    model_spec=VlmModelSpec(
        **GRANITE_DOCLING_MODEL_SPEC_BASE,
        prompt="Convert this page to docling.",
        response_format=ResponseFormat.DOCTAGS,
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_DEEPSEEK_OCR = StageModelPreset(
    preset_id="deepseek_ocr",
    name="DeepSeek-OCR",
    description="DeepSeek OCR model via Ollama/LM Studio for document conversion (3B parameters)",
    model_spec=VlmModelSpec(
        name="DeepSeek-OCR-3B",
        default_repo_id="deepseek-ocr:3b",  # Ollama model name
        prompt="<|grounding|>Convert the document to markdown. ",
        response_format=ResponseFormat.DEEPSEEKOCR_MARKDOWN,
        supported_engines={VlmEngineType.API_OLLAMA, VlmEngineType.API_LMSTUDIO},
        api_overrides={
            VlmEngineType.API_OLLAMA: ApiModelConfig(
                params={"model": "deepseek-ocr:3b", "max_tokens": 4096}
            ),
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "deepseek-ocr", "max_tokens": 4096}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.API_OLLAMA,
)

VLM_CONVERT_GRANITE_VISION = StageModelPreset(
    preset_id="granite_vision",
    name="Granite-Vision",
    description="IBM Granite Vision model for markdown conversion (2B parameters)",
    model_spec=VlmModelSpec(
        **GRANITE_VISION_MODEL_SPEC_BASE,
        prompt="Convert this page to markdown. Do not miss any text and only output the bare markdown!",
        response_format=ResponseFormat.MARKDOWN,
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_PIXTRAL = StageModelPreset(
    preset_id="pixtral",
    name="Pixtral-12B",
    description="Mistral Pixtral model for markdown conversion (12B parameters)",
    model_spec=VlmModelSpec(
        **PIXTRAL_MODEL_SPEC_BASE,
        prompt="Convert this page to markdown. Do not miss any text and only output the bare markdown!",
        response_format=ResponseFormat.MARKDOWN,
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_GOT_OCR = StageModelPreset(
    preset_id="got_ocr",
    name="GOT-OCR-2.0",
    description="GOT OCR 2.0 model for markdown conversion",
    model_spec=VlmModelSpec(
        name="GOT-OCR-2.0",
        default_repo_id="stepfun-ai/GOT-OCR-2.0-hf",
        prompt="",
        response_format=ResponseFormat.MARKDOWN,
        supported_engines={VlmEngineType.TRANSFORMERS},
        stop_strings=["<|im_end|>"],
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                    "transformers_prompt_style": TransformersPromptStyle.NONE,
                    "extra_processor_kwargs": {"format": True},
                }
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.TRANSFORMERS,
)

VLM_CONVERT_PHI4 = StageModelPreset(
    preset_id="phi4",
    name="Phi-4",
    description="Microsoft Phi-4 multimodal model for markdown conversion",
    model_spec=VlmModelSpec(
        name="Phi-4-Multimodal-Instruct",
        default_repo_id="microsoft/Phi-4-multimodal-instruct",
        prompt="Convert this page to MarkDown. Do not miss any text and only output the bare markdown",
        response_format=ResponseFormat.MARKDOWN,
        trust_remote_code=True,
        supported_engines={
            VlmEngineType.TRANSFORMERS,
            VlmEngineType.VLLM,
        },
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_CAUSALLM,
                    "extra_generation_config": {"num_logits_to_keep": 0},
                }
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_QWEN = StageModelPreset(
    preset_id="qwen",
    name="Qwen2.5-VL-3B",
    description="Qwen vision-language model for markdown conversion (3B parameters)",
    model_spec=VlmModelSpec(
        name="Qwen2.5-VL-3B-Instruct",
        default_repo_id="Qwen/Qwen2.5-VL-3B-Instruct",
        prompt="Convert this page to markdown. Do not miss any text and only output the bare markdown!",
        response_format=ResponseFormat.MARKDOWN,
        engine_overrides={
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="mlx-community/Qwen2.5-VL-3B-Instruct-bf16"
            ),
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                }
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_NANONETS_OCR2 = StageModelPreset(
    preset_id="nanonets_ocr2",
    name="Nanonets-OCR2-3B",
    description="Nanonets OCR2 model for text recognition and markdown conversion (3B parameters)",
    model_spec=VlmModelSpec(
        name="Nanonets-OCR2-3B",
        default_repo_id="nanonets/Nanonets-OCR2-3B",
        prompt=(
            "Extract the text from the above document as if you were reading it naturally. "
            "Return the tables in html format. Return the equations in LaTeX representation. "
            "If there is an image in the document and image caption is not present, add a "
            "small description of the image inside the <img></img> tag; otherwise, add the "
            "image caption inside <img></img>. Watermarks should be wrapped in brackets. "
            "Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in "
            "brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. "
            "Prefer using ☐ and ☑ for check boxes."
        ),
        response_format=ResponseFormat.MARKDOWN,
        max_new_tokens=15000,
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                    "transformers_prompt_style": TransformersPromptStyle.CHAT,
                    "torch_dtype": "bfloat16",
                },
            ),
            # MLX uses a qwen2_5_vl-compatible converted checkpoint.
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="mlx-community/Nanonets-OCR2-3B-bf16"
            ),
        },
        api_overrides={
            VlmEngineType.API: ApiModelConfig(
                params={"model": "nanonets/Nanonets-OCR2-3B", "max_tokens": 15000}
            ),
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "nanonets-ocr2-3b", "max_tokens": 15000}
            ),
            VlmEngineType.API_OPENAI: ApiModelConfig(
                params={"model": "nanonets-ocr2-3b", "max_tokens": 15000}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_GEMMA_12B = StageModelPreset(
    preset_id="gemma_12b",
    name="Gemma-3-12B",
    description="Google Gemma-3 vision model for markdown conversion (12B parameters)",
    model_spec=VlmModelSpec(
        name="Gemma-3-12B-IT",
        default_repo_id="google/gemma-3-12b-it",
        prompt="Convert this page to markdown. Do not miss any text and only output the bare markdown!",
        response_format=ResponseFormat.MARKDOWN,
        supported_engines={VlmEngineType.MLX},
        engine_overrides={
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="mlx-community/gemma-3-12b-it-bf16"
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.MLX,
)

VLM_CONVERT_GEMMA_27B = StageModelPreset(
    preset_id="gemma_27b",
    name="Gemma-3-27B",
    description="Google Gemma-3 vision model for markdown conversion (27B parameters)",
    model_spec=VlmModelSpec(
        name="Gemma-3-27B-IT",
        default_repo_id="google/gemma-3-27b-it",
        prompt="Convert this page to markdown. Do not miss any text and only output the bare markdown!",
        response_format=ResponseFormat.MARKDOWN,
        supported_engines={VlmEngineType.MLX},
        engine_overrides={
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="mlx-community/gemma-3-27b-it-bf16"
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.MLX,
)

VLM_CONVERT_DOLPHIN = StageModelPreset(
    preset_id="dolphin",
    name="Dolphin",
    description="ByteDance Dolphin OCR model for markdown conversion",
    model_spec=VlmModelSpec(
        name="Dolphin",
        default_repo_id="ByteDance/Dolphin",
        prompt="<s>Read text in the image. <Answer/>",
        response_format=ResponseFormat.MARKDOWN,
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                    "transformers_prompt_style": TransformersPromptStyle.RAW,
                }
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_GLMOCR = StageModelPreset(
    preset_id="glm_ocr",
    name="GLM-OCR",
    description="Zhipu GLM-OCR model for text recognition and markdown conversion (0.9B parameters)",
    model_spec=VlmModelSpec(
        name="GLM-OCR-0.9B",
        default_repo_id="zai-org/GLM-OCR",
        # GLM-OCR supports three prompts: "Text Recognition:", "Formula Recognition:",
        # "Table Recognition:". We use text-only for full-page document conversion.
        prompt="Text Recognition:",
        response_format=ResponseFormat.MARKDOWN,
        stop_strings=["<|user|>", "<|endoftext|>"],
        engine_overrides={
            # Native GLM-OCR support was added to mlx-vlm in v0.3.11.
            VlmEngineType.MLX: EngineModelConfig(repo_id="mlx-community/GLM-OCR-bf16"),
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                    "transformers_prompt_style": TransformersPromptStyle.CHAT,
                    "transformers_strip_stop_strings": True,
                    "torch_dtype": "bfloat16",
                },
            ),
            VlmEngineType.VLLM: EngineModelConfig(
                extra_config={
                    "enforce_eager": True,
                },
            ),
        },
        api_overrides={
            VlmEngineType.API: ApiModelConfig(
                params={"model": "zai-org/GLM-OCR", "max_tokens": 4096}
            ),
            VlmEngineType.API_OPENAI: ApiModelConfig(
                params={"model": "glm-ocr", "max_tokens": 4096}
            ),
            VlmEngineType.API_OLLAMA: ApiModelConfig(
                params={"model": "glm-ocr", "max_tokens": 4096}
            ),
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "glm-ocr", "max_tokens": 4096}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_FALCON_OCR = StageModelPreset(
    preset_id="falcon_ocr",
    name="Falcon-OCR",
    description="TII Falcon-OCR model for text recognition and markdown conversion (0.3B parameters)",
    model_spec=VlmModelSpec(
        name="Falcon-OCR",
        default_repo_id="tiiuae/Falcon-OCR",
        prompt="",
        response_format=ResponseFormat.MARKDOWN,
        trust_remote_code=True,
        engine_overrides={
            # Native Falcon-OCR support was added to mlx-vlm in v0.4.3.
            # A dedicated mlx-community checkpoint is now available for MLX.
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="mlx-community/Falcon-OCR-bf16"
            ),
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_CAUSALLM,
                    "transformers_prompt_style": TransformersPromptStyle.CHAT,
                    "torch_dtype": "bfloat16",
                },
            ),
        },
        api_overrides={
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "falcon-ocr", "max_tokens": 4096}
            ),
            VlmEngineType.API_OPENAI: ApiModelConfig(
                params={"model": "falcon-ocr", "max_tokens": 4096}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

VLM_CONVERT_LIGHTONOCR = StageModelPreset(
    preset_id="lightonocr",
    name="LightOnOCR-2-1B",
    description="LightOn LightOnOCR-2 model for OCR and markdown conversion (1B parameters)",
    model_spec=VlmModelSpec(
        name="LightOnOCR-2-1B",
        default_repo_id="lightonai/LightOnOCR-2-1B",
        prompt="",
        response_format=ResponseFormat.MARKDOWN,
        max_new_tokens=4096,
        engine_overrides={
            # LightOnOCR currently runs on mlx-vlm via the generic mistral3 handler.
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="mlx-community/LightOnOCR-2-1B-bf16"
            ),
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                    "transformers_prompt_style": TransformersPromptStyle.CHAT,
                    "torch_dtype": "bfloat16",
                },
            ),
        },
        api_overrides={
            VlmEngineType.API: ApiModelConfig(
                params={"model": "lightonai/LightOnOCR-2-1B", "max_tokens": 4096}
            ),
            VlmEngineType.API_OPENAI: ApiModelConfig(
                params={"model": "lightonocr-2-1b", "max_tokens": 4096}
            ),
        },
    ),
    scale=2.0,
    max_size=1540,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

# -----------------------------------------------------------------------------
# PICTURE_DESCRIPTION PRESETS (for image captioning/description)
# -----------------------------------------------------------------------------

PICTURE_DESC_SMOLVLM = StageModelPreset(
    preset_id="smolvlm",
    name="SmolVLM-256M",
    description="Lightweight vision-language model for image descriptions (256M parameters)",
    model_spec=VlmModelSpec(
        name="SmolVLM-256M-Instruct",
        default_repo_id="HuggingFaceTB/SmolVLM-256M-Instruct",
        prompt="Describe this image in a few sentences.",
        response_format=ResponseFormat.PLAINTEXT,
        engine_overrides={
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="moot20/SmolVLM-256M-Instruct-MLX"
            ),
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                },
            ),
        },
        api_overrides={
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "smolvlm-256m-instruct"}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
    stage_options={
        "picture_area_threshold": 0.05,
    },
)

PICTURE_DESC_GRANITE_VISION = StageModelPreset(
    preset_id="granite_vision",
    name="Granite-Vision-3.3-2B",
    description="IBM Granite Vision model for detailed image descriptions (2B parameters)",
    model_spec=VlmModelSpec(
        **GRANITE_VISION_MODEL_SPEC_BASE,
        prompt="What is shown in this image?",
        response_format=ResponseFormat.PLAINTEXT,
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
    stage_options={
        "picture_area_threshold": 0.05,
    },
)

PICTURE_DESC_PIXTRAL = StageModelPreset(
    preset_id="pixtral",
    name="Pixtral-12B",
    description="Mistral Pixtral model for detailed image descriptions (12B parameters)",
    model_spec=VlmModelSpec(
        **PIXTRAL_MODEL_SPEC_BASE,
        prompt="Describe this image in detail.",
        response_format=ResponseFormat.PLAINTEXT,
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
    stage_options={
        "picture_area_threshold": 0.05,
    },
)

PICTURE_DESC_QWEN = StageModelPreset(
    preset_id="qwen",
    name="Qwen2.5-VL-3B",
    description="Qwen vision-language model for image descriptions (3B parameters)",
    model_spec=VlmModelSpec(
        name="Qwen2.5-VL-3B-Instruct",
        default_repo_id="Qwen/Qwen2.5-VL-3B-Instruct",
        prompt="Describe this image.",
        response_format=ResponseFormat.PLAINTEXT,
        engine_overrides={
            VlmEngineType.MLX: EngineModelConfig(
                repo_id="mlx-community/Qwen2.5-VL-3B-Instruct-bf16"
            ),
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                }
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
    stage_options={
        "picture_area_threshold": 0.05,
    },
)

# -----------------------------------------------------------------------------
# CODE_FORMULA PRESETS (for code and formula extraction)
# -----------------------------------------------------------------------------

CODE_FORMULA_CODEFORMULAV2 = StageModelPreset(
    preset_id="codeformulav2",
    name="CodeFormulaV2",
    description="Specialized model for code and formula extraction",
    model_spec=VlmModelSpec(
        name="CodeFormulaV2",
        default_repo_id="docling-project/CodeFormulaV2",
        prompt="",
        response_format=ResponseFormat.PLAINTEXT,
        stop_strings=["</doctag>", "<end_of_utterance>"],
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                    "extra_generation_config": {"skip_special_tokens": False},
                    "torch_dtype": "bfloat16",
                },
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

CODE_FORMULA_GRANITE_DOCLING = StageModelPreset(
    preset_id="granite_docling",
    name="Granite-Docling-CodeFormula",
    description="IBM Granite Docling model for code and formula extraction (258M parameters)",
    model_spec=VlmModelSpec(
        **GRANITE_DOCLING_MODEL_SPEC_BASE,
        prompt="",
        response_format=ResponseFormat.PLAINTEXT,
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

# -----------------------------------------------------------------------------
# CHANDRA / DOTS VLM_CONVERT PRESETS
# -----------------------------------------------------------------------------

VLM_CONVERT_CHANDRA_OCR2 = StageModelPreset(
    preset_id="chandra_ocr2",
    name="Chandra-OCR-2",
    description="Chandra OCR 2 model for document layout parsing with bounding boxes (5.3B parameters)",
    model_spec=VlmModelSpec(
        name="Chandra-OCR-2-5.3B",
        default_repo_id="datalab-to/chandra-ocr-2",
        prompt=CHANDRA_OCR_LAYOUT_PROMPT,
        response_format=ResponseFormat.CHANDRA_HTML,
        max_new_tokens=12384,
        trust_remote_code=True,
        stop_strings=["<|im_end|>", "<|endoftext|>"],
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT,
                    "transformers_prompt_style": TransformersPromptStyle.CHAT,
                    "transformers_strip_stop_strings": True,
                },
            ),
            VlmEngineType.VLLM: EngineModelConfig(
                extra_config={
                    "enforce_eager": True,
                },
            ),
        },
        api_overrides={
            VlmEngineType.API_OPENAI: ApiModelConfig(
                params={"model": "datalab-to/chandra-ocr-2", "max_tokens": 12384}
            ),
            VlmEngineType.API_OLLAMA: ApiModelConfig(
                params={"model": "chandra-ocr-2", "max_tokens": 12384}
            ),
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "chandra-ocr-2", "max_tokens": 12384}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.AUTO_INLINE,
)

# NOTE: dots.ocr/mocr require transformers<=4.57.x for the transformers engine.
# Under transformers 5.x, generation produces incorrect coordinates (single fullpage
# bbox). The vllm engine path is unaffected.
VLM_CONVERT_DOTS_OCR = StageModelPreset(
    preset_id="dots_ocr",
    name="Dots-OCR",
    description="dots.ocr model for multilingual document layout parsing with bounding boxes (3B parameters)",
    model_spec=VlmModelSpec(
        name="dots.ocr-3B",
        default_repo_id="rednote-hilab/dots.ocr",
        prompt=DOTS_LAYOUT_PROMPT,
        response_format=ResponseFormat.DOTS_JSON,
        max_new_tokens=24000,
        trust_remote_code=True,
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_CAUSALLM,
                    "transformers_prompt_style": TransformersPromptStyle.CHAT,
                },
            ),
            VlmEngineType.VLLM: EngineModelConfig(
                extra_config={
                    "enforce_eager": True,
                },
            ),
        },
        api_overrides={
            VlmEngineType.API_OPENAI: ApiModelConfig(
                params={"model": "rednote-hilab/dots.ocr", "max_tokens": 24000}
            ),
            VlmEngineType.API_OLLAMA: ApiModelConfig(
                params={"model": "dots.ocr", "max_tokens": 24000}
            ),
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "dots.ocr", "max_tokens": 24000}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.VLLM,
)

VLM_CONVERT_DOTS_MOCR = StageModelPreset(
    preset_id="dots_mocr",
    name="Dots-MOCR",
    description="dots.mocr multimodal OCR model for document layout parsing with bounding boxes (3B parameters)",
    model_spec=VlmModelSpec(
        name="dots.mocr-3B",
        default_repo_id="rednote-hilab/dots.mocr",
        prompt=DOTS_LAYOUT_PROMPT,
        response_format=ResponseFormat.DOTS_JSON,
        max_new_tokens=24000,
        trust_remote_code=True,
        engine_overrides={
            VlmEngineType.TRANSFORMERS: EngineModelConfig(
                torch_dtype="bfloat16",
                extra_config={
                    "transformers_model_type": TransformersModelType.AUTOMODEL_CAUSALLM,
                    "transformers_prompt_style": TransformersPromptStyle.CHAT,
                },
            ),
            VlmEngineType.VLLM: EngineModelConfig(
                extra_config={
                    "enforce_eager": True,
                },
            ),
        },
        api_overrides={
            VlmEngineType.API_OPENAI: ApiModelConfig(
                params={"model": "rednote-hilab/dots.mocr", "max_tokens": 24000}
            ),
            VlmEngineType.API_OLLAMA: ApiModelConfig(
                params={"model": "dots.mocr", "max_tokens": 24000}
            ),
            VlmEngineType.API_LMSTUDIO: ApiModelConfig(
                params={"model": "dots.mocr", "max_tokens": 24000}
            ),
        },
    ),
    scale=2.0,
    default_engine_type=VlmEngineType.VLLM,
)
