Add custom modeling file

modeling_gptscratch.py

@@ -1,18 +1,52 @@
-from transformers import PreTrainedModel, GPT2Config
+# modeling_gptscratch.py
+import torch
 import torch.nn as nn
-from .gpt_model import GPTModel  # ✅ 実際の場所に合わせて変更
+from transformers import PreTrainedModel, GPT2Config
+from transformers.modeling_outputs import CausalLMOutput
+from .gpt_model import GPTModel  # ← gpt_model.py を同梱済み前提
 
 class GPTScratchForCausalLM(PreTrainedModel):
     config_class = GPT2Config
+
     def __init__(self, config):
         super().__init__(config)
+        # 学習時の cfg に合わせて内部モデルを構築
         self.inner = GPTModel({
             "vocab_size": config.vocab_size,
             "emb_dim": config.n_embd,
             "n_heads": config.n_head,
             "n_layers": config.n_layer,
             "context_length": config.n_positions,
-            "drop_rate": 0.1
+            "drop_rate": 0.1,
         })
+
+        # 互換のために lm_head を生やしておく（重みは inner.out_head を参照）
+        self.lm_head = self.inner.out_head
+
     def forward(self, input_ids, **kwargs):
-        return self.inner(input_ids)
+        logits = self.inner(input_ids)
+        # HF の慣習に合わせて CausalLMOutput で返す
+        return CausalLMOutput(logits=logits)
+
+    # --- これが肝：最小実装の generate（greedy） ---
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens=32, **gen_kwargs):
+        # 非教師ありの最小版（Greedy）。pad/attention_mask 等は省略
+        for _ in range(max_new_tokens):
+            out = self.forward(input_ids)
+            next_token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+        return input_ids
+
+    # （任意）古いチェックポイントからのキー名ズレを吸収
+    @classmethod
+    def _load_state_dict_into_model(cls, model, state_dict, *args, **kwargs):
+        # inner.inner.* → inner.* にリネーム
+        remap = {}
+        for k, v in list(state_dict.items()):
+            if k.startswith("inner.inner."):
+                remap[k.replace("inner.inner.", "inner.", 1)] = v
+                del state_dict[k]
+        state_dict.update(remap)
+        return super()._load_state_dict_into_model(model, state_dict, *args, **kwargs)
+