Fix: prevent infinite “A5A5A5...” repetition loop during text generation (Issue #1008)

Added a simple repetition detection inside the generate() loop to stop the model 
from endlessly producing patterns like “A5A5A5...”. 

- Imported `re` for regex pattern matching.
- Stops generation if “A5” repeats 10+ times or the same token appears 10 times.
- Prints a warning and exits safely instead of looping infinitely.

Fixes: #1008
This commit is contained in:
Ceaser1717 2025-10-21 16:02:03 +05:30 committed by GitHub
parent 9b4e9788e4
commit 21a919d79f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -9,6 +9,7 @@ from transformers import AutoTokenizer
from safetensors.torch import load_model
from model import Transformer, ModelArgs
import re
def sample(logits, temperature: float = 1.0):
@ -67,6 +68,24 @@ def generate(
tokens[:, cur_pos] = next_token
finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id)
prev_pos = cur_pos
# === FIX for DeepSeek-V3 Issue #1008 ===
# Detect infinite repeating loops like "A5A5A5A5..." or same token repetition
# Convert current generated tokens to string for pattern check
decoded_output = "".join([str(t) for t in tokens[0, :cur_pos].tolist()])
# (a) Detect if "A5" repeats 10 or more times continuously
if re.search(r"(A5){10,}", decoded_output):
print("[Warning] Detected excessive 'A5' repetition — stopping generation early.")
break
# (b) Detect same token repeated multiple times (generic loop check)
if cur_pos > 10:
last_10 = tokens[0, cur_pos-10:cur_pos].tolist()
if len(set(last_10)) == 1:
print("[Warning] Detected same token repetition — stopping generation early.")
break
# === END FIX ===
if finished.all():
break
completion_tokens = []