feat: add decoding controls

inference.py

@@ -30,6 +30,33 @@ CTX_LEN     = 2048
 BOS_ID, EOS_ID, EOT_ID = 0, 1, 6
 
 
+def apply_repetition_penalty(
+    logits: np.ndarray,
+    token_ids: list[int],
+    penalty: float,
+) -> np.ndarray:
+    """既出トークンの logit に repetition penalty を適用する。"""
+    if penalty <= 1.0 or not token_ids:
+        return logits
+    adjusted = logits.copy()
+    for token_id in set(token_ids):
+        if adjusted[token_id] > 0:
+            adjusted[token_id] /= penalty
+        else:
+            adjusted[token_id] *= penalty
+    return adjusted
+
+
+def filter_top_k(logits: np.ndarray, top_k: int) -> np.ndarray:
+    """top_k 以外の logit を -inf にする。top_k<=0 なら変更しない。"""
+    if top_k <= 0 or top_k >= logits.shape[-1]:
+        return logits
+    filtered = np.full_like(logits, -np.inf)
+    top_indices = np.argpartition(logits, -top_k)[-top_k:]
+    filtered[top_indices] = logits[top_indices]
+    return filtered
+
+
 # ── Packed Linear Modules ────────────────────────────────────────────────────
 
 class Linear1bit(nn.Module):
@@ -254,23 +281,26 @@ class SSMoELMPacked(nn.Module):
 
     @torch.inference_mode()
     def generate(self, input_ids: list[int], max_new_tokens: int = 200,
-                 temperature: float = 0.8, top_p: float = 0.9,
+                 temperature: float = 0.0, top_p: float = 0.9,
+                 top_k: int = 1, repetition_penalty: float = 1.3,
                  eos_ids: tuple[int, ...] = (EOS_ID, EOT_ID)) -> list[int]:
         ids = list(input_ids)
         generated = []
         for _ in range(max_new_tokens):
             x = torch.tensor([ids[-CTX_LEN:]], dtype=torch.long)
             logits = self(x)[0, -1]
+            logits_np = logits.numpy().astype(np.float64)
+            logits_np = apply_repetition_penalty(logits_np, ids, repetition_penalty)
+            logits_np = filter_top_k(logits_np, top_k)
             if temperature > 0:
-                logits_np = logits.numpy().astype(np.float64)
-                logits_np = (logits_np - logits_np.max()) / temperature
+                logits_np = (logits_np - np.nanmax(logits_np)) / temperature
                 probs = np.exp(logits_np); probs /= probs.sum()
                 idx = np.argsort(-probs); cumsum = np.cumsum(probs[idx])
                 cutoff = np.searchsorted(cumsum, top_p) + 1
                 probs[idx[cutoff:]] = 0.0; probs /= probs.sum()
                 next_id = int(np.random.choice(idx, p=probs))
             else:
-                next_id = int(logits.argmax().item())
+                next_id = int(np.nanargmax(logits_np))
             if next_id in eos_ids:
                 break
             generated.append(next_id)
@@ -369,8 +399,10 @@ def main():
     parser.add_argument("--prompt",      default=None, help="Prompt for single-shot mode (--no-chat)")
     parser.add_argument("--no-chat",     action="store_true", help="Single-shot mode")
     parser.add_argument("--max-tokens",  type=int,   default=200,  help="Max new tokens to generate")
-    parser.add_argument("--temperature", type=float, default=0.8,  help="Sampling temperature")
+    parser.add_argument("--temperature", type=float, default=0.0,  help="Sampling temperature (0 = greedy; recommended for this model)")
     parser.add_argument("--top-p",       type=float, default=0.9,  help="Top-p nucleus sampling")
+    parser.add_argument("--top-k",       type=int,   default=1,    help="Keep only top-k logits before sampling (0 disables)")
+    parser.add_argument("--repetition-penalty", type=float, default=1.3, help="Penalty for tokens already present in the prompt/output")
     args = parser.parse_args()
 
     print(f"Loading {args.ckpt} ...")
@@ -402,7 +434,14 @@ def main():
                 continue
             ids = build_chat_prompt(tok, history, user_input)
             print("Assistant: ", end="", flush=True)
-            out = model.generate(ids, args.max_tokens, args.temperature, args.top_p)
+            out = model.generate(
+                ids,
+                args.max_tokens,
+                args.temperature,
+                args.top_p,
+                args.top_k,
+                args.repetition_penalty,
+            )
             response = tok.decode(out).strip()
             print(response)
             history.append((user_input, response))
@@ -412,7 +451,14 @@ def main():
         ids = [BOS_ID] + tok.encode(prompt).ids
         print(f"\nPrompt: {prompt}")
         print("Output: ", end="", flush=True)
-        out = model.generate(ids, args.max_tokens, args.temperature, args.top_p)
+        out = model.generate(
+            ids,
+            args.max_tokens,
+            args.temperature,
+            args.top_p,
+            args.top_k,
+            args.repetition_penalty,
+        )
         print(tok.decode(out))