from __future__ import annotations

from typing import Any, Dict, Optional

from transformers import PretrainedConfig


class PTDQwen2Config(PretrainedConfig):
    model_type = "ptd_qwen2"

    def __init__(
        self,
        base_model: str = "Qwen/Qwen2.5-0.5B",
        tokenizer: Optional[str] = None,
        ptd_config: Optional[Dict[str, Any]] = None,
        package_type: str = "full_state",
        recommended_keep_rate: float = 0.7,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.base_model = base_model
        self.tokenizer = tokenizer or base_model
        self.ptd_config = ptd_config or {}
        self.package_type = package_type
        self.recommended_keep_rate = float(recommended_keep_rate)