Update modeling_dolphy.py

modeling_dolphy.py

@@ -1,40 +1,16 @@
 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)
+from transformers.models.qwen3.modeling_qwen3 import Qwen3Model, Qwen3Config
+import torch.nn as nn
 
 class Dolphy1ForCausalLM(PreTrainedModel):
-    _auto_class = True
+    config_class = Qwen3Config
 
     def __init__(self, config):
         super().__init__(config)
-        self.model = DolphyModel(config)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.model = Qwen3Model(config)
+
+        # If your router was saved as part of the model, this will load it automatically.
+        # No need to redefine or reattach anything here.
 
     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)
+        return self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)