dataclass config

myolmoe/modeling_myolmoe.py

@@ -18,102 +18,73 @@ from transformers.utils import logging
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_rope_utils import rope_config_validation
 
+from dataclasses import dataclass, field
+from typing import Optional, List, Any
+from transformers import PretrainedConfig
+
+@dataclass
 class MyOlmoeConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`OlmoeModel`].
-    [Previous docstring remains the same...]
-    
-    Args:
-        [Previous args remain the same...]
-        small_expert_intermediate_ratio (`float`, *optional*, defaults to 0.5):
-            Ratio of intermediate size for small experts compared to regular experts.
-        small_expert_count (`int`, *optional*, defaults to 64):
-            Frequency of small experts - every Nth expert will be small.
-        small_expert_sparsity_coef (`float`, *optional*, defaults to 0.1):
-            Coefficient for small expert load balancing loss.
     """
-    model_type = "myolmoe"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size=50304,
-        hidden_size=2048,
-        intermediate_size=2048,
-        num_hidden_layers=16,
-        num_attention_heads=16,
-        num_key_value_heads=None,
-        hidden_act="silu",
-        max_position_embeddings=4096,
-        initializer_range=0.02,
-        rms_norm_eps=1e-05,
-        use_cache=True,
-        pad_token_id=1,
-        bos_token_id=None,
-        eos_token_id=50279,
-        tie_word_embeddings=False,
-        rope_theta=10000.0,
-        rope_scaling=None,
-        attention_bias=False,
-        attention_dropout=0.0,
-        clip_qkv=None,
-        num_experts_per_tok=8,
-        num_experts=64,
-        output_router_logits=False,
-        router_aux_loss_coef=0.01,
-        norm_topk_prob=False,
-        small_expert_intermediate_ratio=64,
-        small_expert_count=64,
-        small_expert_sparsity_coef=0.1,
-        small_expert_strategy="constant",  # increment
-        max_small_expert_count=64,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.attention_bias = attention_bias
-        self.attention_dropout = attention_dropout
-        self.clip_qkv = clip_qkv
-        self.num_experts_per_tok = num_experts_per_tok
-        self.num_experts = num_experts
-        self.output_router_logits = output_router_logits
-        self.router_aux_loss_coef = router_aux_loss_coef
-        self.norm_topk_prob = norm_topk_prob
-        
-        # Small expert parameters
-        self.small_expert_intermediate_ratio = small_expert_intermediate_ratio
-        self.small_expert_count = small_expert_count
-        self.small_expert_sparsity_coef = small_expert_sparsity_coef
-        self.small_expert_strategy = small_expert_strategy
-        self.max_small_expert_count = max_small_expert_count
-        
-        # Validate the correctness of rotary position embeddings parameters
-        if self.rope_scaling is not None and "type" in self.rope_scaling:
-            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
-        rope_config_validation(self)
-
+    Configuration class for MyOlmoe model.
+    """
+    model_type: str = "myolmoe"
+    
+    # Core model parameters
+    vocab_size: int = 50304
+    hidden_size: int = 2048
+    intermediate_size: int = 1024
+    num_hidden_layers: int = 16
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 16
+    max_position_embeddings: int = 4096
+    
+    # Expert parameters
+    num_experts: int = 64
+    num_experts_per_tok: int = 2
+    num_small_experts: int = 0
+    small_expert_count: int = 64
+    small_expert_intermediate_ratio: int = 16
+    small_expert_intermediate_size: int = 0
+    small_expert_sparsity_coef: float = 0.1
+    small_expert_strategy: str = "constant"
+    max_small_expert_count: int = 64
+    
+    # Attention parameters
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    clip_qkv: Optional[float] = None
+    
+    # Normalization and activation
+    hidden_act: str = "silu"
+    rms_norm_eps: float = 1e-05
+    norm_topk_prob: bool = False
+    
+    # Router parameters
+    router_aux_loss_coef: float = 0.01
+    output_router_logits: bool = False
+    
+    # Training parameters
+    initializer_range: float = 0.02
+    tie_word_embeddings: bool = False
+    use_cache: bool = True
+    
+    # RoPE parameters
+    rope_theta: float = 10000.0
+    rope_scaling: Optional[dict] = None
+    
+    # Token IDs
+    pad_token_id: int = 1
+    eos_token_id: int = 50279
+    
+    # Model architecture
+    architectures: List[str] = field(default_factory=lambda: ["MyOlmoeForCausalLM"])
+    
+    def __post_init__(self):
+        """Post-initialization to ensure compatibility with PretrainedConfig."""
         super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
+            pad_token_id=self.pad_token_id,
+            eos_token_id=self.eos_token_id,
+            **{k: v for k, v in self.__dict__.items() if not k.startswith('_')}
         )
 
 logger = logging.get_logger(__name__)