Upload PersadianNanoV4Model.py with huggingface_hub

PersadianNanoV4Model.py

@@ -1,11 +1,10 @@
-# persadian_nano_v4/PersadianNanoV4Model.py
-# Stable GPU-ready implementation with Hugging Face compatibility
+# Simple Hugging Face compatible wrapper without complex inheritance
 
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PreTrainedModel, PretrainedConfig
+from transformers import PretrainedConfig
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 # ============================================================
@@ -14,51 +13,28 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
 
 class PersadianNanoV4Config:
     def __init__(self):
-        # Core architecture (MATCH CHECKPOINT)
         self.hidden_size = 512
         self.intermediate_size = 1024
         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.max_position_embeddings = 2048
         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"
 
 # ============================================================
@@ -127,7 +103,6 @@ class CompressedSparseAttention(nn.Module):
 
     def forward(self, hidden_states, attention_mask=None, position_ids=None):
         batch_size, seq_len, _ = hidden_states.shape
-        
         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:
@@ -172,7 +147,6 @@ class MixtureOfExperts(nn.Module):
         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),
@@ -181,7 +155,6 @@ class MixtureOfExperts(nn.Module):
                 nn.Dropout(config.dropout)
             ) for _ in range(config.num_experts)
         ])
-        
         self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False)
         self.aux_loss_coef = 0.01
 
@@ -211,7 +184,6 @@ class MixtureOfExperts(nn.Module):
         
         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
 
 # ============================================================
@@ -226,7 +198,6 @@ class PersadianNanoV4DecoderLayer(nn.Module):
         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)
@@ -253,7 +224,7 @@ class PersadianNanoV4DecoderLayer(nn.Module):
         return hidden_states, aux_loss
 
 # ============================================================
-# MAIN MODEL (Original)
+# MAIN MODEL
 # ============================================================
 
 class PersadianNanoV4Model(nn.Module):
@@ -301,7 +272,9 @@ class PersadianNanoV4Model(nn.Module):
             if top_k > 0:
                 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)
+                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)
@@ -310,11 +283,10 @@ class PersadianNanoV4Model(nn.Module):
         return input_ids
 
 # ============================================================
-# HF CONFIG WRAPPER
+# HF COMPATIBLE WRAPPER (SIMPLE VERSION)
 # ============================================================
 
 class PersadianNanoV4ConfigHF(PretrainedConfig):
-    """HF-compatible config class"""
     model_type = "persadian_nano_v4"
     
     def __init__(self, **kwargs):
@@ -322,59 +294,69 @@ class PersadianNanoV4ConfigHF(PretrainedConfig):
         for key, value in kwargs.items():
             setattr(self, key, value)
 
-# ============================================================
-# HF MODEL WRAPPER (FIXED)
-# ============================================================
-
-class PersadianNanoV4ForCausalLM(PreTrainedModel):
-    """HF-compatible model class"""
+class PersadianNanoV4ForCausalLM(nn.Module):
+    """Simple HF-compatible wrapper - no complex inheritance"""
     config_class = PersadianNanoV4ConfigHF
     
-    # Add required attributes for modern Transformers
-    _tied_weights_keys = ["lm_head.weight"]
-    _no_split_modules = ["PersadianNanoV4DecoderLayer"]
-    
     def __init__(self, config):
-        super().__init__(config)
-        # Create the original config object
+        super().__init__()
+        # Create original config
         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.config = original_config
         self.model = PersadianNanoV4Model(original_config)
-        self.post_init()
     
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """Load model from Hugging Face hub"""
+        from transformers import AutoConfig
+        import torch
+        
+        # Load config
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        
+        # Create model
+        model = cls(config)
+        
+        # Load weights
+        import os
+        from safetensors.torch import load_file
+        
+        # Try to load weights
+        weight_files = [
+            f"{pretrained_model_name_or_path}/pytorch_model.bin",
+            f"{pretrained_model_name_or_path}/model.safetensors"
+        ]
+        
+        for weight_file in weight_files:
+            if os.path.exists(weight_file):
+                if weight_file.endswith('.safetensors'):
+                    state_dict = load_file(weight_file)
+                else:
+                    state_dict = torch.load(weight_file, map_location='cpu')
+                model.load_state_dict(state_dict, strict=False)
+                break
+        
+        return model
     
-    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+    def forward(self, input_ids, attention_mask=None, labels=None):
         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 prepare_inputs_for_generation(self, input_ids, **kwargs):
-        return {"input_ids": input_ids}
+            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)
+    
+    def eval(self):
+        self.model.eval()
+        return self
+    
+    def to(self, device):
+        self.model = self.model.to(device)
+        return self