Upload modeling_nexus_small.py with huggingface_hub

modeling_nexus_small.py

@@ -4,42 +4,12 @@ from typing import Optional
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel
+from transformers import PreTrainedModel, AutoConfig, AutoModel, AutoModelForCausalLM
 
 try:
-    from config import NexusConfig
+    from config import NexusConfig, NexusSmAllConfig
 except ImportError:
-    from .config import NexusConfig
-
-
-class NexusConfigHF(PretrainedConfig):
-    model_type = "nexus_small"
-
-    def __init__(
-        self,
-        vocab_size=50304,
-        max_seq_len=512,
-        dim=768,
-        num_layers=10,
-        num_heads=12,
-        num_kv_heads=4,
-        multiple_of=256,
-        ff_dim=2048,
-        norm_eps=1e-6,
-        rope_theta=500000.0,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.vocab_size = vocab_size
-        self.max_seq_len = max_seq_len
-        self.dim = dim
-        self.num_layers = num_layers
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.multiple_of = multiple_of
-        self.ff_dim = ff_dim
-        self.norm_eps = norm_eps
-        self.rope_theta = rope_theta
+    from .config import NexusConfig, NexusSmAllConfig
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
@@ -148,12 +118,12 @@ class TransformerBlock(nn.Module):
         return out
 
 class NexusPreTrainedModel(PreTrainedModel):
-    config_class = NexusConfigHF
+    config_class = NexusSmAllConfig
     base_model_prefix = "nexus_small"
     supports_gradient_checkpointing = False
     _no_split_modules = ["TransformerBlock"]
 
-    def __init__(self, config: NexusConfigHF):
+    def __init__(self, config: NexusSmAllConfig):
         super().__init__(config)
         self._nexus_config = NexusConfig(
             vocab_size=config.vocab_size,
@@ -248,9 +218,12 @@ class Nexus(nn.Module):
         return generated, input_ids
 
 class NexusForCausalLM(NexusPreTrainedModel):
-    def __init__(self, config: NexusConfigHF):
+    def __init__(self, config: NexusSmAllConfig):
         super().__init__(config)
         self.model = Nexus(self._nexus_config)
+        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
+        with torch.no_grad():
+            self.lm_head.weight = self.model.output.weight
 
     def forward(self, input_ids, attention_mask=None, **kwargs):
         return self.model(input_ids)
@@ -266,8 +239,9 @@ class NexusForCausalLM(NexusPreTrainedModel):
         return {"input_ids": input_ids}
 
 
-AutoConfig.register("nexus_small", NexusConfigHF)
-AutoModel.register(NexusConfigHF, NexusForCausalLM)
+AutoConfig.register("nexus_small", NexusSmAllConfig)
+AutoModel.register(NexusSmAllConfig, NexusForCausalLM)
+AutoModelForCausalLM.register(NexusSmAllConfig, NexusForCausalLM)
 
 
 def create_nexus_model():