Add custom modeling file

modeling_gptscratch.py

@@ -1,15 +1,13 @@
-# modeling_gptscratch.py
 import torch
 from transformers import PreTrainedModel, GPT2Config
 from transformers.modeling_outputs import CausalLMOutput
-from .gpt_model import GPTModel  # ← repo直下に置いた gpt_model.py を読む
+from .gpt_model import GPTModel
 
 class GPTScratchForCausalLM(PreTrainedModel):
     config_class = GPT2Config
 
     def __init__(self, config, base_model=None):
         super().__init__(config)
-        # 学習時のハイパラに合わせて内部モデルを構築
         self.inner = base_model or GPTModel({
             "vocab_size": config.vocab_size,
             "emb_dim": config.n_embd,
@@ -18,8 +16,7 @@ class GPTScratchForCausalLM(PreTrainedModel):
             "context_length": config.n_positions,
             "drop_rate": 0.1,
         })
-        # HF互換：lm_headを露出（重みはinner側をそのまま共有）
-        self.lm_head = self.inner.out_head
+        self.lm_head = self.inner.out_head  # expose for HF tools
 
     def forward(self, input_ids, **kwargs):
         logits = self.inner(input_ids)
@@ -29,11 +26,9 @@ class GPTScratchForCausalLM(PreTrainedModel):
     def generate(self, input_ids, max_new_tokens=32, eos_token_id=None,
                  do_sample=False, temperature=1.0, top_k=None, top_p=None,
                  repetition_penalty=1.1, **_):
-        # 最小実装（Greedy or 簡易サンプリング）
         for _ in range(max_new_tokens):
             logits = self.forward(input_ids).logits[:, -1, :]
 
-            # 繰り返し抑制
             if repetition_penalty and repetition_penalty != 1.0:
                 for b in range(input_ids.size(0)):
                     logits[b, input_ids[b]] /= repetition_penalty
@@ -48,14 +43,14 @@ class GPTScratchForCausalLM(PreTrainedModel):
                     probs = torch.where(probs >= thresh, probs, torch.zeros_like(probs))
                     probs = probs / probs.sum(dim=-1, keepdim=True)
                 if top_p is not None:
-                    sorted_probs, sorted_idx = probs.sort(descending=True, dim=-1)
-                    cum = sorted_probs.cumsum(dim=-1)
+                    sp, si = probs.sort(descending=True, dim=-1)
+                    cum = sp.cumsum(dim=-1)
                     mask = cum > top_p
                     mask[:, 0] = False
-                    sorted_probs[mask] = 0
-                    probs = torch.zeros_like(probs).scatter(-1, sorted_idx, sorted_probs)
+                    sp[mask] = 0
+                    probs = torch.zeros_like(probs).scatter(-1, si, sp)
                     probs = probs / probs.sum(dim=-1, keepdim=True)
-                next_token = torch.multinomial(probs, num_samples=1)
+                next_token = torch.multinomial(probs, 1)
             else:
                 next_token = torch.argmax(logits, dim=-1, keepdim=True)
 
@@ -64,7 +59,7 @@ class GPTScratchForCausalLM(PreTrainedModel):
                 break
         return input_ids
 
-    # 旧チェックポイントの 'inner.inner.' → 'inner.' を吸収
+    # absorb old checkpoints whose keys start with 'inner.inner.'
     @classmethod
     def _load_state_dict_into_model(cls, model, state_dict, *args, **kwargs):
         remap = {}