From e2b6747e26983c5e65ddb2ee047a8ab020e4738f Mon Sep 17 00:00:00 2001
From: Ronald Tse <ronald.tse@gmail.com>
Date: Fri, 28 Aug 2026 13:33:15 +0800
Subject: [PATCH] fix(imf): decoded-text repetition guard for
 varying-punctuation loops

The token-window guard bounds periodic loops, but live int8 output
loops with rotating punctuation (phrase + varied separator) never
repeats a verbatim window. Cut when the recent 16 decoded chars echo
3+ times. decode_tokens mirrors the TS runtime (% 256).
---
 src/imf/export.py          | 14 ++++++++++++++
 tests/test_decode_guard.py | 27 +++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/src/imf/export.py b/src/imf/export.py
index e6775c1..2c85130 100644
--- a/src/imf/export.py
+++ b/src/imf/export.py
@@ -38,6 +38,13 @@
 EOS_ID = 1
 
 
+def decode_tokens(tokens: list[int]) -> str:
+    """Inverse of encode_bytes (byte-3 offsets, EOS-terminated)."""
+    return bytes((t - BYTE_OFFSET) % 256 for t in tokens if t >= BYTE_OFFSET).decode(
+        "utf-8", "replace"
+    )
+
+
 def encode_bytes(text: str) -> list[int]:
     """Canonical byte-level tokenization: byte ids + trailing EOS."""
     return [b + BYTE_OFFSET for b in text.encode("utf-8")] + [EOS_ID]
@@ -418,6 +425,13 @@ def onnx_greedy_kv(encoder_sess, kv_sess, text: str, max_len: int = 256) -> list
                 break
         else:
             joined = ",".join(str(t) for t in generated)
+        # Decoded-text guard: loops with varying punctuation never repeat
+        # a verbatim token window — catch the phrase itself echoing.
+        if len(generated) % 8 == 0:
+            text = decode_tokens(generated)
+            suffix = text[-16:]
+            if len(suffix) == 16 and text.count(suffix) >= 3:
+                break
     return generated
 
 
diff --git a/tests/test_decode_guard.py b/tests/test_decode_guard.py
index db969b4..679d85e 100644
--- a/tests/test_decode_guard.py
+++ b/tests/test_decode_guard.py
@@ -74,3 +74,30 @@ def run(self, _, feeds):
 
     out = onnx_greedy_kv(_FakeEncoder(), _StopKV(), "x", max_len=256)
     assert out == [20, 21, 22, 23, 24]
+
+
+def test_varying_separator_loop_is_cut():
+    """Phrase + rotating punctuation never repeats a verbatim token
+    window — the live int8 failure mode. The decoded-text guard cuts it."""
+
+    class _RotateKV(_FakeKV):
+        seps = ['"', " ", "\n", ":"]
+
+        def run(self, _, feeds):
+            import numpy as np
+
+            logits = np.full((1, 1, 260), -1e9)
+            if self.step == 0:
+                logits[0, -1, 100] = 1e9  # phrase token
+            else:
+                mod = self.step % 4
+                if mod == 0:
+                    logits[0, -1, 100] = 1e9  # phrase again
+                else:
+                    # rotating separator tokens 200..203
+                    logits[0, -1, 200 + ((self.step // 4) % 4)] = 1e9
+            self.step += 1
+            return [logits, np.zeros((1,)), np.zeros((1,))]
+
+    out = onnx_greedy_kv(_FakeEncoder(), _RotateKV(), "x", max_len=4096)
+    assert len(out) < 300