Fix custom model loading

modeling_custom_gpt.py

@@ -1,4 +1,5 @@
 from transformers import PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
 
 class CustomGPTConfig(PretrainedConfig):
     model_type = "custom_gpt"
@@ -194,8 +195,36 @@ class CustomGPTModel(PreTrainedModel, GenerationMixin):
                 ignore_index=-100
             )
 
-        return {"loss": loss, "logits": logits}
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=logits,
+        )
+
+def port_weights(custom_state, hf_state, config):
+    """
+    Map weights from pretrained modgpt into HF CustomGPTModel.
+    """
+    new_state = hf_state.copy()
+
+    # --- Embeddings ---
+    new_state['embed_tokens.weight'] = custom_state['transformer.wte.weight']
+
+    # --- LM head ---
+    new_state['lm_head.weight'] = custom_state['lm_head.weight']
+
+    # --- Transformer blocks ---
+    for i in range(config.n_layer):
+        # Attention Q, K, V
+        new_state[f'layers.{i}.attn.c_q.weight'] = custom_state[f'transformer.h.{i}.attn.c_q.weight']
+        new_state[f'layers.{i}.attn.c_k.weight'] = custom_state[f'transformer.h.{i}.attn.c_k.weight']
+        new_state[f'layers.{i}.attn.c_v.weight'] = custom_state[f'transformer.h.{i}.attn.c_v.weight']
+        new_state[f'layers.{i}.attn.c_proj.weight'] = custom_state[f'transformer.h.{i}.attn.c_proj.weight']
+
+        # MLP
+        new_state[f'layers.{i}.mlp.c_fc.weight'] = custom_state[f'transformer.h.{i}.mlp.c_fc.weight']
+        new_state[f'layers.{i}.mlp.c_proj.weight'] = custom_state[f'transformer.h.{i}.mlp.c_proj.weight']
 
+    return new_state
 
 # Auto-registration for HuggingFace
 from transformers import AutoConfig, AutoModel, AutoModelForCausalLM