fixed inference.py

inference.py

@@ -11,51 +11,65 @@ def generate_stream(
     top_k=None,
 ):
     """
-    最小・安全なストリーミング生成
-    - batch size = 1 前提
-    - KV cache 利用
+    ストリーミング生成（batch size = 1 固定）
+    - GPT.generate と同じロジック
+    - KV cache 使用
+    - top-k / top-p 対応
     """
     model.eval()
     next_token = None
 
     with torch.no_grad():
         for i in range(max_new_tokens):
+
+            # ===== forward =====
             if i == 0:
                 logits, _ = model(input_ids, None, use_cache=True)
             else:
                 logits, _ = model(next_token, None, use_cache=True)
 
-            logits = logits[:, -1, :] / temperature
+            # last token logits
+            last_logits = logits[:, -1, :] / temperature  # [1, vocab]
 
-            # top-k
+            # ===== top-k =====
             if top_k is not None:
-                top_k = min(top_k, logits.size(-1))
-                values, _ = torch.topk(logits, top_k)
-                min_val = values[:, -1].unsqueeze(-1)
-                logits = torch.where(
-                    logits < min_val,
-                    torch.full_like(logits, float("-inf")),
-                    logits,
+                top_k = min(top_k, last_logits.size(-1))
+                values, _ = torch.topk(last_logits, top_k)
+                min_value = values[:, -1].unsqueeze(-1)
+                last_logits = torch.where(
+                    last_logits < min_value,
+                    torch.full_like(last_logits, float("-inf")),
+                    last_logits,
                 )
 
-            # top-p (nucleus)
+            # ===== top-p (nucleus) =====
             if top_p is not None:
-                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-                probs = F.softmax(sorted_logits, dim=-1)
-                cum_probs = torch.cumsum(probs, dim=-1)
+                sorted_logits, sorted_indices = torch.sort(
+                    last_logits, descending=True
+                )
+                sorted_probs = F.softmax(sorted_logits, dim=-1)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
 
-                mask = cum_probs > top_p
-                mask[..., 1:] = mask[..., :-1]
-                mask[..., 0] = False
+                sorted_mask = cumulative_probs > top_p
+                # ★ ここが重要：clone() を入れる
+                sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+                sorted_mask[..., 0] = False
+
+                sorted_logits = torch.where(
+                    sorted_mask,
+                    torch.full_like(sorted_logits, float("-inf")),
+                    sorted_logits,
+                )
 
-                sorted_logits[mask] = -float("inf")
-                logits = torch.zeros_like(logits).scatter(
+                last_logits = torch.zeros_like(last_logits).scatter(
                     -1, sorted_indices, sorted_logits
                 )
 
-            probs = F.softmax(logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
+            # ===== sample =====
+            probs = F.softmax(last_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)  # [1, 1]
 
             yield int(next_token.item())
 
+            # 次ステップ用に連結
             input_ids = torch.cat([input_ids, next_token], dim=1)