Update modeling_dolphy.py

modeling_dolphy.py

@@ -1,67 +1,40 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
-class RMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.eps = eps
-
-    def forward(self, x):
-        norm = x.pow(2).mean(-1, keepdim=True)
-        return self.weight * x * torch.rsqrt(norm + self.eps)
-
-class MLP(nn.Module):
-    def __init__(self, hidden_size, intermediate_size):
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-
-    def forward(self, x):
-        return self.down_proj(F.gelu(self.gate_proj(x)) * self.up_proj(x))
-
-class DolphyBlock(nn.Module):
-    def __init__(self, hidden_size, intermediate_size, num_heads, fused=False):
-        super().__init__()
-        self.norm1 = RMSNorm(hidden_size)
-        self.attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
-        self.norm2 = RMSNorm(hidden_size)
-        self.mlp = None if fused else MLP(hidden_size, intermediate_size)
-
-    def forward(self, x, attn_mask=None):
-        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x), attn_mask=attn_mask)[0]
-        if self.mlp:
-            x = x + self.mlp(self.norm2(x))
-        return x
-
-class Dolphy1ForCausalLM(nn.Module):
-    def __init__(self, vocab_size=32000, hidden_size=4096, intermediate_size=16384, num_layers=32, num_heads=32, moe_fused=True):
-        super().__init__()
-        self.embed = nn.Embedding(vocab_size, hidden_size)
-        self.blocks = nn.ModuleList([
-            DolphyBlock(hidden_size, intermediate_size, num_heads, fused=moe_fused) for _ in range(num_layers)
-        ])
-        self.norm = RMSNorm(hidden_size)
-        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
-
-    def forward(self, input_ids, attention_mask=None, labels=None):
-        x = self.embed(input_ids)
-        for block in self.blocks:
-            x = block(x, attention_mask)
-        x = self.norm(x)
-        logits = self.lm_head(x)
-
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100)
-
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=None,
-            hidden_states=None,
-            attentions=None,
-        )
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from torch import nn
+
+class DolphyBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attn = nn.Linear(config.hidden_size, config.hidden_size)  # placeholder
+        self.mlp = nn.Linear(config.hidden_size, config.hidden_size)   # placeholder
+
+    def forward(self, x):
+        x = self.attn(x)
+        x = self.mlp(x)
+        return x
+
+class DolphyModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([DolphyBlock(config) for _ in range(config.num_hidden_layers)])
+        self.norm = nn.LayerNorm(config.hidden_size)
+
+    def forward(self, input_ids):
+        x = self.embed_tokens(input_ids)
+        for layer in self.layers:
+            x = layer(x)
+        return self.norm(x)
+
+class Dolphy1ForCausalLM(PreTrainedModel):
+    _auto_class = True
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = DolphyModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        hidden_states = self.model(input_ids)
+        logits = self.lm_head(hidden_states)
+        return CausalLMOutputWithPast(logits=logits)