Fix model registration for AutoModel compatibility

__init__.py

@@ -0,0 +1,31 @@
+"""NanoGPT HuggingFace Integration"""
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+# Import our classes
+try:
+    from .configuration_nanogpt import NanoGPTConfig
+    from .modeling_nanogpt import (
+        NanoGPTModel,
+        NanoGPTForCausalLM,
+        NanoGPTPreTrainedModel
+    )
+except ImportError:
+    from configuration_nanogpt import NanoGPTConfig
+    from modeling_nanogpt import (
+        NanoGPTModel,
+        NanoGPTForCausalLM,
+        NanoGPTPreTrainedModel
+    )
+
+# Register the model with Auto* classes
+AutoConfig.register("nanogpt", NanoGPTConfig)
+AutoModel.register(NanoGPTConfig, NanoGPTModel)
+AutoModelForCausalLM.register(NanoGPTConfig, NanoGPTForCausalLM)
+
+__all__ = [
+    "NanoGPTConfig",
+    "NanoGPTModel", 
+    "NanoGPTForCausalLM",
+    "NanoGPTPreTrainedModel"
+]

configuration_nanogpt.py

@@ -1,8 +1,13 @@
 """NanoGPT model configuration"""
 
 from transformers import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
 
 class NanoGPTConfig(PretrainedConfig):
+    """Configuration for NanoGPT model"""
+    
     model_type = "nanogpt"
     
     def __init__(

modeling_nanogpt.py

@@ -1,4 +1,4 @@
-"""NanoGPT model implementation"""
+"""NanoGPT model implementation for HuggingFace"""
 
 import torch
 import torch.nn as nn
@@ -6,7 +6,15 @@ import torch.nn.functional as F
 import math
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
-from .configuration_nanogpt import NanoGPTConfig
+from transformers.utils import logging
+
+# Import configuration
+try:
+    from .configuration_nanogpt import NanoGPTConfig
+except ImportError:
+    from configuration_nanogpt import NanoGPTConfig
+
+logger = logging.get_logger(__name__)
 
 class ExactNanoGPTAttention(nn.Module):
     def __init__(self, config):
@@ -80,8 +88,26 @@ class ExactNanoGPTBlock(nn.Module):
         x = x + self.mlp(self.ln_2(x))
         return x
 
-class NanoGPTModel(PreTrainedModel):
+class NanoGPTPreTrainedModel(PreTrainedModel):
+    """Base class for NanoGPT models"""
     config_class = NanoGPTConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["ExactNanoGPTBlock"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.zeros_(module.bias)
+            torch.nn.init.ones_(module.weight)
+
+class NanoGPTModel(NanoGPTPreTrainedModel):
+    """The main NanoGPT model"""
     
     def __init__(self, config):
         super().__init__(config)
@@ -96,8 +122,15 @@ class NanoGPTModel(PreTrainedModel):
         ))
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
         
+        # Initialize weights
         self.post_init()
     
+    def get_input_embeddings(self):
+        return self.transformer.wte
+    
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+    
     def forward(self, input_ids, attention_mask=None, **kwargs):
         device = input_ids.device
         b, t = input_ids.size()
@@ -115,7 +148,8 @@ class NanoGPTModel(PreTrainedModel):
         
         return CausalLMOutputWithCrossAttentions(logits=logits)
     
-    def generate(self, input_ids, max_length=None, max_new_tokens=None, temperature=1.0, top_k=None, do_sample=True, top_p=None, **kwargs):
+    def generate(self, input_ids, max_length=None, max_new_tokens=None, temperature=1.0, 
+                top_k=None, do_sample=True, top_p=None, pad_token_id=None, eos_token_id=None, **kwargs):
         if max_new_tokens is None:
             max_new_tokens = max_length - input_ids.shape[1] if max_length else 50
         
@@ -153,3 +187,6 @@ class NanoGPTModel(PreTrainedModel):
             input_ids = torch.cat((input_ids, idx_next), dim=1)
         
         return input_ids
+
+# For backward compatibility
+NanoGPTForCausalLM = NanoGPTModel