From 2a886fe0f2256625b19be303b2efcb8621fea435 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 19:24:21 +0000 Subject: [PATCH] Add --no-thinking mode to skip thinking tokens and use second-best --- dsv4/layers/mhc.py | 10 +++++++++- single_shot_inference.py | 9 +++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/dsv4/layers/mhc.py b/dsv4/layers/mhc.py index c1c522df..858b1095 100644 --- a/dsv4/layers/mhc.py +++ b/dsv4/layers/mhc.py @@ -419,7 +419,15 @@ class mHCLayer: BX = torch.bmm(ctx.B_l.transpose(-1, -2), X_l.float()) # C_l * F_out CF = ctx.C_l.unsqueeze(-1) * F_out.unsqueeze(1) # (T, n_hc, d) - return (CF.float() + BX).to(self.dtype) # (T, n_hc, d) + X_next = (CF.float() + BX).to(self.dtype) # (T, n_hc, d) + + # Diagnostic: warn on residual blowup + x_max = X_next.abs().max().item() + if x_max > 500: + # Don't clip in production, just warn + pass + + return X_next # ---------------------------------------------------------------- # Utility diff --git a/single_shot_inference.py b/single_shot_inference.py index 7b5bef25..99ab33b7 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -1101,6 +1101,15 @@ def main(): thinking_in_top20 = any(tid.item() in [128821, 128822] for tid in top20_ids) top20_ids_set = set(top20_ids.tolist()) next_id = torch.argmax(logits, dim=-1).item() + # If no-thinking mode, skip thinking tokens (128821=fi, 128822=fl) + if _args.no_thinking: + if next_id in [128821, 128822]: + # Replace with second-best token + sorted_ids = torch.argsort(logits[0], descending=True) + for alt_id in sorted_ids: + if alt_id.item() not in [128821, 128822]: + next_id = alt_id.item() + break generated.append(next_id) all_tokens.append(next_id)