"""BaguettotronVLM configuration."""
from __future__ import annotations
from transformers import PretrainedConfig


class BaguettotronVLMConfig(PretrainedConfig):
    model_type = "baguettotron_vlm"

    def __init__(
        self,
        vit_model_id: str = "OpenGVLab/InternViT-300M-448px-V2_5",
        llm_model_id: str = "PleIAs/Baguettotron",
        vit_hidden: int = 1024,
        vit_tokens: int = 1024,
        llm_hidden: int = 576,
        num_visual_tokens: int = 256,
        unshuffle_factor: int = 2,
        image_token: str = "<image>",
        chat_style: str = "answer",
        stage: int = 2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vit_model_id = vit_model_id
        self.llm_model_id = llm_model_id
        self.vit_hidden = vit_hidden
        self.vit_tokens = vit_tokens
        self.llm_hidden = llm_hidden
        self.num_visual_tokens = num_visual_tokens
        self.unshuffle_factor = unshuffle_factor
        self.image_token = image_token
        # chat_style controls the assistant-turn prefix emitted by the
        # processor when add_generation_prompt=True:
        #   "base"   → <|im_start|>assistant\n        (stage 1, no think tokens)
        #   "answer" → <|im_start|>assistant\n</think>\n (stage 2, answer-only)
        #   "think"  → <|im_start|>assistant\n<think>\n  (stage 3, reasoning)
        self.chat_style = chat_style
        self.stage = stage