Upload PersadianNanoV4Model.py with huggingface_hub

PersadianNanoV4Model.py

@@ -1,11 +1,12 @@
 # persadian_nano_v4/PersadianNanoV4Model.py
-# Stable GPU-ready implementation
-# persadian-Nano-V4 (~160M total params)
+# Stable GPU-ready implementation with Hugging Face compatibility
 
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
 
 # ============================================================
 # CONFIG
@@ -13,110 +14,81 @@ import torch.nn.functional as F
 
 class PersadianNanoV4Config:
     def __init__(self):
-
-        # Core architecture
+        # Core architecture (MATCH CHECKPOINT)
         self.hidden_size = 512
         self.intermediate_size = 1024
-        self.num_hidden_layers = 6
-
+        self.num_hidden_layers = 12
+        
         # Attention
         self.num_attention_heads = 8
         self.num_key_value_heads = 4
-
+        
         # MoE
         self.num_experts = 4
         self.num_experts_per_tok = 2
-
+        
         # Progressive experts
         self.progressive_experts = True
         self.min_experts = 1
         self.max_experts = 2
-
+        
         # Hyper connections
         self.use_hyper_connection = True
-
+        
         # Compression
         self.use_compressed_attention = True
         self.compress_ratio = 4
-
+        
         # Context
         self.max_position_embeddings = 2048
-
+        
         # Tokenizer
         self.vocab_size = 50257
         self.bos_token_id = 50256
         self.eos_token_id = 50256
-
+        
         # RoPE
         self.rope_theta = 10000.0
-
+        
         # Regularization
         self.dropout = 0.1
         self.layer_norm_eps = 1e-5
-
+        
         # Attention
         self.attention_dropout = 0.0
         self.use_flash_attention = True
-
+        
         # Precision
         self.torch_dtype = "float16"
 
-
 # ============================================================
 # ROTARY EMBEDDING
 # ============================================================
 
 class RotaryEmbedding(nn.Module):
-
     def __init__(self, dim, base=10000.0):
         super().__init__()
-
-        inv_freq = 1.0 / (
-            base ** (
-                torch.arange(0, dim, 2).float() / dim
-            )
-        )
-
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
         self.register_buffer("inv_freq", inv_freq)
 
     def forward(self, x, position_ids):
-
-        # x shape:
-        # [batch, seq, heads, head_dim]
-
         batch, seq_len, num_heads, head_dim = x.shape
-
-        freqs = torch.einsum(
-            "bi,j->bij",
-            position_ids.float(),
-            self.inv_freq
-        )
-
+        freqs = torch.einsum("bi,j->bij", position_ids.float(), self.inv_freq)
         emb = torch.cat([freqs, freqs], dim=-1)
-
         cos = emb.cos()[:, :, None, :]
         sin = emb.sin()[:, :, None, :]
-
         x1 = x[..., ::2]
         x2 = x[..., 1::2]
-
-        rotated = torch.stack(
-            (-x2, x1),
-            dim=-1
-        ).flatten(-2)
-
+        rotated = torch.stack((-x2, x1), dim=-1).flatten(-2)
         return (x * cos) + (rotated * sin)
 
-
 # ============================================================
 # ADAPTIVE HYPER CONNECTION
 # ============================================================
 
 class AdaptiveHyperConnection(nn.Module):
-
     def __init__(self, hidden_size):
         super().__init__()
-
         self.router = nn.Sequential(
             nn.Linear(hidden_size, hidden_size // 4),
             nn.SiLU(),
@@ -125,726 +97,263 @@ class AdaptiveHyperConnection(nn.Module):
         )
 
     def forward(self, x, residual):
-
-        weight = self.router(
-            x.mean(dim=1, keepdim=True)
-        )
-
+        weight = self.router(x.mean(dim=1, keepdim=True))
         return residual + (x * weight)
 
-
 # ============================================================
 # ATTENTION
 # ============================================================
 
 class CompressedSparseAttention(nn.Module):
-
     def __init__(self, config):
         super().__init__()
-
         self.config = config
-
         self.hidden_size = config.hidden_size
-
         self.num_heads = config.num_attention_heads
         self.num_key_value_heads = config.num_key_value_heads
-
-        self.head_dim = (
-            config.hidden_size //
-            config.num_attention_heads
-        )
-
-        # Q projection
-        self.q_proj = nn.Linear(
-            config.hidden_size,
-            config.hidden_size,
-            bias=False
-        )
-
-        # KV projection
-        self.k_proj = nn.Linear(
-            config.hidden_size,
-            self.num_key_value_heads * self.head_dim,
-            bias=False
-        )
-
-        self.v_proj = nn.Linear(
-            config.hidden_size,
-            self.num_key_value_heads * self.head_dim,
-            bias=False
-        )
-
-        # Output
-        self.o_proj = nn.Linear(
-            config.hidden_size,
-            config.hidden_size,
-            bias=False
-        )
-
-        # Compression
-        compressed_dim = (
-            config.hidden_size //
-            config.compress_ratio
-        )
-
-        self.compressor = nn.Linear(
-            config.hidden_size,
-            compressed_dim
-        )
-
-        self.k_proj_compressed = nn.Linear(
-            compressed_dim,
-            self.num_key_value_heads * self.head_dim,
-            bias=False
-        )
-
-        self.v_proj_compressed = nn.Linear(
-            compressed_dim,
-            self.num_key_value_heads * self.head_dim,
-            bias=False
-        )
-
-        # Rotary
-        self.rotary_emb = RotaryEmbedding(
-            self.head_dim,
-            config.rope_theta
-        )
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask=None,
-        position_ids=None,
-    ):
-
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        
+        compressed_dim = config.hidden_size // config.compress_ratio
+        self.compressor = nn.Linear(config.hidden_size, compressed_dim)
+        self.k_proj_compressed = nn.Linear(compressed_dim, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj_compressed = nn.Linear(compressed_dim, self.num_key_value_heads * self.head_dim, bias=False)
+        
+        self.rotary_emb = RotaryEmbedding(self.head_dim, config.rope_theta)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None):
         batch_size, seq_len, _ = hidden_states.shape
-
-        # ====================================================
-        # Q
-        # ====================================================
-
-        q = self.q_proj(hidden_states)
-
-        q = q.view(
-            batch_size,
-            seq_len,
-            self.num_heads,
-            self.head_dim
-        )
-
-        # ====================================================
-        # COMPRESSED OR NORMAL KV
-        # ====================================================
-
-        if (
-            self.config.use_compressed_attention and
-            seq_len > 512
-        ):
-
+        
+        q = self.q_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim)
+        
+        if self.config.use_compressed_attention and seq_len > 512:
             compressed = self.compressor(hidden_states)
-
             k = self.k_proj_compressed(compressed)
             v = self.v_proj_compressed(compressed)
-
         else:
-
             k = self.k_proj(hidden_states)
             v = self.v_proj(hidden_states)
-
-        k = k.view(
-            batch_size,
-            seq_len,
-            self.num_key_value_heads,
-            self.head_dim
-        )
-
-        v = v.view(
-            batch_size,
-            seq_len,
-            self.num_key_value_heads,
-            self.head_dim
-        )
-
-        # ====================================================
-        # GQA EXPANSION
-        # ====================================================
-
-        repeat_factor = (
-            self.num_heads //
-            self.num_key_value_heads
-        )
-
-        k = k.repeat_interleave(
-            repeat_factor,
-            dim=2
-        )
-
-        v = v.repeat_interleave(
-            repeat_factor,
-            dim=2
-        )
-
-        # ====================================================
-        # ROPE
-        # ====================================================
-
+        
+        k = k.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
+        v = v.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
+        
+        repeat_factor = self.num_heads // self.num_key_value_heads
+        k = k.repeat_interleave(repeat_factor, dim=2)
+        v = v.repeat_interleave(repeat_factor, dim=2)
+        
         q = self.rotary_emb(q, position_ids)
         k = self.rotary_emb(k, position_ids)
-
-        # ====================================================
-        # TRANSPOSE
-        # ====================================================
-
+        
         q = q.transpose(1, 2)
         k = k.transpose(1, 2)
         v = v.transpose(1, 2)
-
-        # ====================================================
-        # FLASH ATTENTION
-        # ====================================================
-
-        if (
-            self.config.use_flash_attention and
-            hasattr(F, "scaled_dot_product_attention")
-        ):
-
-            attn_output = F.scaled_dot_product_attention(
-                q,
-                k,
-                v,
-                attn_mask=attention_mask,
-                dropout_p=self.config.attention_dropout,
-                is_causal=True
-            )
-
+        
+        if self.config.use_flash_attention and hasattr(F, "scaled_dot_product_attention"):
+            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, is_causal=True)
         else:
-
-            scores = torch.matmul(
-                q,
-                k.transpose(-2, -1)
-            ) / math.sqrt(self.head_dim)
-
+            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
             if attention_mask is not None:
                 scores = scores + attention_mask
-
-            probs = F.softmax(scores, dim=-1)
-
-            attn_output = torch.matmul(
-                probs,
-                v
-            )
-
-        # ====================================================
-        # OUTPUT
-        # ====================================================
-
-        attn_output = attn_output.transpose(1, 2)
-
-        attn_output = attn_output.contiguous().view(
-            batch_size,
-            seq_len,
-            self.hidden_size
-        )
-
+            attn_output = torch.matmul(F.softmax(scores, dim=-1), v)
+        
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_size)
         return self.o_proj(attn_output)
 
-
 # ============================================================
 # MIXTURE OF EXPERTS
 # ============================================================
 
 class MixtureOfExperts(nn.Module):
-
     def __init__(self, config):
         super().__init__()
-
         self.config = config
-
         self.num_experts = config.num_experts
-
+        
         self.experts = nn.ModuleList([
             nn.Sequential(
-                nn.Linear(
-                    config.hidden_size,
-                    config.intermediate_size
-                ),
+                nn.Linear(config.hidden_size, config.intermediate_size),
                 nn.SiLU(),
-                nn.Linear(
-                    config.intermediate_size,
-                    config.hidden_size
-                ),
+                nn.Linear(config.intermediate_size, config.hidden_size),
                 nn.Dropout(config.dropout)
-            )
-            for _ in range(config.num_experts)
+            ) for _ in range(config.num_experts)
         ])
-
-        self.router = nn.Linear(
-            config.hidden_size,
-            config.num_experts,
-            bias=False
-        )
-
+        
+        self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False)
         self.aux_loss_coef = 0.01
 
-    def forward(
-        self,
-        hidden_states,
-        progressive_factor=1.0
-    ):
-
+    def forward(self, hidden_states, progressive_factor=1.0):
         batch_size, seq_len, hidden_size = hidden_states.shape
-
-        hidden_states_flat = hidden_states.view(
-            -1,
-            hidden_size
-        )
-
-        router_logits = self.router(
-            hidden_states_flat
-        )
-
-        routing_weights = F.softmax(
-            router_logits,
-            dim=-1
-        )
-
-        # ====================================================
-        # PROGRESSIVE EXPERT ACTIVATION
-        # ====================================================
-
+        hidden_states_flat = hidden_states.view(-1, hidden_size)
+        router_logits = self.router(hidden_states_flat)
+        routing_weights = F.softmax(router_logits, dim=-1)
+        
         if self.config.progressive_experts:
-
-            k = int(
-                round(
-                    self.config.min_experts +
-                    (
-                        self.config.max_experts -
-                        self.config.min_experts
-                    ) * progressive_factor
-                )
-            )
-
-            k = max(
-                self.config.min_experts,
-                min(k, self.config.max_experts)
-            )
-
+            k = int(round(self.config.min_experts + (self.config.max_experts - self.config.min_experts) * progressive_factor))
+            k = max(self.config.min_experts, min(k, self.config.max_experts))
         else:
             k = self.config.num_experts_per_tok
-
-        # ====================================================
-        # TOP-K ROUTING
-        # ====================================================
-
-        top_k_weights, top_k_indices = torch.topk(
-            routing_weights,
-            k,
-            dim=-1
-        )
-
-        top_k_weights = (
-            top_k_weights /
-            top_k_weights.sum(
-                dim=-1,
-                keepdim=True
-            )
-        )
-
-        # ====================================================
-        # EXPERT COMPUTE
-        # ====================================================
-
-        final_hidden = torch.zeros_like(
-            hidden_states_flat
-        )
-
+        
+        top_k_weights, top_k_indices = torch.topk(routing_weights, k, dim=-1)
+        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
+        
+        final_hidden = torch.zeros_like(hidden_states_flat)
         for expert_idx in range(self.num_experts):
-
-            expert_mask = (
-                top_k_indices == expert_idx
-            ).any(dim=-1)
-
+            expert_mask = (top_k_indices == expert_idx).any(dim=-1)
             if expert_mask.any():
-
-                expert_input = hidden_states_flat[
-                    expert_mask
-                ]
-
-                expert_output = self.experts[
-                    expert_idx
-                ](expert_input)
-
-                expert_weights = top_k_weights[
-                    expert_mask
-                ].mean(dim=-1, keepdim=True)
-
-                final_hidden[expert_mask] += (
-                    expert_output * expert_weights
-                )
-
-        # ====================================================
-        # AUX LOSS
-        # ====================================================
-
+                expert_input = hidden_states_flat[expert_mask]
+                expert_output = self.experts[expert_idx](expert_input)
+                expert_weights = top_k_weights[expert_mask].mean(dim=-1, keepdim=True)
+                final_hidden[expert_mask] += expert_output * expert_weights
+        
         router_probs = routing_weights.mean(dim=0)
-
         aux_loss = torch.var(router_probs)
-
-        return (
-            final_hidden.view(
-                batch_size,
-                seq_len,
-                hidden_size
-            ),
-            aux_loss * self.aux_loss_coef
-        )
-
+        
+        return final_hidden.view(batch_size, seq_len, hidden_size), aux_loss * self.aux_loss_coef
 
 # ============================================================
 # DECODER LAYER
 # ============================================================
 
 class PersadianNanoV4DecoderLayer(nn.Module):
-
     def __init__(self, config, layer_idx):
         super().__init__()
-
         self.config = config
-
-        self.input_layernorm = nn.LayerNorm(
-            config.hidden_size,
-            eps=config.layer_norm_eps
-        )
-
-        self.self_attn = CompressedSparseAttention(
-            config
-        )
-
-        self.post_attention_layernorm = nn.LayerNorm(
-            config.hidden_size,
-            eps=config.layer_norm_eps
-        )
-
+        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.self_attn = CompressedSparseAttention(config)
+        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.moe = MixtureOfExperts(config)
-
+        
         if config.use_hyper_connection:
+            self.hc_attention = AdaptiveHyperConnection(config.hidden_size)
+            self.hc_moe = AdaptiveHyperConnection(config.hidden_size)
 
-            self.hc_attention = AdaptiveHyperConnection(
-                config.hidden_size
-            )
-
-            self.hc_moe = AdaptiveHyperConnection(
-                config.hidden_size
-            )
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask=None,
-        position_ids=None,
-        progressive_factor=1.0
-    ):
-
-        # ====================================================
-        # ATTENTION
-        # ====================================================
-
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, progressive_factor=1.0):
         residual = hidden_states
-
-        hidden_states = self.input_layernorm(
-            hidden_states
-        )
-
-        attn_output = self.self_attn(
-            hidden_states,
-            attention_mask,
-            position_ids
-        )
-
+        hidden_states = self.input_layernorm(hidden_states)
+        attn_output = self.self_attn(hidden_states, attention_mask, position_ids)
+        
         if self.config.use_hyper_connection:
-
-            hidden_states = self.hc_attention(
-                attn_output,
-                residual
-            )
-
+            hidden_states = self.hc_attention(attn_output, residual)
         else:
             hidden_states = residual + attn_output
-
-        # ====================================================
-        # MOE
-        # ====================================================
-
+        
         residual = hidden_states
-
-        hidden_states = self.post_attention_layernorm(
-            hidden_states
-        )
-
-        moe_output, aux_loss = self.moe(
-            hidden_states,
-            progressive_factor
-        )
-
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        moe_output, aux_loss = self.moe(hidden_states, progressive_factor)
+        
         if self.config.use_hyper_connection:
-
-            hidden_states = self.hc_moe(
-                moe_output,
-                residual
-            )
-
+            hidden_states = self.hc_moe(moe_output, residual)
         else:
             hidden_states = residual + moe_output
-
+        
         return hidden_states, aux_loss
 
-
 # ============================================================
-# MAIN MODEL
+# MAIN MODEL (Original)
 # ============================================================
 
 class PersadianNanoV4Model(nn.Module):
-
     def __init__(self, config):
         super().__init__()
-
         self.config = config
-
-        self.embed_tokens = nn.Embedding(
-            config.vocab_size,
-            config.hidden_size
-        )
-
-        self.layers = nn.ModuleList([
-            PersadianNanoV4DecoderLayer(
-                config,
-                i
-            )
-            for i in range(config.num_hidden_layers)
-        ])
-
-        self.norm = nn.LayerNorm(
-            config.hidden_size,
-            eps=config.layer_norm_eps
-        )
-
-        self.lm_head = nn.Linear(
-            config.hidden_size,
-            config.vocab_size,
-            bias=False
-        )
-
-        # Tie weights
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([PersadianNanoV4DecoderLayer(config, i) for i in range(config.num_hidden_layers)])
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.lm_head.weight = self.embed_tokens.weight
-
-        # Init
         self.apply(self._init_weights)
-
+    
     def _init_weights(self, module):
-
         if isinstance(module, nn.Linear):
-
-            nn.init.normal_(
-                module.weight,
-                mean=0.0,
-                std=0.02
-            )
-
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
             if module.bias is not None:
                 nn.init.zeros_(module.bias)
-
         elif isinstance(module, nn.Embedding):
-
-            nn.init.normal_(
-                module.weight,
-                mean=0.0,
-                std=0.02
-            )
-
-    def forward(
-        self,
-        input_ids,
-        attention_mask=None,
-        progressive_factor=1.0
-    ):
-
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+    
+    def forward(self, input_ids, attention_mask=None, progressive_factor=1.0):
         batch_size, seq_len = input_ids.shape
-
-        hidden_states = self.embed_tokens(
-            input_ids
-        )
-
-        position_ids = torch.arange(
-            seq_len,
-            device=input_ids.device
-        ).unsqueeze(0)
-
-        # ====================================================
-        # CAUSAL MASK
-        # ====================================================
-
+        hidden_states = self.embed_tokens(input_ids)
+        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
+        
         if attention_mask is None:
-
-            mask = torch.triu(
-                torch.full(
-                    (seq_len, seq_len),
-                    float("-inf"),
-                    device=input_ids.device
-                ),
-                diagonal=1
-            )
-
+            mask = torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=input_ids.device), diagonal=1)
             attention_mask = mask.unsqueeze(0).unsqueeze(0)
-
-        total_aux_loss = torch.tensor(
-            0.0,
-            device=input_ids.device
-        )
-
+        
+        total_aux_loss = torch.tensor(0.0, device=input_ids.device)
         for layer in self.layers:
-
-            hidden_states, aux_loss = layer(
-                hidden_states,
-                attention_mask,
-                position_ids,
-                progressive_factor
-            )
-
+            hidden_states, aux_loss = layer(hidden_states, attention_mask, position_ids, progressive_factor)
             total_aux_loss += aux_loss
-
+        
         hidden_states = self.norm(hidden_states)
-
         logits = self.lm_head(hidden_states)
-
         return logits, total_aux_loss
-
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids,
-        max_new_tokens=50,
-        temperature=0.7,
-        top_k=50
-    ):
-
+    
+    def generate(self, input_ids, max_new_tokens=50, temperature=0.7, top_k=50):
         self.eval()
-
         for _ in range(max_new_tokens):
-
             logits, _ = self.forward(input_ids)
-
-            next_token_logits = (
-                logits[:, -1, :] / temperature
-            )
-
+            next_token_logits = logits[:, -1, :] / temperature
             if top_k > 0:
-
-                values, _ = torch.topk(
-                    next_token_logits,
-                    top_k
-                )
-
+                values, _ = torch.topk(next_token_logits, top_k)
                 min_values = values[:, -1].unsqueeze(-1)
-
-                next_token_logits = torch.where(
-                    next_token_logits < min_values,
-                    torch.full_like(
-                        next_token_logits,
-                        float("-inf")
-                    ),
-                    next_token_logits
-                )
-
-            probs = F.softmax(
-                next_token_logits,
-                dim=-1
-            )
-
-            next_token = torch.multinomial(
-                probs,
-                num_samples=1
-            )
-
-            input_ids = torch.cat(
-                [input_ids, next_token],
-                dim=1
-            )
-
-            if (
-                next_token.item() ==
-                self.config.eos_token_id
-            ):
+                next_token_logits = torch.where(next_token_logits < min_values, torch.full_like(next_token_logits, float("-inf")), next_token_logits)
+            probs = F.softmax(next_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+            if next_token.item() == self.config.eos_token_id:
                 break
-
         return input_ids
 
-
 # ============================================================
-# PARAMETER COUNT
+# HF COMPATIBLE WRAPPER
 # ============================================================
 
-def count_parameters(model):
-
-    return sum(
-        p.numel()
-        for p in model.parameters()
-        if p.requires_grad
-    )
-
-
-# ============================================================
-# TEST
-# ============================================================
-
-if __name__ == "__main__":
-
-    print("=" * 60)
-    print("Initializing Persadian-Nano-V4")
-    print("=" * 60)
-
-    config = PersadianNanoV4Config()
-
-    model = PersadianNanoV4Model(config)
-
-    total_params = count_parameters(model)
-
-    print(f"Total Parameters: {total_params:,}")
-    print(f"Approx Size: {total_params / 1e6:.2f}M")
-
-    device = (
-        "cuda"
-        if torch.cuda.is_available()
-        else "cpu"
-    )
-
-    print(f"Device: {device}")
-
-    model = model.to(device)
-
-    dummy_input = torch.randint(
-        0,
-        config.vocab_size,
-        (1, 128)
-    ).to(device)
-
-    print("\nRunning forward pass...")
-
-    with torch.no_grad():
-
-        logits, aux_loss = model(dummy_input)
-
-    print("Forward pass successful")
-    print(f"Logits Shape: {logits.shape}")
-    print(f"Aux Loss: {aux_loss.item():.6f}")
-
-    print("\nTesting generation...")
-
-    generated = model.generate(
-        dummy_input[:, :16],
-        max_new_tokens=10
-    )
-
-    print(f"Generated Shape: {generated.shape}")
-
-    print("\nPersadian-Nano-V4 is operational")
-    print("=" * 60)
+class PersadianNanoV4ConfigHF(PretrainedConfig):
+    """HF-compatible config class"""
+    model_type = "persadian_nano_v4"
+    
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+class PersadianNanoV4ForCausalLM(PreTrainedModel):
+    """HF-compatible model class"""
+    config_class = PersadianNanoV4ConfigHF
+    
+    def __init__(self, config):
+        super().__init__(config)
+        # Create the original config object
+        original_config = PersadianNanoV4Config()
+        
+        # Copy all attributes from HF config to original config
+        for key, value in config.__dict__.items():
+            if hasattr(original_config, key):
+                setattr(original_config, key, value)
+        
+        self.model = PersadianNanoV4Model(original_config)
+        
+    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        logits, aux_loss = self.model(input_ids, attention_mask)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+        
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+        )
+    
+    def generate(self, input_ids, **kwargs):
+        return self.model.generate(input_ids, **kwargs)