Add early-stop on special tokens + repetition penalty controls in standalone runner

run_standalone_inference.py

@@ -49,6 +49,24 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--temperature", type=float, default=0.8)
     p.add_argument("--top-k", type=int, default=50)
     p.add_argument("--max-new-tokens", type=int, default=96)
+    p.add_argument("--min-new-tokens", type=int, default=1, help="Do not stop before this many generated tokens")
+    p.add_argument(
+        "--repetition-penalty",
+        type=float,
+        default=1.08,
+        help="Penalty >1.0 discourages repeats (1.0 disables).",
+    )
+    p.add_argument(
+        "--stop-tokens",
+        default="<|assistant_end|>,<|output_end|>,<|user_start|>",
+        help="Comma-separated special tokens that trigger early stop.",
+    )
+    p.add_argument(
+        "--stop-on-special",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Enable early stop when generated token matches --stop-tokens.",
+    )
     p.add_argument("--seed", type=int, default=42)
     p.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto")
     p.add_argument(
@@ -160,6 +178,17 @@ def _top_k_filter(logits, top_k: int):
     return masked
 
 
+def _apply_repetition_penalty(logits, token_ids, penalty: float):
+    import torch
+
+    if penalty <= 1.0 or token_ids.numel() == 0:
+        return logits
+    uniq = torch.unique(token_ids)
+    penalized = logits.clone()
+    penalized[:, uniq] = penalized[:, uniq] / penalty
+    return penalized
+
+
 def _has_meta_tensors(model) -> bool:
     import torch
 
@@ -224,7 +253,43 @@ def _load_model_from_pretrained(local_dir: str, torch_dtype, device, use_device_
     return model
 
 
-def generate_compat(model, input_ids, max_new_tokens: int, temperature: float, top_k: int, seed: int):
+def _parse_stop_tokens(raw: str) -> List[str]:
+    out = []
+    for part in raw.split(","):
+        token = part.strip()
+        if token:
+            out.append(token)
+    return out
+
+
+def _resolve_stop_ids(tokenizer, stop_tokens: List[str]) -> List[int]:
+    ids: List[int] = []
+    for token in stop_tokens:
+        tid = _special_id(tokenizer, token)
+        if tid is not None:
+            ids.append(int(tid))
+    # preserve order but deduplicate
+    seen = set()
+    uniq = []
+    for tid in ids:
+        if tid not in seen:
+            seen.add(tid)
+            uniq.append(tid)
+    return uniq
+
+
+def generate_compat(
+    model,
+    input_ids,
+    max_new_tokens: int,
+    min_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    repetition_penalty: float,
+    stop_ids: List[int],
+    stop_on_special: bool,
+    seed: int,
+):
     import torch
     import torch.nn.functional as F
 
@@ -237,6 +302,7 @@ def generate_compat(model, input_ids, max_new_tokens: int, temperature: float, t
     for _ in range(max_new_tokens):
         attention_mask = torch.ones_like(ids)
         logits = model(input_ids=ids, attention_mask=attention_mask, return_dict=True).logits[:, -1, :]
+        logits = _apply_repetition_penalty(logits, ids[0], penalty=repetition_penalty)
         logits = _top_k_filter(logits, top_k)
         if temperature > 0:
             probs = F.softmax(logits / temperature, dim=-1)
@@ -244,6 +310,10 @@ def generate_compat(model, input_ids, max_new_tokens: int, temperature: float, t
         else:
             next_ids = torch.argmax(logits, dim=-1, keepdim=True)
         ids = torch.cat((ids, next_ids), dim=1)
+        if stop_on_special and stop_ids and (ids.shape[1] - input_ids.shape[1]) >= max(0, min_new_tokens):
+            next_id = int(next_ids.item())
+            if next_id in stop_ids:
+                break
     return ids
 
 
@@ -300,6 +370,8 @@ def main() -> None:
 
     prompt_style = _resolve_prompt_style(args.prompt_style, tok)
     max_prompt_tokens = max(1, context_window - args.max_new_tokens)
+    stop_tokens = _parse_stop_tokens(args.stop_tokens)
+    stop_ids = _resolve_stop_ids(tok, stop_tokens) if args.stop_on_special else []
 
     print(f"repo={args.repo_id} revision={args.revision} snapshot_sha={sha}")
     print(
@@ -307,6 +379,10 @@ def main() -> None:
         f"use_device_map={args.use_device_map} load_mode={args.load_mode}"
     )
     print(f"context_window={context_window} max_prompt_tokens={max_prompt_tokens}")
+    print(
+        f"generation: temperature={args.temperature} top_k={args.top_k} "
+        f"repetition_penalty={args.repetition_penalty} stop_on_special={args.stop_on_special} stop_ids={stop_ids}"
+    )
 
     torch.manual_seed(args.seed)
     if torch.cuda.is_available():
@@ -321,8 +397,12 @@ def main() -> None:
                 model=model,
                 input_ids=input_ids,
                 max_new_tokens=args.max_new_tokens,
+                min_new_tokens=args.min_new_tokens,
                 temperature=args.temperature,
                 top_k=args.top_k,
+                repetition_penalty=args.repetition_penalty,
+                stop_ids=stop_ids,
+                stop_on_special=args.stop_on_special,
                 seed=args.seed + i,
             )