Add HuggingFace transformers integration (AutoModelForCausalLM support)

configuration_eve.py

@@ -1,24 +1,55 @@
+"""
+Eve-2-MoE Configuration
+========================
+HuggingFace-compatible configuration for the Eve-2-MoE architecture.
+
+Usage:
+    from transformers import AutoConfig
+    config = AutoConfig.from_pretrained("anthonym21/Eve-2-MoE-272M", trust_remote_code=True)
+"""
 
 from transformers import PretrainedConfig
 
+
 class EveConfig(PretrainedConfig):
-    model_type = "eve_moe"
-    
+    """Configuration for the Eve-2-MoE model.
+
+    This is a DeepSeek-V3 style Mixture of Experts architecture with a shared
+    expert, top-k routed experts, RoPE positional encoding, and SwiGLU activations.
+
+    Args:
+        vocab_size: Vocabulary size (padded for efficiency). Default: 50304.
+        n_layer: Number of transformer blocks. Default: 12.
+        n_embd: Hidden dimension / embedding size. Default: 512.
+        n_head: Number of attention heads. Default: 8.
+        head_dim: Dimension per attention head. Default: 64.
+        block_size: Maximum sequence length (context window). Default: 2048.
+        num_experts: Number of routed MoE experts. Default: 8.
+        top_k: Number of experts activated per token. Default: 2.
+        expert_intermediate_size: FFN hidden dim for each expert (SwiGLU). Default: 1408.
+        shared_expert_intermediate_size: FFN hidden dim for the shared expert. Default: 1408.
+        router_aux_loss_coef: Weight of the load-balancing auxiliary loss. Default: 0.01.
+        rope_theta: Base frequency for RoPE. Default: 10000.0.
+        use_checkpointing: Enable gradient checkpointing to save VRAM. Default: False.
+    """
+
+    model_type = "eve-moe"
+
     def __init__(
         self,
-        vocab_size=50304,
-        n_layer=12,
-        n_embd=512,
-        n_head=8,
-        head_dim=64,
-        block_size=2048,
-        num_experts=8,
-        top_k=2,
-        expert_intermediate_size=1408,
-        shared_expert_intermediate_size=1408,
-        router_aux_loss_coef=0.01,
-        use_checkpointing=False,
-        rope_theta=10000.0,
+        vocab_size: int = 50304,
+        n_layer: int = 12,
+        n_embd: int = 512,
+        n_head: int = 8,
+        head_dim: int = 64,
+        block_size: int = 2048,
+        num_experts: int = 8,
+        top_k: int = 2,
+        expert_intermediate_size: int = 1408,
+        shared_expert_intermediate_size: int = 1408,
+        router_aux_loss_coef: float = 0.01,
+        rope_theta: float = 10000.0,
+        use_checkpointing: bool = False,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -32,6 +63,10 @@ class EveConfig(PretrainedConfig):
         self.expert_intermediate_size = expert_intermediate_size
         self.shared_expert_intermediate_size = shared_expert_intermediate_size
         self.router_aux_loss_coef = router_aux_loss_coef
-        self.use_checkpointing = use_checkpointing
         self.rope_theta = rope_theta
+        self.use_checkpointing = use_checkpointing
+
+        # Default tie_word_embeddings to True (Eve-2 ties embedding + lm_head)
+        kwargs.setdefault("tie_word_embeddings", True)
+
         super().__init__(**kwargs)